如何在 numpy 中有效地将 3D 矩阵与二维矩阵相乘

Question

I have two multidimensional arrays, which I want to multiply with each other. One has the shape N,N,3 and the other has the shape N,N.

Let me set the stage:

I have an array of atom positions of the shape N,3 :

atom_positions = [[x1,y1,z1],
                  [x2,y2,z2],
                  [x3,y3,z3],
                  ...       
                            ]

From these I calculate an upper triangular matrix of distance vectors so that the resulting N,N,3 matrix contains all unique pair distance vectors r_ij of the vectors inside atom_positions :

pair_distance_vectors = [[[0,0,0],[x2-x1,y2-y1,z2-z1],[x3-x1,y3-y1,z3-z1],...],
                         [[0,0,0],[0,0,0]            ,[x3-x2,y3-y2,z3-z2],...],
                          ...
                                                                              ]

Now I want to normalize each of these pair distance vectors. For that I want to use my N,N pair_distances array, which contains the length of every vector inside pair_distance_vectors . The formula for a single vector is: r_ij/|r_ij|

I want to do that by doing a matrix multiplication, where every entry in the N,N array becomes a scalar by which a vector inside the N,N,3 array is multiplied. I'm pretty sure that this can be achieved somehow with numpy by using numpy.dot() or a different function, but I just can't find the answer myself. Also, I'm afraid if I do find a transformation which allows for this, that my maths will be faulty.

Here's some demonstration code, which achieves what I want in a very inefficient fashion:

import numpy as np

pair_distance_vectors = np.ones(shape=(2,2,3))
pair_distances = np.array(((1,2),(3,4)))
normalized_pair_distance_vectors  = np.zeros(shape=(2,2,3))

for i,vec_list in enumerate(pair_distance_vectors):
    for j,vec in enumerate(vec_list):
        normalized_pair_distance_vectors[i,j] = vec*pair_distances[i,j]

print(normalized_pair_distance_vectors)

Thanks in advance.

EDIT: Maybe this is clearer:

distance_vectors = [[[x11,y11,z11],[x12,y12,z12],[x13,y13,z13],...],
                    [[x21,y21,z21],[x22,y22,z22],[x23,y23,z23],...],
                    ...                                            ]

distance_matrix = [[r_11,r_12,r_13,...],
                   [r_21,r_22,r_23,...],
                    ...                 ]

norm_distance_vectors = some_operation(distance_vectors,distance_matrix)

norm_distance_vectors = [[r_11*[x11,y11,z11],r_12*[x12,y12,z12],r_13*[x13,y13,z13],...],
                         [r_21*[x21,y21,z21],r_22*[x22,y22,z22],r_23*[x23,y23,z23],...],
                          ...                                                    ]

Answer 1

You won't need a loop. Trick is to expand your pair_distance in the 3rd dimension by repeating it m times ( m being the dimension of your vectors, here 3D) and then divide two arrays element wise (works for any m -dimensional vectors, replace 3 with m ):

pair_distances = np.repeat(pair_distances[:,:,None], 3, axis=2)
normalized_pair_distance_vectors = np.nan_to_num(pair_distance_vectors/ pair_distances)

Output for your example inputs:

[[[1.         1.         1.        ]
  [0.5        0.5        0.5       ]]

 [[0.33333333 0.33333333 0.33333333]
  [0.25       0.25       0.25      ]]]