如何从pandas DataFrame的滚动窗口访问两列的值？

Question

我现在的目标是根据经纬度系统找到两个点之间的距离，以便追踪航班的轨迹。我有一个 pandas 数据框，里面包含不断变化的经纬度值。为了找到这些点之间的距离，我使用了 haversine 距离函数，这个函数可以接收这些值作为输入，计算出公里数。

我最开始尝试用一个循环，遍历航班的长度，计算距离，类似于下面的代码：

    for i in range(len(df) - 1):
        row1 = df.iloc[i]
        row2 = df.iloc[i + 1]
        result = haversine_distance(row1, row2)

但是数据集非常大，计算效率不高，所以我换了个方法。

接着，我尝试使用 pandas 的 df.rolling 函数来实现一个滚动窗口，并结合一个带有 lambda 函数的 .apply，像下面这样：

df['DISTANCE'] = df[['Latitude', 'Longitude']].rolling(window=2).apply(lambda x: haversine_distance(x), raw = True)

我理解这里发生的事情是，一个二维数组（来自 raw = True）被传递给 haversine 函数，里面包含了窗口中的 4 个经纬度值。

然而，我得到的是一个一维数组，而不是我想要的从 2 列中提取的 4 个值的二维数组。我的意思是：

df = pd.DataFrame({'Latitude': [40.7128, 37.7749, 34.0522],
                   'Longitude': [-74.0060, -122.4194, -118.2437]})

如果数据框像上面那样，我应该得到数组 [[40.7128, -74.0060],[37.7749,-122.4194]]。

我该如何修正我的代码，或者换个方法来获取这些值呢？下面是 haversine 函数：

def haversine_distance(ndarray):
    lat1, lat2 = ndarray[0][0], ndarray[0][1]
    lon1, lon2 = ndarray[1][0], ndarray[1][1]

    # Convert latitude and longitude from degrees to radians
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])

    # Haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat / 2.0) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2.0) ** 2
    c = 2 * np.arcsin(np.sqrt(a))
    km = 6371 * c
    return km

这是我想要的输出：

df = pd.DataFrame({'Latitude': [40.7128, 37.7749, 34.0522],
                   'Longitude': [-74.0060, -122.4194, -118.2437],
                   'DISTANCE': [0, 4129.0861, 559.1205]})

Answer 1

你可以直接使用 numpy：

def haversine_distance(arr):
  arr = np.radians(arr).to_numpy()
  dlat, dlon = np.diff(arr, axis = 0).T
  a = np.sin(dlat / 2.0) ** 2 + \
      np.cos(arr[1:,0]) * np.cos(arr[:-1,0]) * np.sin(dlon / 2.0) ** 2
  return np.r_[np.nan, 6371 * 2 * np.arcsin(np.sqrt(a))]

df.assign(dist = haversine_distance(df))
   Latitude  Longitude         dist
0   40.7128   -74.0060          NaN
1   37.7749  -122.4194  4129.086165
2   34.0522  -118.2437   559.120577

Answer 2

你需要把你的haversine函数进行向量化处理，然后创建一个有4列的数组，顺序要正确（可以用shift、concat和to_numpy来实现），然后把这个数组传给函数：

df = pd.DataFrame({'Latitude': [40.7128, 37.7749, 34.0522],
                   'Longitude': [-74.0060, -122.4194, -118.2437]})

def haversine_distance(ndarray):
    # get the coordinates as 4 vectors
    # Convert latitude and longitude from degrees to radians
    lat1, lat2, lon1, lon2 = np.radians(ndarray.T)

    # Haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat / 2.0) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2.0) ** 2
    c = 2 * np.arcsin(np.sqrt(a))
    km = 6371 * c
    return km

a = (pd.concat([df[['Latitude', 'Longitude']],
                df[['Latitude', 'Longitude']].shift()
               ], axis=1)
       .iloc[:, [0,2,1,3]].to_numpy()
    )

df['DISTANCE'] = haversine_distance(a)

输出结果：

   Latitude  Longitude     DISTANCE
0   40.7128   -74.0060          NaN
1   37.7749  -122.4194  4129.086165
2   34.0522  -118.2437   559.120577

注意：与其用 .iloc[:, [0,2,1,3]] 来重新排列列的顺序，你也可以在函数中使用 lat1, lon1, lat2, lon2 = ndarray.T。

中间变量 a：

#            lat1        lat2      lon1       lon2
array([[  40.7128,       nan,  -74.006 ,       nan],
       [  37.7749,   40.7128, -122.4194,  -74.006 ],
       [  34.0522,   37.7749, -118.2437, -122.4194]])

另外，你也可以写一个函数，直接把 df 作为输入：

def haversine_distance_df(df, cols=['Latitude', 'Longitude']):
    # get the coordinates as 4 vectors    
    # Convert latitude and longitude from degrees to radians
    lat1, lon1, lat2, lon2 = np.radians(
        pd.concat([df[cols], df[cols].shift()], axis=1)
          .to_numpy().T
    )
    
    # Haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat / 2.0) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2.0) ** 2
    c = 2 * np.arcsin(np.sqrt(a))
    km = 6371 * c
    return km

df['DISTANCE'] = haversine_distance_df(df)