python - 计算按前两列中的索引分组的 numpy 数组条目的第 N 列的总和？

Question

我想循环以下 check_matrix ，以便代码识别第一个和第二个元素是否为 1 和 1 或1 和 2 等？然后，对于每个单独的类对，即 1,1 或 1,2 或 2,2，代码应存储在新矩阵中，最后一个元素(在本例中索引为 8)乘以 exp(-i*q(check_matrix[k][2:5]-check_matrix[k][5:8])) ，其中 i 是 iota(复数)，k 是 check_matrix 上的运行索引，q 是定义如下的向量。所以有 20 个 q 向量。

import numpy as np

q= []
for i in np.linspace(0, 10, 20):
    q.append(np.array((0, 0, i)))
q = np.array(q)

check_matrix = np.array([[1, 1, 0, 0, 0, 0,    0,       -0.7977, -0.243293],
                         [1, 1, 0, 0, 0, 0,    0,       1.5954,  0.004567],
                         [1, 2, 0, 0, 0, -1,   0,       0,       1.126557],
                         [2, 1, 0, 0, 0, 0.5,  0.86603, 1.5954,  0.038934],
                         [2, 1, 0, 0, 0, 2,    0,       -0.7977, -0.015192],
                         [2, 2, 0, 0, 0, -0.5, 0.86603, 1.5954,  0.21394]])

这意味着原则上我必须有 20 个形状为 2x2 的矩阵，对应于每个 q 向量。

目前，我的代码只给出了一个矩阵，它似乎是最后一个矩阵，即使我附加在矩阵中。我的代码如下所示，

for i in range(2):
    i = i+1
    for j in range(2):
        j= j +1
        j_list = []
        Matrices = []
        for k in range(len(check_matrix)):
            if check_matrix[k][0] == i and check_matrix[k][1] == j:
                j_list.append(check_matrix[k][8]*np.exp(-1J*np.dot(q,(np.subtract(check_matrix[k][2:5],check_matrix[k][5:8])))))
                j_11 = np.sum(j_list)
                I_matrix[i-1][j-1] = j_11
                Matrices.append(I_matrix)

I_matrix 定义如下:

I_matrix= np.zeros((2,2),dtype=np.complex_)

目前我得到以下输出。

Matrices = [array([[-0.66071446-0.77603624j, -0.29038112+2.34855023j],         [-0.31387562-0.08116629j,  4.2788    +0.j        ]])]

但是，我希望获得与每个 q 值对应的矩阵，这意味着在这种情况下总共应该有 20 个矩阵，其中每个 2x2 矩阵元素将包含总和，使得元素属于1,1 和 1,2 和 2,2 按以下方式配对

 array([[11., 12.],
       [21., 22.]])

我将非常感谢您的纠正建议。提前致谢!

Answer 1

我很确定你可以用更简单的方法解决这个问题，但我不能 100% 确定我正确理解了你的意思，但这里有一些代码可以实现我认为你想要的功能。如果您有机会检查结果是否有效，我建议您这样做。

import numpy as np

n = 20
q = np.zeros((20, 3))
q[:, -1] = np.linspace(0, 10, n)

check_matrix = np.array([[1, 1, 0, 0, 0, 0,    0,       -0.7977, -0.243293],
                         [1, 1, 0, 0, 0, 0,    0,       1.5954,  0.004567],
                         [1, 2, 0, 0, 0, -1,   0,       0,       1.126557],
                         [2, 1, 0, 0, 0, 0.5,  0.86603, 1.5954,  0.038934],
                         [2, 1, 0, 0, 0, 2,    0,       -0.7977, -0.015192],
                         [2, 2, 0, 0, 0, -0.5, 0.86603, 1.5954,  0.21394]])
check_matrix[:, :2] -= 1  # python indexing is zero based

matrices = np.zeros((n, 2, 2), dtype=np.complex_)

for i in range(2):
    for j in range(2):
        k_list = []
        for k in range(len(check_matrix)):
            if check_matrix[k][0] == i and check_matrix[k][1] == j:
                k_list.append(check_matrix[k][8] *
                              np.exp(-1J * np.dot(q, check_matrix[k][2:5] 
                                                     - check_matrix[k][5:8])))

        matrices[:, i, j] = np.sum(k_list, axis=0)

注意:我更改了您的索引以保持一致 从零开始的索引。

---

这是我用向量版本替换 k 循环的另一种方法:

for i in range(2):
    for j in range(2):
        k = np.logical_and(check_matrix[:, 0] == i, check_matrix[:, 1] == j)
        temp = np.dot(check_matrix[k, 2:5] - check_matrix[k, 5:8], q[:, :, np.newaxis])[..., 0]
        temp = check_matrix[k, 8:] * np.exp(-1J * temp)
        matrices[:, i, j] = np.sum(temp, axis=0)