Numpy einsum() for rotation of meshgrid

Question

I have a set of 3d coordinates that was generated using meshgrid(). I want to be able to rotate these about the 3 axes.

I tried unraveling the meshgrid and doing a rotation on each point but the meshgrid is large and I run out of memory.

This question addresses this in 2d with einsum(), but I can't figure out the string format when extending it to 3d.

I have read several other pages about einsum() and its format string but haven't been able to figure it out.

EDIT:

I call my meshgrid axes X, Y, and Z, each is of shape (213, 48, 37). Also, the actual memory error came when I tried to put the results back into a meshgrid.

When I attempted to 'unravel' it to do point by point rotation I used the following function:

def mg2coords(X, Y, Z):
    return np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T

I looped over the result with the following:

def rotz(angle, point):
    rad = np.radians(angle)
    sin = np.sin(rad)
    cos = np.cos(rad)
    rot = [[cos, -sin, 0],
           [sin,  cos, 0],
           [0, 0, 1]]

    return np.dot(rot, point)

After the rotation I will be using the points to interpolate onto.

Answer 1

Working with your definitions:

In [840]: def mg2coords(X, Y, Z):
        return np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T

In [841]: def rotz(angle):
        rad = np.radians(angle)
        sin = np.sin(rad)
        cos = np.cos(rad)
        rot = [[cos, -sin, 0],
               [sin,  cos, 0],
               [0, 0, 1]]
        return np.array(rot)
        # just to the rotation matrix

define a sample grid:

In [842]: X,Y,Z=np.meshgrid([0,1,2],[0,1,2,3],[0,1,2],indexing='ij')    
In [843]: xyz=mg2coords(X,Y,Z)

rotate it row by row:

In [844]: xyz1=np.array([np.dot(rot,i) for i in xyz])

equivalent einsum row by row calculation:

In [845]: xyz2=np.einsum('ij,kj->ki',rot,xyz)

They match:

In [846]: np.allclose(xyz2,xyz1)
Out[846]: True

Alternatively I could collect the 3 arrays into one 4d array, and rotate that with einsum. Here np.array adds a dimension at the start. So the dot sum j dimension is 1st, and the 3d of the arrays follow:

In [871]: XYZ=np.array((X,Y,Z))
In [872]: XYZ2=np.einsum('ij,jabc->iabc',rot,XYZ)

In [873]: np.allclose(xyz2[:,0], XYZ2[0,...].ravel())
Out[873]: True

Similary for the 1 and 2.

Alternatively I could split XYZ2 into 3 component arrays:

In [882]: X2,Y2,Z2 = XYZ2
In [883]: np.allclose(X2,xyz2[:,0].reshape(X.shape))
Out[883]: True

Use ji instead of ij if you want to rotate in the other direction, i.e. use rot.T.