根据先前的函数结果将函数应用于 Dataframe 的行

Question

我正在尝试迭代 pandas DataFrame 的行并将函数一一应用于行。函数的输入值取决于上一行的结果。

这是一个例子：

import numpy as np
import pandas as pd
import math


def predict_loc(df, lon, lat):
    R = 6378.1  # Radius of the Earth
    brng = np.deg2rad(df.wdir)  # Bearing is radians.
    d = df.wspd * df.delta * 60 / 1e3  # Distance in km

    lat2 = math.asin(
        math.sin(lat) * math.cos(d / R)
        + math.cos(lat) * math.sin(d / R) * math.cos(brng)
    )

    lon2 = lon + math.atan2(
        math.sin(brng) * math.sin(d / R) * math.cos(lat),
        math.cos(d / R) - math.sin(lat) * math.sin(lat2),
    )

    lat2 = np.rad2deg(lat2)
    lon2 = np.rad2deg(lon2)

    return lon2, lat2


dates = pd.date_range("20130101", periods=6, freq="1H")
df = pd.DataFrame(
    np.random.randn(6, 3),
    index=dates,
    columns=[
        "wdir",
        "wspd",
        "delta",
    ],
)

lon = 0
lat = 1

for index, row in df.iterrows():
    lon, lat = predict_loc(row, lon, lat)

在本例中，lon 和 lat 的初始值分别为 0 和 1。 然后，由predict_loc 函数预测位置。新的 lon 和 lat 是下一行的输入。我想要的是最后的经纬度。

有没有更快的方法来完成这项任务？谢谢。

Answer 1

因此，您的代码平均运行时间 0.0002 秒：

import statistics
import time

np.random.seed(0)  # In order to get consistant results

iterations = 25_000

elapsed_time = []
for i in range(iterations):
    start_time = time.time()
    lon = 0
    lat = 1
    for index, row in df.iterrows():
        lon, lat = predict_loc(row, lon, lat)
    elapsed_time.append(time.time() - start_time)

print(lon, lat)
# 6861.350646683788 -63.005854847412145
print(f"--- {statistics.mean(elapsed_time):2f} seconds in average ---")
# --- 0.000233 seconds in average ---

在你工作的这一点上，鉴于你只对最终结果感兴趣，你不需要 Pandas，我建议使用 Python，稍微修改 predict_loc 并定义一个操作 @ 的辅助函数987654323@ 而不是Series，像这样：

def new_predict_loc(*args):
    wdir, wspd, delta, lon, lat = args
    R = 6378.1  # Radius of the Earth
    brng = np.deg2rad(wdir)  # Bearing is radians.
    d = wspd * delta * 60 / 1e3  # Distance in km

    lat2 = math.asin(
        math.sin(lat) * math.cos(d / R)
        + math.cos(lat) * math.sin(d / R) * math.cos(brng)
    )

    lon2 = lon + math.atan2(
        math.sin(brng) * math.sin(d / R) * math.cos(lat),
        math.cos(d / R) - math.sin(lat) * math.sin(lat2),
    )

    lat2 = np.rad2deg(lat2)
    lon2 = np.rad2deg(lon2)

    return lon2, lat2


def compute_coordinates(df):
    n = 0
    wdir = df["wdir"].to_list()
    wspd = df["wspd"].to_list()
    delta = df["delta"].to_list()
    lon, lat = 0, 1

    while n < df.shape[0]:
        lon, lat = new_predict_loc(wdir[n], wspd[n], delta[n], lon, lat)
        n += 1
    return lon, lat

这样，计算的执行时间平均为 0.00003 秒，比以前快了近 7 倍：

elapsed_time = []
for i in range(iterations):
    start_time = time.time()
    lon, lat = compute_coordinates(df)
    elapsed_time.append(time.time() - start_time)

print(lon, lat)
# Same results as before
# 6861.350646683788 -63.005854847412145
print(f"--- {statistics.mean(elapsed_time):2f} seconds in average ---")
# --- 0.000027 seconds in average ---