代码加速需求：统计指定半径内的搁浅水母点位数量

嘿，我来帮你搞定这个效率问题！你的核心痛点是用了Python双重循环处理百万级粒子，这肯定慢得离谱——Python循环在这种量级的数据面前完全发挥不了优势，咱们得把计算转到numpy的向量化操作上，再加上一些小技巧，速度能提升几十甚至上百倍。

问题回顾 #

我需要统计某海滩指定半径范围内搁浅的水母数量，现有经纬度的masked数组（Xpos1和Ypos1结构类似），数组维度为124个时间步×800,000个粒子。当前代码仅计算3个时间步就耗时5分钟，多区域计算效率极低，希望优化代码生成每个时间步对应半径内水母数量的数组Blen（长度124）。

原数组结构示例：

Xpos1 = masked_array(
    data=[[50.0441, 50.0640, ..., --, --], ...],
    mask=[[False, False, ..., True, True], ...],
    fill_value=9.96921e+36, dtype=float32
)
# len(Xpos1)=124（时间步），len(Xpos1[0])=800000（粒子数）

原低效代码的问题就在于嵌套的for循环——遍历80万次/时间步的Python循环，速度根本跟不上。

优化方案 #

1. 用Numpy向量化操作替代Python循环 #

Numpy的底层是C实现的向量化运算，能一次性处理整个数组的判断，完全避免Python层面的循环开销。同时我们要利用masked数组的compressed()方法快速过滤掉无效的掩码值。

2. 可选：从矩形框过滤升级为精确圆形半径判断 #

原代码用矩形框近似半径范围，如果你需要更准确的圆形半径统计，可以先通过矩形框快速过滤掉大部分粒子，再对剩余粒子计算精确的球面距离，这样既保证速度又提升精度。

完整优化代码 #

import numpy as np
import math

# 待检测海滩位置
lon_beach = 2.631547
lat_beach = 51.120983
radius_km = 10  # 检测半径（公里）
total_timesteps = 124  # 现在可以轻松处理全部124个时间步

# 计算目标纬度处，1度经/纬度对应的公里数
earth_radius = 6371  # 地球平均半径（km）
lat_per_km = 360 / (2 * math.pi * earth_radius)  # 每公里对应的纬度度数
lon_per_km = 360 / (2 * math.pi * earth_radius * math.cos(math.radians(lat_beach)))  # 每公里对应的经度度数

# 生成矩形过滤框（用于快速粗筛）
lon_min = lon_beach - radius_km * lon_per_km
lon_max = lon_beach + radius_km * lon_per_km
lat_min = lat_beach - radius_km * lat_per_km
lat_max = lat_beach + radius_km * lat_per_km

# 初始化结果数组
count_per_timestep = np.zeros(total_timesteps, dtype=np.int32)

# 向量化处理每个时间步
for step in range(total_timesteps):
    # 获取当前时间步的有效经纬度（自动跳过masked值）
    lons = Xpos1[step].compressed() if Xpos1[step].mask.any() else Xpos1[step]
    lats = Ypos1[step].compressed() if Ypos1[step].mask.any() else Ypos1[step]
    
    # 第一步：快速粗筛，找出矩形框内的粒子
    in_bounds = (lats >= lat_min) & (lats <= lat_max) & (lons >= lon_min) & (lons <= lon_max)
    filtered_lons = lons[in_bounds]
    filtered_lats = lats[in_bounds]
    
    # 第二步：精确计算球面距离（Haversine公式），筛选圆形半径内的粒子
    # 向量化实现Haversine公式，避免循环
    def haversine(lon1, lat1, lon2, lat2):
        lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])
        dlon = lon2 - lon1
        dlat = lat2 - lat1
        a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
        c = 2 * np.arcsin(np.sqrt(a))
        return c * earth_radius
    
    distances = haversine(filtered_lons, filtered_lats, lon_beach, lat_beach)
    in_radius = distances <= radius_km
    
    # 记录当前时间步的有效计数
    count_per_timestep[step] = np.sum(in_radius)

# 最终结果就是你要的Blen
Blen = count_per_timestep.tolist()
print(Blen)

额外提速技巧 #

• Numba加速：如果还想再快一点，用numba编译Haversine函数，能把距离计算速度再提升几倍：

  from numba import jit
  
  @jit(nopython=True)
  def haversine_numba(lon1, lat1, lon2, lat2):
      lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])
      dlon = lon2 - lon1
      dlat = lat2 - lat1
      a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
      c = 2 * np.arcsin(np.sqrt(a))
      return c * earth_radius
  
  # 使用时直接替换原haversine函数即可
  distances = haversine_numba(filtered_lons, filtered_lats, lon_beach, lat_beach)
  

• 并行处理时间步：如果你的CPU有多个核心，可以用joblib或者multiprocessing把时间步的计算并行化，进一步缩短总耗时。
• 预处理Mask：如果所有时间步的mask是相同的，可以一次性把整个Xpos1和Ypos1压缩成二维数组，避免每个时间步重复处理mask。

内容的提问来源于stack exchange，提问作者Jellyse