Эффективное вращение векторного поля в Python ⇐ Python

[Anonymous]

Мне нужно повернуть трехмерное векторное поле на 90 градусов. Я делаю это, сначала используя функцию NumPy rot90 (см. документацию здесь). Затем я поворачиваю каждый компонент в каждой позиции очень неэффективным способом: используя вложенные циклы for. Это работает нормально, как видно, например, из следующих двух рисунков: слева вы видите исходное векторное поле, где все векторы просто указывают в направлении x, справа все поле повернуто на 90 градусов. Он работает так, как задумано, но реализован очень неэффективно, что делает его практически бесполезным, когда я подхожу к большому массиву (например, 100x200x400).


Это моя реализация:

Код: Выделить всё

# import standard modules
import matplotlib.pyplot as plt
import numpy as np
from mayavi import mlab

def define_3D_vector_field( Nx=30, Ny=40, Nz=50,
direction='x'
):
# initialze empty arrays, one for each vector component
vecfield_x  = np.zeros( (Nx, Ny, Nz) )
vecfield_y  = np.zeros( (Nx, Ny, Nz) )
vecfield_z  = np.zeros( (Nx, Ny, Nz) )

if direction == 'x':
vecfield_x[ round(.25*Nx):round(.75*Nx),
round(.25*Ny):round(.75*Ny),
round(.25*Nz):round(.75*Nz) ] = 1.

return vecfield_x, vecfield_y, vecfield_z

def make_simple_3Dvecfield_plot( vecfield_x, vecfield_y, vecfield_z ):
fig1    = mlab.figure( bgcolor=(1,1,1), fgcolor=(0,0,0),
size=(800,600),
)

src = mlab.pipeline.vector_field( vecfield_x, vecfield_y, vecfield_z )
mlab.pipeline.vectors(src, mask_points=20, scale_factor=4.)

mlab.xlabel('x')
mlab.ylabel('y')
mlab.zlabel('z')
mlab.outline()

mlab.show()

def rot90vecfield(vecfield_x, vecfield_y, vecfield_z, rot_axis='y'):
theta   = np.radians(90.)

if rot_axis == 'y':
vf_x = np.rot90(vecfield_x, k=1, axes=(2,0))
vf_y = np.rot90(vecfield_y, k=1, axes=(2,0))
vf_z = np.rot90(vecfield_z, k=1, axes=(2,0))

# NOTE: this requires the transposed matrix to work correctly
rot_mat = np.array( [ [np.cos(theta) , 0., np.sin(theta)  ],
[0.            , 1., .0             ],
[-np.sin(theta), 0.,  np.cos(theta) ] ] ).T

vf_x_tmp = np.copy(vf_x)
vf_y_tmp = np.copy(vf_y)
vf_z_tmp = np.copy(vf_z)

for xx in range(vecfield_x.shape[0]):
for yy in range(vecfield_x.shape[1]):
for zz in range(vecfield_x.shape[2]):
vec_xyz = np.dot( rot_mat, np.array( [vf_x_tmp[zz,yy,xx],
vf_y_tmp[zz,yy,xx],
vf_z_tmp[zz,yy,xx]] ) )
vf_x[ zz, yy, xx ] = vec_xyz[0]
vf_y[ zz, yy, xx ] = vec_xyz[1]
vf_z[ zz, yy, xx ] = vec_xyz[2]

return vf_x, vf_y, vf_z

def main():
vf_x, vf_y, vf_z = define_3D_vector_field()
make_simple_3Dvecfield_plot( vf_x, vf_y, vf_z )

vf_x, vf_y, vf_z = rot90vecfield( vf_x, vf_y, vf_z )
make_simple_3Dvecfield_plot( vf_x, vf_y, vf_z )

if __name__ == '__main__':
main()

Я почти уверен, что есть какая-то хорошая маленькая алгебра, которую я могу применить, но я просто не могу уложиться в этом. Будем очень признательны за любые подсказки о том, как сделать это более эффективным (с точки зрения памяти и вычислений)!

Подробнее здесь: https://stackoverflow.com/questions/786 ... -in-python