numpy 方法连接两个网格网格及其结果数组

Question

考虑两个 n 维，可能重叠，numpy meshgrids，比如说

m1 = (x1, y1, z1, ...)
m2 = (x2, y2, z2, ...)

在m1 和m2 内没有重复的坐标元组。每个meshgrid都有一个结果数组，可能是不同函数产生的：

r1 = f1(m1)
r2 = f2(m2)

这样f1(m) != f2(m)。现在我想加入这两个meshgrids 及其结果数组，例如m=m1&m2 和r=r1&r2（其中& 表示某种联合），这样m 中的坐标元组仍然是排序的，r 中的值仍然对应于原始坐标元组。新创建的坐标元组应该是可识别的（例如具有特殊值）。

为了详细说明我所追求的，我有两个例子可以用简单的for 和if 语句来做我想做的事。这是一个 1D 示例：

x1 = [1, 5, 7]
r1 = [i**2 for i in x1]

x2 = [2, 4, 6]
r2 = [i*3 for i in x2]

x,r = list(zip(*sorted([(i,j) for i,j in zip(x1+x2,r1+r2)],key=lambda x: x[0])))

给了

x = (1, 2, 4, 5, 6, 7)
r = (1, 6, 12, 25, 18, 49)

对于 2D，它开始变得相当复杂：

import numpy as np
a1 = [1, 5, 7]
b1 = [2, 5, 6]

x1,y1 = np.meshgrid(a1,b1)
r1 = x1*y1

a2 = [2, 4, 6]
b2 = [1, 3, 8]

x2, y2 = np.meshgrid(a2,b2)
r2 = 2*x2

a = [1, 2, 4, 5, 6, 7]
b = [1, 2, 3, 5, 6, 8]

x,y = np.meshgrid(a,b)

r = np.ones(x.shape)*-1

for i in range(x.shape[0]):
    for j in range(x.shape[1]):
        if   x[i,j] in a1 and y[i,j] in b1:
            r[i,j] = r1[a1.index(x[i,j]),b1.index(y[i,j])]

        elif x[i,j] in a2 and y[i,j] in b2:
            r[i,j] = r2[a2.index(x[i,j]),b2.index(y[i,j])]

这给出了所需的结果，新坐标对的值为-1：

x=
[[1 2 4 5 6 7]
 [1 2 4 5 6 7]
 [1 2 4 5 6 7]
 [1 2 4 5 6 7]
 [1 2 4 5 6 7]
 [1 2 4 5 6 7]]
y=
[[1 1 1 1 1 1]
 [2 2 2 2 2 2]
 [3 3 3 3 3 3]
 [5 5 5 5 5 5]
 [6 6 6 6 6 6]
 [8 8 8 8 8 8]]
r=
[[ -1.   4.   4.  -1.   4.  -1.]
 [  2.  -1.  -1.   5.  -1.   6.]
 [ -1.   8.   8.  -1.   8.  -1.]
 [ 10.  -1.  -1.  25.  -1.  30.]
 [ 14.  -1.  -1.  35.  -1.  42.]
 [ -1.  12.  12.  -1.  12.  -1.]]

但是随着维度和数组大小的增加，这也会很快变慢。所以最后的问题是：如何仅使用numpy 函数来完成。如果不可能，在python 中实现此功能的最快方法是什么。如果无论如何相关，我更喜欢使用 Python 3。请注意，我在示例中使用的函数并不是我实际使用的函数。

Answer 1

我们可以使用一些掩码来替换A in B 部分，从而为我们提供1D 掩码。然后，我们可以将这些掩码与np.ix_ 一起使用以扩展到所需的维数。

因此，对于2D 的情况，应该是这样的 -

# Initialize o/p array
r_out = np.full([len(a), len(b)],-1)           

# Assign for the IF part
mask_a1 = np.in1d(a,a1)
mask_b1 = np.in1d(b,b1)
r_out[np.ix_(mask_b1, mask_a1)] = r1.T

# Assign for the ELIF part
mask_a2 = np.in1d(a,a2)
mask_b2 = np.in1d(b,b2)
r_out[np.ix_(mask_b2, mask_a2)] = r2.T

a 可以这样创建 -

a = np.concatenate((a1,a2))
a.sort()

同样，对于b。

此外，我们可以使用索引而不是掩码来与np.ix_ 一起使用。同样，我们可以使用np.searchsorted。因此，代替掩码np.in1d(a,a1)，我们可以使用np.searchsorted(a,a1) 获得相应的索引，以此类推其余的掩码。这应该会快很多。

---

对于3D 的情况，我假设我们会有另一个数组，比如c。因此，初始化部分将涉及使用len(c)。将有一个对应于c 的掩码/索引数组，因此还有一个术语到np.ix_，并且将有r1 和r2 的转置。

Answer 2

Divakar 的回答正是我所需要的。但是，我想仍然尝试该答案中的第二个建议，并且最重要的是我做了一些分析。我认为结果可能对其他人来说很有趣。这是我用于分析的代码：

import numpy as np
import timeit
import random

def for_join_2d(x1,y1,r1, x2,y2,r2):
    """
    The algorithm from the question.
    """

    a = sorted(list(x1[0,:])+list(x2[0,:]))
    b = sorted(list(y1[:,0])+list(y2[:,0]))

    x,y = np.meshgrid(a,b)
    r = np.ones(x.shape)*-1

    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            if   x[i,j] in a1 and y[i,j] in b1:
                r[i,j] = r1[a1.index(x[i,j]),b1.index(y[i,j])]

            elif x[i,j] in a2 and y[i,j] in b2:
                r[i,j] = r2[a2.index(x[i,j]),b2.index(y[i,j])]
    return x,y,r


def mask_join_2d(x1,y1,r1,x2,y2,r2):
    """
    Divakar's original answer.
    """
    a = np.sort(np.concatenate((x1[0,:],x2[0,:])))
    b = np.sort(np.concatenate((y1[:,0],y2[:,0])))

    # Initialize o/p array
    x,y = np.meshgrid(a,b)
    r_out = np.full([len(a), len(b)],-1)           

    # Assign for the IF part
    mask_a1 = np.in1d(a,a1)
    mask_b1 = np.in1d(b,b1)
    r_out[np.ix_(mask_b1, mask_a1)] = r1.T

    # Assign for the ELIF part
    mask_a2 = np.in1d(a,a2)
    mask_b2 = np.in1d(b,b2)
    r_out[np.ix_(mask_b2, mask_a2)] = r2.T

    return x,y,r_out


def searchsort_join_2d(x1,y1,r1,x2,y2,r2):
    """
    Divakar's second suggested solution using searchsort.
    """

    a = np.sort(np.concatenate((x1[0,:],x2[0,:])))
    b = np.sort(np.concatenate((y1[:,0],y2[:,0])))

    # Initialize o/p array
    x,y = np.meshgrid(a,b)
    r_out = np.full([len(a), len(b)],-1)           

    #the IF part
    ind_a1 = np.searchsorted(a,a1)
    ind_b1 = np.searchsorted(b,b1)
    r_out[np.ix_(ind_b1,ind_a1)] = r1.T

    #the ELIF part
    ind_a2 = np.searchsorted(a,a2)
    ind_b2 = np.searchsorted(b,b2)
    r_out[np.ix_(ind_b2,ind_a2)] = r2.T

    return x,y,r_out

##the profiling code:
if __name__ == '__main__':

    N1 = 100
    N2 = 100

    coords_a = [i for i in range(N1)]
    coords_b = [i*2 for i in range(N2)]

    a1 = random.sample(coords_a, N1//2)
    b1 = random.sample(coords_b, N2//2)

    a2 = [i for i in coords_a if i not in a1]
    b2 = [i for i in coords_b if i not in b1]

    x1,y1 = np.meshgrid(a1,b1)
    r1 = x1*y1
    x2,y2 = np.meshgrid(a2,b2)
    r2 = 2*x2

    print("original for loop")
    print(min(timeit.Timer(
        'for_join_2d(x1,y1,r1,x2,y2,r2)',
        setup = 'from __main__ import for_join_2d,x1,y1,r1,x2,y2,r2',
    ).repeat(7,1000)))

    print("with masks")
    print(min(timeit.Timer(
        'mask_join_2d(x1,y1,r1,x2,y2,r2)',
        setup = 'from __main__ import mask_join_2d,x1,y1,r1,x2,y2,r2',
    ).repeat(7,1000)))

    print("with searchsort")
    print(min(timeit.Timer(
        'searchsort_join_2d(x1,y1,r1,x2,y2,r2)',
        setup = 'from __main__ import searchsort_join_2d,x1,y1,r1,x2,y2,r2',
    ).repeat(7,1000)))

对于每个函数，我使用了 7 组 1000 次迭代，并选择了最快的一组进行评估。两个 10x10 数组的结果是：

original for loop
0.5114614190533757

with masks
0.21544912096578628

with searchsort
0.12026709201745689

对于两个 100x100 数组，它是：

original for loop
247.88183582702186

with masks
0.5245905339252204

with searchsort
0.2439237720100209

对于大型矩阵，使用numpy 功能毫无疑问会产生巨大的差异，实际上searchsort 和索引而不是屏蔽大约可以将运行时间减半。