6

I am wanting to produce a meshgrid whose coordinates have been rotated. I have to do the rotation in a double loop and I'm sure there is a better way to vectorize it. The code goes as so:

# Define the range for x and y in the unrotated matrix
xspan = linspace(-2*pi, 2*pi, 101)
yspan = linspace(-2*pi, 2*pi, 101)

# Generate a meshgrid and rotate it by RotRad radians.
def DoRotation(xspan, yspan, RotRad=0):

    # Clockwise, 2D rotation matrix
    RotMatrix = np.array([  [np.cos(RotRad),  np.sin(RotRad)],
                            [-np.sin(RotRad), np.cos(RotRad)]])
    print RotMatrix

    # This makes two 2D arrays which are the x and y coordinates for each point.
    x, y = meshgrid(xspan,yspan)

    # After rotating, I'll have another two 2D arrays with the same shapes.
    xrot = zeros(x.shape)
    yrot = zeros(y.shape)

    # Dot the rotation matrix against each coordinate from the meshgrids.
    # I BELIEVE THERE IS A BETTER WAY THAN THIS DOUBLE LOOP!!!
    # I BELIEVE THERE IS A BETTER WAY THAN THIS DOUBLE LOOP!!!
    # I BELIEVE THERE IS A BETTER WAY THAN THIS DOUBLE LOOP!!!
    # I BELIEVE THERE IS A BETTER WAY THAN THIS DOUBLE LOOP!!!
    # I BELIEVE THERE IS A BETTER WAY THAN THIS DOUBLE LOOP!!!
    # I BELIEVE THERE IS A BETTER WAY THAN THIS DOUBLE LOOP!!!
    for i in range(len(xspan)):
        for j in range(len(yspan)):
            xrot[i,j], yrot[i,j] = dot(RotMatrix, array([x[i,j], y[i,j]]))

    # Now the matrix is rotated
    return xrot, yrot

# Pick some arbitrary function and plot it (no rotation)
x, y = DoRotation(xspan, yspan, 0)
z = sin(x)+cos(y)
imshow(z)

enter image description here

# And now with 0.3 radian rotation so you can see that it works.
x, y = DoRotation(xspan, yspan, 0.3)
z = sin(x)+cos(y)
figure()
imshow(z)

enter image description here

It seems silly to have to write a double loop over the two meshgrids. Do one of the wizards out there have an idea how to vectorize this?

Divakar
  • 218,885
  • 19
  • 262
  • 358
ZSG
  • 1,329
  • 14
  • 18

4 Answers4

5

Einstein summation (np.einsum) is very quick for this sort of thing. I got 97 ms for 1001x1001.

def DoRotation(xspan, yspan, RotRad=0):
    """Generate a meshgrid and rotate it by RotRad radians."""

    # Clockwise, 2D rotation matrix
    RotMatrix = np.array([[np.cos(RotRad),  np.sin(RotRad)],
                          [-np.sin(RotRad), np.cos(RotRad)]])

    x, y = np.meshgrid(xspan, yspan)
    return np.einsum('ji, mni -> jmn', RotMatrix, np.dstack([x, y]))
wilywampa
  • 1,280
  • 1
  • 15
  • 15
  • 1
    Terrific! Yes, I'm getting 97 ms too. I had no idea this function existed but it will become my new best friend. – ZSG Apr 17 '15 at 21:58
  • 1
    Just as a side-comment, this could not work with `ogrid()` because, in general, the rotated `grid` could not be expressed in terms of `ogrid()` – norok2 Aug 20 '18 at 09:53
4

Maybe I misunderstand the question, but I usually just...

import numpy as np

pi = np.pi

x = np.linspace(-2.*pi, 2.*pi, 1001)
y = x.copy()

X, Y = np.meshgrid(x, y)

Xr   =  np.cos(rot)*X + np.sin(rot)*Y  # "cloclwise"
Yr   = -np.sin(rot)*X + np.cos(rot)*Y

z = np.sin(Xr) + np.cos(Yr)

~100ms also

uhoh
  • 3,713
  • 6
  • 42
  • 95
3

You can get rid of those two nested loops with some reshaping & flattening with np.ravel and keeping that matrix multiplication with np.dot like so -

mult = np.dot( RotMatrix, np.array([x.ravel(),y.ravel()]) )
xrot = mult[0,:].reshape(xrot.shape)
yrot = mult[1,:].reshape(yrot.shape)
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • I used your suggestion but just switched the reshape within the dot to ravel, and increased the size of the matrix up to 1001x1001. The execution time is now 1/2 second instead of 10 seconds. Thank you! – ZSG Apr 17 '15 at 20:58
  • @ZSG Ah yes, ravel() must be a better option! Edited with it. – Divakar Apr 17 '15 at 20:59
  • Sorry, your answer is great, but the einsum really is faster, a single line, and easier to understand mathematically. – ZSG Apr 17 '15 at 21:58
  • @ZSG Ah that's okay and wow what a great speedup with that new function! – Divakar Apr 17 '15 at 22:03
0

Just in case you want to go for 3D, scipy.spatial.transform.Rotation may be useful

import numpy as np
from scipy.spatial.transform import Rotation as R

# define lines for x- and y-subdivision
x = np.linspace(-5, 5)
y = np.linspace(-5, 5)

# create meshgrid for a plane surface (just as example)
X, Y  = np.meshgrid(x, y)
Z = np.zeros(X.shape) # alternatively Z may come from a 3D-meshgrid

# define rotation by rotation angle and axis, here 45DEG around z-axis
r = R.from_rotvec(np.pi/4 * np.array([0, 0, 1]))

# arrange point coordinates in shape (N, 3) for vectorized processing
XYZ = np.array([X.ravel(), Y.ravel(), Z.ravel()]).transpose()

# apply rotation
XYZrot = r.apply(XYZ)

# return to original shape of meshgrid
Xrot = XYZrot[:, 0].reshape(X.shape)
Yrot = XYZrot[:, 1].reshape(X.shape)
Zrot = XYZrot[:, 2].reshape(X.shape)
Dominik Kern
  • 146
  • 8