I want to use a high-dimensional numpy array to store the norms of weighted sums of matrices. For example:

 mat1, mat2, mat3, mat4 = np.random.rand(3, 3), np.random.rand(3, 3), np.random.rand(3, 3), np.random.rand(3, 3)

 res = np.empty((8, 7, 6, 5))

 for i in range(8):
    for j in range(7):
        for p in range(6):
            for q in range(5):
                 res[i, j, p, q] = np.linalg.norm(i * mat1   j * mat2   p * mat3   q * mat4)

I would like to ask that are there any methods to avoid this nested loop?

CodePudding user response：

Solution

Here's one way you can do it, via adding axes with None (equivalent to np.newaxis):

def weighted_norms(mat1, mat2, mat3, mat4):
    P = mat1 * np.arange(8)[:, None, None]
    Q = mat2 * np.arange(7)[:, None, None]
    R = mat3 * np.arange(6)[:, None, None]
    S = mat4 * np.arange(5)[:, None, None]
    
    summation = S   R[:, None]   Q[:, None, None]   P[:, None, None, None]
    return np.linalg.norm(summation, axis=(4, 5))

Veracity and a simple benchmark

In [6]: output = weighted_norms(mat1, mat2, mat3, mat4)

In [7]: np.allclose(output, res)
Out[7]: True

In [8]: %timeit weighted_norms(mat1, mat2, mat3, mat4)
71.3 µs ± 446 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Explanation

By adding two new axes to the np.arange objects, you can force the broadcasting you want, producing 0 * mat1, 1 * mat1, 2 * mat1 ....

The real tricky bit is then constructing the (8, 7, 6, 5, 3, 3) array (which is the shape before evaluating the norm which collapses the last two dimensions).

Notice that the summation of all the weighted 3D arrays starts with the last array, S, and progressively adds more weighted 3D arrays. The way it does this is by adding a new axis to broadcast over at each step.

For example, the shape of S is (5, 3, 3) and in order to correctly add R you need to insert a new axis. So the shape of R goes from (6, 3, 3) to (6, 1, 3, 3). This second dimension of size 1 is what allows us to broadcast the sum of S over R such that each array in the 3D S is added to each array in R (that's one level of nested loop).

Then we need to add Q (for every array in Q, for every array in R, for every array in S), so we need to insert two new axes turning Q from (7, 3, 3) to (7, 1, 1, 3, 3).

Finally, P goes from (8, 3, 3) to (8, 1, 1, 1, 3, 3).

It may help to "visualize" this by overlaying the shapes:

            (5, 3, 3)  <-  S
             :
         (6, 1, 3, 3)  <-  R[:, None]
---------------------
         (6, 5, 3, 3)
          :  :
      (7, 1, 1, 3, 3)  <-  Q[:, None, None]
---------------------
      (7, 6, 5, 3, 3)
       :  :  :
   (8, 1, 1, 1, 3, 3)  <-  P[:, None, None, None]
---------------------
   (8, 7, 6, 5, 3, 3)

Generalizing

Here's a generalized version using a helper function for adding axes just to clean up the code a little:

from typing import Tuple
import numpy as np


def add_axes(x: np.ndarray, n: int) -> np.ndarray:
    """
    Inserts `n` number of new axes into `x` from axis 1 onward.
    e.g., for `x.shape == (3, 3)`, `add_axes(x, 2) -> (3, 1, 1, 3)`
    """
    return np.expand_dims(x, axis=(*range(1, n   1),))


def weighted_norms(arrs: Tuple[np.ndarray], weights: Tuple[int]) -> np.ndarray:
    if len(arrs) != len(weights):
        raise ValueError("Number of arrays must match number of weights")
    summation = np.empty((weights[-1], *arrs[-1].shape))
    for i, (x, w) in enumerate(zip(arrs[::-1], weights[::-1])):
        summation = summation   add_axes(x * add_axes(np.arange(w), 2), i)
    return np.linalg.norm(summation, axis=(-1, -2))

Usage:

In [10]: arrs = (mat1, mat2, mat3, mat4)

In [11]: weights = (8, 7, 6, 5)

In [12]: output = weighted_norms(arrs, weights)

In [13]: np.allclose(output, res)
Out[13]: True

In [14]: %timeit weighted_norms(arrs, weights)
109 µs ± 3.07 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

CodePudding user response：

Here is a completely vectorized way to solve what you are looking for. Comments inline for explanation.

import itertools
import numpy as np

#Generate the i,j,p,q multipliers as a matrix
res_new = np.array(list(itertools.product(range(8), range(7), range(6), range(5))))

#(1680,4) @ (4,3,3) -> (1680,3,3)
prod = np.einsum('ij,jkl -> ikl', res_new, np.stack([mat1, mat2, mat3, mat4]))

#calculating norm over last 2 axis and reshaping 1680 -> 8*7*6*5
output = np.linalg.norm(prod, axis=(-1,-2)).reshape((8,7,6,5))
output.shape

(8,7,6,5)

Just to confirm, if the output that this code generates is the same as the output you are getting from your nested for loops -

np.allclose(output, res)

True