熊猫:在两个数据帧上使用函数进行矢量化

Question

我在pandas中实现矢量化时遇到了麻烦.让我先说一下我是矢量化的全新手,所以我极有可能得到一些错误的语法.

假设我有两只熊猫数据帧.

数据帧1描述了具有半径R的一些圆的x,y坐标,具有唯一ID.

>>> data1 = {'ID': [1, 2], 'x': [1, 10], 'y': [1, 10], 'R': [4, 5]}
>>> df_1=pd.DataFrame(data=data1)
>>>
>>> df_1
   ID  x   y   R
   1   1   1   4
   2   10  10  5

Dataframe 2描述了某些点的x,y坐标,也有唯一ID.

>>> data2 = {'ID': [3, 4, 5], 'x': [1, 3, 9], 'y': [2, 5, 9]}
>>> df_2=pd.DataFrame(data=data2)
>>>
>>> df_2
   ID  x  y
   3   1  2
   4   3  5
   5   9  9

现在,想象一下绘制2D平面上的圆和点.一些点将驻留在圆圈内.见下图.

我想要做的就是在df_2中创建一个名为"host_circle"的新列,它指示每个点所在的圆的ID.如果粒子不在一个圆中,则该值应为"None".

我想要的输出是

>>> df_2
   ID  x  y   host_circle
   3   1  2   1 
   4   3  5   None 
   5   9  9   2

首先,定义一个函数,检查给定的粒子(x2,y2)是否位于给定的圆(x1,y1,R1,ID_1)内.如果是,则返回圆圈的ID; 否则,返回无.

>>> def func(x1,y1,R1,ID_1,x2,y2):
...     dist = np.sqrt( (x1-x2)**2 + (y1-y2)**2 )
...     if dist < R:
...         return ID_1
...     else:
...        return None

接下来,实际矢量化.我有点迷失在这里.我认为应该是这样的

df_2['host']=func(df_1['x'],df_1['y'],df_1['R'],df_1['ID'],df_2['x'],df_2['y'])

但这只会引发错误.有人能帮我吗？

最后一点说明:我正在使用的实际数据非常大; 数千万行.速度至关重要,因此我为什么要尝试进行矢量化工作.

Answer 1

Numba v1

您可能必须安装numba有

pip install numba

然后numba通过njit函数装饰器使用s jit编译器

from numba import njit

@njit
def distances(point, points):
  return ((points - point) ** 2).sum(1) ** .5

@njit
def find_my_circle(point, circles):
  points = circles[:, :2]
  radii = circles[:, 2]
  dist = distances(point, points)
  mask = dist < radii
  i = mask.argmax()
  return i if mask[i] else -1

@njit
def find_my_circles(points, circles):
  n = len(points)
  out = np.zeros(n, np.int64)
  for i in range(n):
    out[i] = find_my_circle(points[i], circles)
  return out

ids = np.append(df_1.ID.values, np.nan)

i = find_my_circles(points, df_1[['x', 'y', 'R']].values)
df_2['host_circle'] = ids[i]

df_2

   ID  x  y  host_circle
0   3  1  2          1.0
1   4  3  5          NaN
2   5  9  9          2.0

这会逐行迭代...这意味着它一次尝试找到主机圈.现在,那部分仍然是矢量化的. 和循环应该是非常快的.巨大的好处是你不会占用大量的内存.

Numba v2

当它找到主机时,这个更环路但是短路

from numba import njit

@njit
def distance(a, b):
  return ((a - b) ** 2).sum() ** .5

@njit
def find_my_circles(points, circles):
  n = len(points)
  m = len(circles)

  out = -np.ones(n, np.int64)

  centers = circles[:, :2]
  radii = circles[:, 2]

  for i in range(n):
    for j in range(m):
      if distance(points[i], centers[j]) < radii[j]:
        out[i] = j
        break

  return out

ids = np.append(df_1.ID.values, np.nan)

i = find_my_circles(points, df_1[['x', 'y', 'R']].values)
df_2['host_circle'] = ids[i]

df_2

矢量化

但仍然有问题

c = ['x', 'y']
centers = df_1[c].values
points = df_2[c].values
radii = df_1['R'].values

i, j = np.where(((points[:, None] - centers) ** 2).sum(2) ** .5 < radii)

df_2.loc[df_2.index[i], 'host_circle'] = df_1['ID'].iloc[j].values

df_2

   ID  x  y  host_circle
0   3  1  2          1.0
1   4  3  5          NaN
2   5  9  9          2.0

说明

距离圆心的任何点的距离是

((x1 - x0) ** 2 + (y1 - y0) ** 2) ** .5

如果我将我的一个数组扩展到第三维,我可以使用广播

points[:, None] - centers

array([[[ 0,  1],
        [-9, -8]],

       [[ 2,  4],
        [-7, -5]],

       [[ 8,  8],
        [-1, -1]]])

这就是矢量差异的所有六种组合.现在计算距离.

((points[:, None] - centers) ** 2).sum(2) ** .5

array([[ 1.        , 12.04159458],
       [ 4.47213595,  8.60232527],
       [11.3137085 ,  1.41421356]])

这是所有6种距离组合,我可以与半径进行比较,看看哪些是在圆圈内

((points[:, None] - centers) ** 2).sum(2) ** .5 < radii

array([[ True, False],
       [False, False],
       [False,  True]])

好的,我想找到True值的位置.这是一个完美的用例np.where.它会给我两个数组,第一个是行位置,第二个是这些True值所在的列位置.事实证明,行位置是points和列位置是圆圈.

i, j = np.where(((points[:, None] - centers) ** 2).sum(2) ** .5 < radii)

现在我只需要切片df_2以i某种方式和值赋给它,我从获得df_1使用j某种方式......但我发现上面.