How to fill a 4D matrix with a 2D matrix for every 2D cross section in Python

Question

So I have a matrix call it Vjunk which is 70x70x70x70. I have another matrix which is 70x70 call it V.

What I wanna do is that for every i, j the matrix Vjunk[:,:,i,j] is 70 by 70. I want to change this matrix so that it is replaced by itself + V[i,j] where V[i,j] is the ij-th element of my matrix V.

I tried

[Vjunk[:,:,i,j]=Vjunk[:,:,i,j]-beta*V[i,j] for i in range(humangrid_size) for j in range(assetgrid_size)]

but this command was unsucessful.

Answer 1

Let's use this indice on Vjunk : (m, n, i, j)

If I'm correct, you want that for every m, n combination, Vjunk(m,n,i,j) get replaced by Vjunk(m,n,i,j) -beta * V[i,j]. If that's the goal, this loop should do the trick :

for m in range(70):
    for n in range(70):
        for i in range(70):
            for j in range(70):
                Vjunk[m,n,i,j] = Vjunk[m,n,i,j] - beta * V[i,j]

Dunno if it's gonna be fast enough, even if it's only a 70*70*70*70 matrix. Still more than 20M operations.

The loop on i, j could probably get replaced by list comprehension.

Answer 2

First of all, you can't put assignments in a list comprehension.

Second of all: you're in luck, because Vjunk and V readily broadcast when you subtract them. Here is an example with non-trivial shapes to make it easier to spot bugs:

import numpy as np

Vjunk = np.random.rand(2, 3, 4, 5) 
V = np.random.rand(4, 5) 

# naive version: loop
res1 = Vjunk.copy() 
for i in range(2): 
    for j in range(3): 
        for l in range(4): 
            for m in range(5): 
                res1[i,j,l,m] -= V[l,m]

# vectorized, broadcasting version:
res2 = Vjunk - V

print(np.array_equal(res1, res2))
# True

Here Vjunk has shape (2, 3, 4, 5), and V has shape (4, 5). The latter is compatible with shape (1, 1, 4, 5) for the purposes for broadcasting, which is then compatible with the shape of Vjunk.

Performing the broadcasting subtraction Vjunk - V will exactly do what you want: for each element along the last two dimensions each value of Vjunk (a 2d array along its first two dimensions) will be decreased by V.

It's then trivial to throw in a scalar factor:

res = Vjunk - beta * V