在熊貓中更快的應用方法

Question

我有一個函數，我正在嘗試應用於位置的數據幀。 具體來說，我想附加一個新列，其中包含每個站點最近的10個站點。 以下似乎有效，但速度極慢。

def distance(first_lat, first_lon, second_lat, second_lon):
    return ((first_lat - second_lat) ** 2 + (first_lon - second_lon) ** 2) ** 0.5


def load_site_list():
    '''
    This function generates a dataframe with all the available sites
    '''
    url = 'ftp://ftp.ncdc.noaa.gov/pub/data/noaa/isd-history.csv'
    cols = ["STATION NAME",
            "LAT",
            "LON"]
    df = pd.read_csv(url, parse_dates=False, usecols=cols)
    df = df.dropna(subset=['LAT'])
    df = df.dropna(subset=['LON'])
    df['LAT'] = df['LAT'].astype(float)
    df['LON'] = df['LON'].astype(float)
    return df

sites = load_site_list()
sites['closest'] = ""
for index, row in sites.iterrows():
    sites['dist'] = sites.apply(lambda line: distance(line['LAT'], line['LON'], row['LAT'], row['LON']), axis=1)
    sites.sort_values('dist', inplace=True)
    sites['closest'][index] = sites['STATION NAME'].iloc[1:11].tolist()

看起來for循環中第一行生成距當前列的距離，每個循環占用一秒。 在這里有超過10,000行循環...是否有更快的方法這樣做？

Answer 1

請注意，您的代碼的時間復雜度為O（n ^ 2）：在這種情況下，您在for循環中的apply函數中計算30k * 30k = 9億個距離，即純Python。

大熊貓中的向量操作在C中實現，因此如果計算單個向量操作中的所有距離，則可以獲得相對加速。

如果你有足夠的RAM，你可以進行笛卡爾連接，計算所有成對距離，然后進行排序，分組，然后取頭，如下：

# code to reduce memory usage
sites['site_code'] = pd.Categorical(sites['STATION NAME']).codes
sites['LAT'] = sites.LAT.astype(np.float16)
sites['LON'] = sites.LAT.astype(np.float16)
sites_small = sites[['site_code','LAT','LON']].copy()
sites_small.index = [0]*len(sites_small)

pairs = sites_small.join(sites_small,lsuffix='_x',rsuffix='_y')
pairs['dist'] = (pairs['LAT_x'] - pairs['LAT_y'])**2 + (pairs['LON_x'] - pairs['LON_y'])**2
pairs.sort_values(['STATION NAME_x','dist'], inplace = True) # actually, just sorting by dist is sufficient
pairs.groupby('STATION NAME_x').head(10)

不幸的是，您可能沒有足夠的RAM：如果您將站點名稱編碼為16位整數，並將坐標編碼為16位浮點數，則每行需要12個字節（因為您正在查看對），以及另一個索引的8個字節（pandas將這些帶入連接中的longint;我不知道如何解決這個問題），最終數據幀的大約20個字節* 900m行= 18GB。 它可能在實踐中更多，並且在操作期間的峰值內存使用率高於此值（特別是，排序將花費最長時間，並使用大量內存）。

我在我的機器上嘗試了這個：我使用了大約30GB，放棄了等待完整排序，並對dist小於100的子集進行了排序。 花了不到5分鍾，大部分時間花在了加入上。

在一天結束時，你正在研究接近十億計算的計算; 如果你想以C的速度執行此操作而不必存儲所有成對數據（在pandas中使用直接方法就是這種情況），你很可能必須使用numpy數組在Cython中編寫代碼，和/或多。

更智能的方法是避免進行十億次計算，這包括了解您不需要打擾計算的距離。 這需要一些聰明的邏輯，但幸運的是，這是k-Nearest Neighbors的一個研究得很好的主題，它有專門針對這種問題設計的高效算法：

from sklearn.neighbors import NearestNeighbors
data = sites[['LAT','LON']].values
nbrs = NearestNeighbors(n_neighbors=10, algorithm='auto', metric = 'euclidean').fit(data)
distances, indices = nbrs.kneighbors(data)
indices

計算時間不到一秒。 恢復最近鄰居的名稱需要更長時間：

df = pd.DataFrame(indices, index = sites['STATION NAME'].values)
df.replace(dict(enumerate(sites['STATION NAME'].values)), inplace = True)

（實際上你可以通過使用.merge()方法和一些堆疊/卸載來大大提高速度，但是在這種情況下它會稍微復雜一些，因為你的數據包含重復數據。）