用numpy旋转网格

Question

我想要生成一个坐标已经旋转的网格。我必须做一个双循环的旋转，我相信有一个更好的方法来矢量化它。代码是这样的：

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

​

​

要在两个网格上写一个双循环似乎很愚蠢。有一个巫师知道如何将它向量化吗？

Answer 1

爱因斯坦求和(np.einsum)对于这类事情来说是非常快的。我的1001x1001有97毫秒。

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

Answer 2

也许我误解了这个问题，但我通常只是.

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)

~100

Answer 3

你可以用一些reshaping & flattening with np.ravel去掉这两个嵌套循环，然后用np.dot保持矩阵乘法，就像这样-

mult = np.dot( RotMatrix, np.array([x.ravel(),y.ravel()]) )
xrot = mult[0,:].reshape(xrot.shape)
yrot = mult[1,:].reshape(yrot.shape)