[Numpy-discussion] help with applying a transform

[Gary Ruben]

I'm trying to work out how to apply a 2x2 transform matrix to a spinor, e.g.

[psi1']   [a  b][psi1]
[     ] = [    ][    ]
[psi2']   [c  d][psi2]

where [[a,b],[c,d]] is a transform matrix and psi1 and psi2 are
i x j x k complex arrays representing complex scalar field data. I 
worked that one way to do it with 2D fields (psi1 and psi2 being 2D,
i x j arrays) is

def transform(tx_matrix, psi1, psi2):
     psi = np.dot(tx_matrix, np.rollaxis(np.dstack((psi1,psi2)),2,1))
     return psi[0], psi[1]

or, equivalently

def transform(tx_matrix, psi1, psi2):
     psi = np.dot(tx_matrix, np.rollaxis(
             np.concatenate((psi1[newaxis],psi2[newaxis])),1))
     return psi[0], psi[1]

but, as seems usual for me, I'm confused with trying to extend this to 
the next higher dimension; 3D in this case. It seems to me that there 
might be a neater way to do this, i.e. some built-in feature of an 
existing numpy function. How do I extend this for 3D psi1 and psi2 
arrays? Are there any general tips or guides for helping to minimise 
confusion when it comes to manipulating axes like this, such as standard 
ways to extend 2D recipes to 3D?

Gary R.