python 网格计算_python – 用numpy旋转网格

我想要生成一个坐标已旋转的网格.我必须在双循环中进行旋转,我确信有更好的方法来矢量化它.代码如下：

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

在两个网格物体上写一个双循环似乎很愚蠢.那里的一个巫师有一个想法如何矢量化这个？

最佳答案 爱因斯坦求和(

np.einsum)对于这种事情非常快. 1001×1001我得到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]))