用 numpy 旋转网格 [英] Rotate meshgrid with numpy

问题描述

我想生成一个坐标已旋转的网格.我必须在双循环中进行旋转，并且我确信有更好的方法来对其进行矢量化.代码如下:

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)

# 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)

在两个网格上写一个双循环似乎很愚蠢.那里的一位向导是否知道如何将其矢量化?

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?

推荐答案

Einstein summation (np.einsum) 处理这类事情非常快.对于 1001x1001，我得到了 97 毫秒.

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]))

这篇关于用 numpy 旋转网格的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！