Nump einsum () to rotate meshgrid
I have a set of 3d coordinates that was generated using meshgrid (). I want to be able to rotate them about 3 axes.
I tried to untangle the meshgrid and rotate at every point, but the meshgrid is large and I ran out of memory.
This question resolves this in 2d using einsum (), but I cannot figure out the format of the string when expanding it to 3d.
I've read several other pages about einsum () and format strings but couldn't figure it out.
EDIT:
I name my meshgrid axes X, Y and Z, each with the shape (213, 48, 37). Also, there was an actual memory error when I tried to return the results to the meshgrid.
When I tried to figure it out for pivoting by points, I used the following function:
def mg2coords(X, Y, Z):
return np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T
I have looped the result to 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 rotating, I will use the points to interpolate on.
source to share
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
set the 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 line by line:
In [844]: xyz1=np.array([np.dot(rot,i) for i in xyz])
equivalent to einsum
counting lines by line:
In [845]: xyz2=np.einsum('ij,kj->ki',rot,xyz)
They correspond to:
In [846]: np.allclose(xyz2,xyz1)
Out[846]: True
Alternatively, I could collect 3 arrays into one 4d array and rotate it with einsum
. Here np.array
adds dimension at the beginning. So dot
sum size j
is 1 and 3d arrays:
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
Similarly for 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 ij
if you want to turn in a different direction, i.e. use rot.T
.
source to share