Dave31415 Dave31415 - 8 months ago 75
Python Question

Can you write numpy slicing generically?

I want to do something like

x[i, :, :] = (rhs[i, :, :]-diag[i] * x[i+1, :, :])/diag[i]

where x and rhs are 3D numpy arrays of size (T,L,S). diag is a 1D array of size T.

This will broadcast properly.

But now I'd like to write a similar function to work on 2D arrays or some other number of dimensions. How can I write this generically so that it will work on any array that has first dimension of size T. I don't want to duplicate code with just a different number of colons since there are a lot of these kinds of lines in the function.

x[i] = (rhs[i] - diag[i] * x[i+1])/diag[i]

Those colons are completely unnecessary.