Question

由于显而易见的原因，我有两个不同大小的numpy数组，一个索引列和x y z坐标，另一个只包含坐标。（请忽略第一个序列号，我无法弄清楚格式化。）第二个数组少了no。坐标，我需要来自第一个数组的那些坐标的索引（atomID）。

Array1（带索引列）：

    serialNo. moleculeID atomID x y z

1 1 2 0 7.7590151 7.2925348 12.5933323
2 1 2 0 7.123642 6.1970949 11.5622416
3 1 6 0 6.944543 7.0390449 12.0713224
4 1 2 0 8.8900348 11.5477333 13.5633965
5 1 2 0 7.857268 12.8062735 13.4357052
6 1 6 0 8.2124357 12.1004238 14.0486889

Array2（只是坐标）：

x          y             z

7.7590151 7.2925348 12.5933323
7.123642 6.1970949 11.5622416
6.944543 7.0390449 12.0713224
8.8900348 11.5477333 13.5633965

具有索引列（atomID）的数组的索引为2,2,6,2,2和6.如何获取Array1和Array2中常见坐标的索引。我希望将2 2 6 2作为列表返回，然后将其与第二个数组连接起来。任何简单的想法？

更新

尝试使用以下代码，但它似乎无法正常工作。

import numpy as np

a = np.array([[4, 2.2, 5], [2, -6.3, 0], [3, 3.6, 8], [5, -9.8, 50]])

b = np.array([[2.2, 5], [-6.3, 0], [3.6, 8]])

print a
print b

for i in range(len(b)):
 for j in range(len(a)):
    if a[j,1]==b[i,0]:
        x = np.insert(b, 0, a[i,0], axis=1) #(input array, position to insert, value to insert, axis)
        #continue
    else:
        print 'not true'
print x

输出以下内容：

not true
not true
not true
not true
not true
not true
not true
not true
not true
[[ 3.   2.2  5. ]
 [ 3.  -6.3  0. ]
 [ 3.   3.6  8. ]]

但期望是：

    [[ 4.   2.2  5. ]
     [ 2.  -6.3  0. ]
     [ 3.   3.6  8. ]]

Answer 1

使用cdist -

进行两种简洁的矢量化方法

from scipy.spatial.distance import cdist

out = a[np.any(cdist(a[:,1:],b)==0,axis=1)]

或者，如果您不介意获取一些 voodoo-ish ，请np.einsum替换np.any -

out = a[np.einsum('ij->i',cdist(a[:,1:],b)==0)]

示例运行 -

In [15]: from scipy.spatial.distance import cdist

In [16]: a
Out[16]: 
array([[  4. ,   2.2,   5. ],
       [  2. ,  -6.3,   0. ],
       [  3. ,   3.6,   8. ],
       [  5. ,  -9.8,  50. ]])

In [17]: b
Out[17]: 
array([[ 2.2,  5. ],
       [-6.3,  0. ],
       [ 3.6,  8. ]])

In [18]: a[np.any(cdist(a[:,1:],b)==0,axis=1)]
Out[18]: 
array([[ 4. ,  2.2,  5. ],
       [ 2. , -6.3,  0. ],
       [ 3. ,  3.6,  8. ]])

In [19]: a[np.einsum('ij->i',cdist(a[:,1:],b)==0)]
Out[19]: 
array([[ 4. ,  2.2,  5. ],
       [ 2. , -6.3,  0. ],
       [ 3. ,  3.6,  8. ]])

Answer 2

numpy_indexed包（免责声明：我是它的作者）包含以优雅，高效/矢量化的方式解决此类问题的功能：

import numpy_indexed as npi
print(a[npi.contains(b, a[:, 1:])])

目前接受的答案让我觉得对于后一坐标不同的点不正确。这里的表现也应该大大改善;这个解决方案不仅被矢量化，而且最坏情况下的性能是NlogN，而不是当前接受的答案的二次时间复杂度。

Answer 3

这只是您问题的伪代码：

import numpy as np
for i in range(len(array2)):
    for element in array1:
        if array2[i]xyz == elementxyz: #compare the coordinates of the two elements
            np.insert(array2[i], 0, element_coord) #insert the atomid at the beginning of the coordinate array
            break

Answer 4

使用列表而不是数组来表示np.insert的值。

import numpy as np

a = np.array([[4, 2.2, 5], [2, -6.3, 0], [3, 3.6, 8], [5, -9.8, 50]])

b = np.array([[2.2, 5], [-6.3, 0], [3.6, 8]])

print a
print b
x = []

for i in range(len(b)):
 for j in range(len(a)):
    if a[j,1]==b[i,0]:
        x.append(a[j,0])
    else:
        x = x
print np.insert(b,0,x,axis=1)

将输出：

[[ 4.   2.2  5. ]
 [ 2.  -6.3  0. ]
 [ 3.   3.6  8. ]]

如何在两个不同大小的numpy数组之间进行比较，并使用公共元素返回索引列？

4 个答案: