Dave31415 - 1 year ago 91

Python Question

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.

Answer Source

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

Those colons are completely unnecessary.