在pandas中优化循环

Question

我有一个数据框，有两列是点的坐标。 如果某个点位于特定位置，我需要用特定值填充一列（全部为无）。该位置和标签存储在另一个df中

这不容易解释，但我希望通过一个例子可以清楚地说明： DF1

   latitude  longitude  LABEL
0    1.3       2.7      None
1    3.5       3.6      None
2    2.8       3.0      None
3    9.7       1.9      None
4    6.2       5.7      None
5    1.7       3.4      None
6    3.5       1.4      None
7    2.7       6.6      None
8    1.7       2.7      None
9    1.3       1.3      None

DF2

   minlat     maxlat    minlong   maxlong  STRING
0    1.0       2.0        1.0       3.0     AAA
1    3.0       4.0        1.0       2.0     BBB
2    3.0       4.0        3.0       4.0     CCC
3    5.0       7.0        2.0       3.0     DDD

最终结果是：

   latitude  longitude  LABEL
0    1.3       2.7      AAA
1    3.5       3.6      CCC
2    2.8       3.0      None
3    9.7       1.9      None
4    6.2       5.7      None
5    1.7       3.4      None
6    3.5       1.4      BBB
7    2.7       6.6      None
8    1.7       2.7      AAA
9    1.3       1.3      AAA

目前的代码是：

for i in range(len(df2)-1):
DF1.loc[(DF1['latitude']>=DF2.loc[i:i,'minlat'].at[i]) & (DF1['latitude']<DF2.loc[i:i,'maxlat'].at[i]) &
   (DF1['longitude']>=DF2.loc[i:i,'minlong'].at[i]) & (DF1['longitude']<DF2.loc[i:i,'maxlong'].at[i]),'LABEL'] = DF2.loc[i:i,'STRING'].at[i]

屏幕以获得更好的缩进：

因此，对于DF2的每一行，我检查DF1的值是否在中间，并分配一个字符串

但是像这样需要很多时间。你对我能做什么有什么建议吗？ 我的问题是，必须用DF2的每一行检查数字_1的每个值，而不仅仅是用具有相同索引的行

编辑：我正在尝试其他方法：

(二)

for i in range(len(xlsx_fact_maneuver_specialareas)-1):
    minLat=DF2.loc[i:i,'minLat'].at[i]
    maxLat=DF2.loc[i:i,'maxLat'].at[i]
    minLong=DF2.loc[i:i,'maxLat'].at[i]
    maxLong=DF2.loc[i:i,'maxLong'].at[i]
    DF1.loc[(DF1['latitude']>=minLat) & (DF1['latitude']<maxLat) &
   (DF1['longitude']>=minLong) & (DF1['longitude']<maxLong),'LABEL'] = DF2.loc[i:i,'STRING'].at[i]

这让我在本地感觉不太好，但当我在机器上尝试时，感觉更好

及

for i in range(len(xlsx_fact_maneuver_specialareas)-1):
    minLat=DF2.loc[i:i,'minLat'].at[i]
    maxLat=DF2.loc[i:i,'maxLat'].at[i]
    minLong=DF2.loc[i:i,'maxLat'].at[i]
    maxLong=DF2.loc[i:i,'maxLong'].at[i]
    DF1 = DF1.assign(
        label =  np.select(
          [(DF1['latitude']>=minLat) & (DF1['latitude']<maxLat) & (DF1['longitude']>=minLong) & (DF1['longitude']<maxLong)],
          [DF2.loc[i:i,'STRING'].at[i]],
          [None]))

这让我在本地感觉更好，但在机器上感觉更差

Answer 1

这个问题并不特别适合在Pandas本身中解决，因为没有简单的原语来处理您需要进行的计算。 更好的方法是转移到NumPy或Numba域，在较低级别上解决问题

我将提供生成最后一列的函数，假设将最后一列复制到数据帧中相对容易

最初的做法是：

def locate_in_regions_OP(points, regions):
    result = np.full(len(points), None, dtype=object)
    for i in range(len(regions) - 1):
        result[
            (points['lat'] >= regions.loc[i:i, 'min_lat'].at[i])
            & (points['lat'] < regions.loc[i:i,'max_lat'].at[i])
            & (points['lon'] >= regions.loc[i:i, 'min_lon'].at[i])
            & (points['lon'] < regions.loc[i:i, 'max_lon'].at[i])
        ] = regions.loc[i:i, 'lbl'].at[i]
    return result

这将为最后一列生成正确的结果。 （OP中提出的其他方法要么不相关，要么仅对仅使用一次的数量使用独立赋值，要么我没有设法让它们工作）

一种相对简单的方法涉及广播，在@PierreD answer中介绍，可以进一步简化为：

import numpy as np


def locate_in_regions_bc(points, regions):
    pos_arr = points[['lat', 'lon']].values
    min_arr = regions[['min_lat', 'min_lon']].values
    max_arr = regions[['max_lat', 'max_lon']].values
    mask = (
        np.all(min_arr[None, :, :] <= pos_arr[:, None, :], axis=2)
        & np.all(pos_arr[:, None, :] < max_arr[None, :, :], axis=2))
    has_any = np.any(mask, axis=1)
    first = np.argmax(mask, axis=1)
    result = np.full(len(points), None, dtype=object)
    result[has_any] = regions.loc[first[has_any], 'lbl'].values
    return result

在假设每个位置只属于一个区域的情况下，这可以稍微简化：

import numpy as np


def locate_in_regions_bcu(points, regions):
    pos_arr = points[['lat', 'lon']].values
    min_arr = regions[['min_lat', 'min_lon']].values
    max_arr = regions[['max_lat', 'max_lon']].values
    mask = (
        np.all(min_arr[None, :, :] <= pos_arr[:, None, :], axis=2)
        & np.all(pos_arr[:, None, :] < max_arr[None, :, :], axis=2))
    pos, lbl = np.where(mask)    
    result = np.full(len(points), None, dtype=object)
    result[pos] = regions.loc[lbl, 'lbl'].values
    return result

但是，有大量不必要的内存分配和比较正在进行。 一种更快的方法是使用Numba显式循环，您可以显式添加短路。该守则的内容如下：

import numpy as np
import numba as nb


def locate_in_regions_nb(points, regions):
    pos_arr = points[['lat', 'lon']].values
    min_arr = regions[['min_lat', 'min_lon']].values
    max_arr = regions[['max_lat', 'max_lon']].values
    found_arr = _locate_in_regions_nb(pos_arr, min_arr, max_arr)
    mask = found_arr >= 0
    result = np.full(len(points), None, dtype=object)
    result[mask] = regions.loc[found_arr[mask], 'lbl'].values
    return result
    

@nb.jit
def _locate_in_regions_nb(pos_arr, min_arr, max_arr):
    n, l = pos_arr.shape
    m = min_arr.shape[0]
    found_arr = np.full((n,), -1)
    for i in range(n):
        for j in range(m):
            contained = True
            for k in range(l):
                if min_arr[j, k] > pos_arr[i, k] or pos_arr[i, k] >= max_arr[j, k]:
                    contained = False
                    break
            if contained:
                found_arr[i] = j
                break
    return found_arr

使用稍微干净但在其他方面具有可比性的输入：

import numpy as np
import pandas as pd
import numba as nb


df1 = pd.DataFrame(
    columns=['lat', 'lon', 'lbl'],
    data=[
        [1.3, 2.7, None],
        [3.5, 3.6, None],
        [2.8, 3.0, None],
        [9.7, 1.9, None],
        [1.7, 3.4, None],
        [3.5, 1.4, None],
        [2.7, 6.6, None],
        [1.7, 2.7, None],
        [1.3, 1.3, None],
    ])

df2 = pd.DataFrame(
    columns=['min_lat', 'max_lat', 'min_lon', 'max_lon', 'lbl'],
    data=[
        [1.0, 2.0, 1.0, 3.0, 'AAA'],
        [3.0, 4.0, 1.0, 2.0, 'BBB'],
        [3.0, 4.0, 3.0, 4.0, 'CCC'],
        [5.0, 7.0, 2.0, 3.0, 'DDD'],
    ])


print(df1)
#    lat  lon   lbl
# 0  1.3  2.7  None
# 1  3.5  3.6  None
# 2  2.8  3.0  None
# 3  9.7  1.9  None
# 4  1.7  3.4  None
# 5  3.5  1.4  None
# 6  2.7  6.6  None
# 7  1.7  2.7  None
# 8  1.3  1.3  None

print(df2)
#    min_lat  max_lat  min_lon  max_lon  lbl
# 0      1.0      2.0      1.0      3.0  AAA
# 1      3.0      4.0      1.0      2.0  BBB
# 2      3.0      4.0      3.0      4.0  CCC
# 3      5.0      7.0      2.0      3.0  DDD

在所有情况下都可以获得预期输出：

print(locate_in_regions_OP(df1, df2))
# ['AAA' 'CCC' None None None 'BBB' None 'AAA' 'AAA']
print(locate_in_regions_bc(df1, df2))
# ['AAA' 'CCC' None None None 'BBB' None 'AAA' 'AAA']
print(locate_in_regions_bcu(df1, df2))
# ['AAA' 'CCC' None None None 'BBB' None 'AAA' 'AAA']
print(locate_in_regions_nb(df1, df2))
# ['AAA' 'CCC' None None None 'BBB' None 'AAA' 'AAA']

虽然产生代表问题的任意大的输入并不容易，但天真的时间安排表明Numba方法要快得多：

np.random.seed(0)
df3 = pd.DataFrame(
    columns=['lat', 'lon', 'lbl'],
    data=np.random.random((1000000, 3)) * 5)
df3.loc[:, 'lbl'] = None


%timeit locate_in_regions_OP(df3, df2)
# 209 ms ± 7.99 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit locate_in_regions_bc(df3, df2)
# 161 ms ± 2.92 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit locate_in_regions_bcu(df3, df2)
# 115 ms ± 2.43 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit locate_in_regions_nb(df3, df2)
# 66.6 ms ± 3.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

将此操作矢量化的一个解决方案是使用Numpy及其出色的广播功能。这为中小型数据帧提供了一个快速的解决方案，但它会随着mask的O[n*m]（对于df1的n行和df2行的m行的mask而增长（在时间和内存上），因此最终对于大型数据帧来说速度会变慢

a = df1[['latitude', 'longitude']].values
vmin = df2[['minlat', 'minlong']].values
vmax = df2[['maxlat', 'maxlong']].values
mask = (vmin[None, :, :] <= a[:, None, :]).all(2) & (a[:, None, :] <= vmax[None, :, :]).all(2)
has_any = mask.any(1)
first = mask.argmax(axis=1)
label = np.full(len(df1), None, dtype=object)
label[has_any] = df2['STRING'].values[first[has_any]]

>>> df1.assign(LABEL=label)
   latitude  longitude LABEL
0       1.3        2.7   AAA
1       3.5        3.6   CCC
2       2.8        3.0  None
3       9.7        1.9  None
4       6.2        5.7  None
5       1.7        3.4  None
6       3.5        1.4   BBB
7       2.7        6.6  None
8       1.7        2.7   AAA
9       1.3        1.3   AAA

解释

关键部分是mask的构造。有必要对其进行细分，以了解其机制及其如何使用Numpy的广播：

>>> vmin[None, :, :] <= a[:, None, :]
[[[ True  True]
  [False  True]
  [False False]
  [False  True]]

 [[ True  True]
  [ True  True]
  [ True  True]
  [False  True]]
 ...
 [[ True  True]
  [False  True]
  [False False]
  [False False]]]

如您所见，上面将a和vmin之间的所有比较扩展到第三维。然后我们用逻辑“所有第三轴（经度和纬度）都必须为真”投射回2D：

>>> (vmin[None, :, :] <= a[:, None, :]).all(2)
[[ True False False False]
 [ True  True  True False]
 [ True False False False]
 [ True  True False False]
 [ True  True  True  True]
 [ True False False False]
 [ True  True False False]
 [ True False False False]
 [ True False False False]
 [ True False False False]]

以上表示高于df2.iloc[j]最小值的所有点df1.iloc[i]为...[i, j]

我们对vmax做了同样的处理，得到的mask是df1.iloc[i]的所有点都在df2.iloc[j]的边界框中

接下来的两位是has_any和first。前者表示df1中的哪些点至少位于一个边界框中。后者是第一个这样的边界框（如df2中的索引）

其余的都是不言自明的

注释

请注意，这使用了O[n*m]比较（对于df1的n行和df2的m行），这对于大型矩阵来说可能太慢（尽管因为它是矢量化的，所以对于中型矩阵来说速度非常快）

对于大型矩阵，更好的方法包括排序或使用KD树。见this other answer

Answer 3

这是另一个答案，这个答案使用了^{}。如果边界框之间没有太多重叠，并且它们的“半径”分布也不太“随意”（radius.max() / np.median(radius)不太大，例如在1和2之间），则该方法特别有效

它适用于p-范数1（曼哈顿）或2（欧几里德），尽管在实践中{}更快，因为平均每个点看到的候选点更少（圆的总面积小于钻石的总面积）

为什么这么快KD-trees非常适合处理这类问题。它们通过沿维度和中点在每个节点上分割空间来划分空间。一旦构建了它们，由于它们提供了分而治之的方法，查询它们的速度很快

关键功能如下：

def find(kd_boxes, kd_points, bb, r, p):
    # bb: m bounding boxes, in the form [x0,x1,y0,y1]
    # find one bb for each point (the first one), or m if none
    xy = kd_points.data
    found = np.full(len(xy), len(bb), dtype=int)
    cand = pad_jagged(kd_points.query_ball_tree(kd_boxes, r * 1.001, p=p))

    remaining = None
    for c in range(cand.shape[1]):
        remaining = np.nonzero(cand[:, c] >= 0)[0] if remaining is None else remaining[np.nonzero(cand[remaining, 1] >= 0)[0]]
        if remaining.size == 0:
            break
        i = remaining
        j = cand[i, c]
        ok = (bb[j, 0::2] <= xy[i]).all(1) & (xy[i] <= bb[j, 1::2]).all(1)
        k = np.nonzero(ok)[0]
        found[i[k]] = j[k]
        remaining = remaining[~ok]
    return found

在调用此函数之前，我们为每个边界框计算一个“半径”，该半径是对角线p范数的一半。然后，我们使用总体最大半径r作为KDTree查询的最大距离。KDTree查询（kd_points.query_ball_tree()）有效地筛选所有边界框中心，并在一次调用中查找半径内的所有边界框中心。这是实际匹配的超集，但速度很快，大大减少了搜索空间。然后，过滤实际上在边界框中的点，并跟踪每个点的（第一个）匹配边界框

作为优化（这里没有实现），我们可以考虑边界框的大小（^ {< CD8>}数组）。如果判断差异太大，则可以将边界框分为两组（例如围绕np.median(radius)），并且可以对每一半进行相同的搜索（如果必要，再次递归）

对于OP示例，在准备以更易于使用的形式（所有Numpy数组）获取中心、边界框和半径后，该查询返回：

# preparation
xy = df1[['longitude', 'latitude']].values
bb = df2[['minlong', 'maxlong', 'minlat', 'maxlat']].values
centers = bb @ np.array([[1,1,0,0], [0,0,1,1]]).T / 2
radius = np.linalg.norm(bb @ np.array([[-1,1,0,0], [0,0,-1,1]]).T, axis=1) / 2
kd_boxes = KDTree(centers)
kd_points = KDTree(xy)

# query
>>> kd_points.query_ball_tree(kd_boxes, radius.max() * 1.001)
[[0], [2], [2], [], [], [], [1], [], [0], [0]]

使用曼哈顿范数可以得到类似的结果（在本例中，相同）：

radius = bb @ np.array([-1,1,-1,1]) / 2
>>> kd_points.query_ball_tree(kd_boxes, radius.max() * 1.001, p=1)
[[0], [2], [2], [], [], [], [1], [], [0], [0]]

下图中绘制了这两种解决方案，其中突出显示了所有候选边界框及其中心距离r内的实际面积：

请注意，在这两种情况下，点#2是如何错误地包含在bbox#2中的。第一步，使用KD树，只是一个过滤步骤。下一步是检查每个点的候选是否实际包含该点。使用该筛选（用于ok数组的表达式），解决方案是：

让我们看看如果我们有更多的点和边界框（可能有重叠），会发生什么：

def gen_example(n, m, seed=None):
    # let's make a more complex problem, with overlaps
    if seed is not None:
        np.random.seed(seed)
    df1 = pd.DataFrame({
        'latitude': np.random.uniform(0, 1 + 1 / np.sqrt(m), size=n),
        'longitude': np.random.uniform(0, 1 + 1 / np.sqrt(m), size=n),
    })

    df2 = pd.DataFrame({
        'minlat': np.random.uniform(size=m),
        'maxlat': np.random.uniform(.5, 1, size=m) / np.sqrt(m),
        'minlong': np.random.uniform(size=m),
        'maxlong': np.random.uniform(.5, 1, size=m) / np.sqrt(m),
        'STRING': [f'b_{i}' for i in range(m)],
    })
    df2['maxlat'] += df2['minlat']
    df2['maxlong'] += df2['minlong']
    return df1, df2

对于df1, df2 = gen_example(32, 12, 0)，相应的图片为：

候选项（查询结果，此处为p=2）为：

>>> kd_cand = kd_points.query_ball_tree(kd_boxes, r * 1.001, p=2)
[[], [], [9], [], [], [10], [], [], [], [], [], [9], [10], [], [11], [11], [], [2, 6], [8], [8], [], [4], [7, 9], [4], [3, 5], [2, 4], [0], [], [7, 9], [7, 9], [0, 3, 5], [4]]

因为这是一个锯齿状数组，所以我们将其转换为带有fill_value=-1的矩形数组：

def pad_jagged(a, fill_value=-1):
    lens = np.array([len(r) for r in a])
    v = np.array([e for r in a for e in r])
    w = np.max(lens)
    return np.insert(
        v,
        np.repeat(np.cumsum(lens), w - lens),
        fill_value
    ).reshape(-1, w)

对于上面的填充阵列，这将提供：

>>> pad_jagged(kd_cand)
[[-1 -1 -1]
 [ 8 -1 -1]
 [ 9 -1 -1]
 ...
 [ 7  9 -1]
 [ 0  3  5]
 [ 4 -1 -1]]

现在，我们迭代这个数组的列，但是以贪婪的方式从以前的迭代中删除任何成功的匹配项（这就是remaining的原因）

处理预处理等的其他功能包括：

def find_regions(xy, bb, p=2):
    # xy: numpy array (n, 2)
    # bb: numpy array (m, 4): the four columns are xmin, xmax, ymin, ymax
    # for each point in xy, return the index of a region that contains it, or -1
    centers = bb @ np.array([[1,1,0,0], [0,0,1,1]]).T / 2
    assert p in {1,2}, "only 1- or 2-norm supported"
    radius = bb @ np.array([-1,1,-1,1]) / 2 if p == 1 else np.linalg.norm(bb @ np.array([[-1,1,0,0], [0,0,-1,1]]).T, axis=1) / 2
    kd_boxes = KDTree(centers)
    kd_points = KDTree(xy)
    return find(kd_boxes, kd_points, bb, radius.max(), p=p)

def find_region_label(df1, df2, nomatch=None, p=2):
    found = find_regions(
        df1[['longitude', 'latitude']].values,
        df2[['minlong', 'maxlong', 'minlat', 'maxlat']].values,
        p=p,
    )
    lbl = np.r_[df2['STRING'].values, [nomatch]]
    return lbl[found]

关于OP示例：

>>> df1.assign(LABEL=find_region_label(df1, df2))
   latitude  longitude LABEL
0       1.3        2.7   AAA
1       3.5        3.6   CCC
2       2.8        3.0  None
3       9.7        1.9  None
4       6.2        5.7  None
5       1.7        3.4  None
6       3.5        1.4   BBB
7       2.7        6.6  None
8       1.7        2.7   AAA
9       1.3        1.3   AAA

速度测试

现在进行速度测试，其中点数n和边界框数m都较大：

df1, df2 = gen_example(100_000, 10_000, 0)

与@norok2's numba solution的速度比较（稍微调整以跟随OP的列名）：

a = %timeit -o find_region_label(df1, df2)
# 222 ms ± 904 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

b = %timeit -o locate_in_regions_nb(df1, df2)
# 5.38 s ± 40.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

>>> b.average / a.average
24.255

这是该数据的24倍加速。

验证我们是否得到相同的解决方案：

sa = find_region_label(df1, df2)
sb = locate_in_regions_nb(df1, df2)
>>> (sa == sb).all()
True