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.

What are the dimensions of the mesh and the rotation matrix? By 'unraveling' do you mean you tried the 2nd answer at your link? With a little more information I could help with the `einsum`, but I doubt if it is more memory efficient. Stll we can try. — hpaulj, Aug 04 '15 at 19:04
@hpaulj, more specifically, I was trying to recreate the meshgrid from the points. Since this set of points was rotated by 15 degrees I bet it made meshgrid go crazy. I am now looking at whether I can use the points themselves to interpolate my 3d data onto. — Chris, Aug 04 '15 at 19:49
I found that I can use the points I rotated using the loop talked about in the original post to interpolate onto, without having to put them into a meshgrid. But I still am quite interested in your ideas on how to do the 3d rotation with einsum(). Thanks! — Chris, Aug 04 '15 at 21:19

hpaulj · Accepted Answer · 2015-08-04T21:43:28.803

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.

I especially appreciate the slight changes you made to the code such as returning the matrix rather than the rotated coords, saves many calculations. Also the list comprehension... I keep forgetting to include these in my code. And thanks for showing the extension to 4d, I may need to used homogeneous coordinates for my transforms. — Chris, Aug 05 '15 at 13:09

Numpy einsum() for rotation of meshgrid

1 Answers1

Linked