如何对Numpy生成的n×n×n三维数组进行任意角度三轴旋转？

Got it, let's tackle this problem step by step. Rotating a 3D NumPy array by arbitrary angles isn't as straightforward as using rot90, but it's totally doable with rotation matrices, coordinate mapping, and interpolation. Here's how you can pull it off:

1. Start with 3D Rotation Matrices #

First, you need standard rotation matrices for each axis—these define how every point in your 3D array gets shifted when rotated. Let θ be your rotation angle (in radians):

• 绕X轴旋转：Keeps x-values fixed, rotates around the x-axis (right-hand rule applies here)

  def rotation_matrix_x(theta):
      return np.array([
          [1, 0, 0],
          [0, np.cos(theta), -np.sin(theta)],
          [0, np.sin(theta), np.cos(theta)]
      ])
  

• 绕Y轴旋转：Keeps y-values fixed, rotates around the y-axis

  def rotation_matrix_y(theta):
      return np.array([
          [np.cos(theta), 0, np.sin(theta)],
          [0, 1, 0],
          [-np.sin(theta), 0, np.cos(theta)]
      ])
  

• 绕Z轴旋转：Keeps z-values fixed, rotates around the z-axis

  def rotation_matrix_z(theta):
      return np.array([
          [np.cos(theta), -np.sin(theta), 0],
          [np.sin(theta), np.cos(theta), 0],
          [0, 0, 1]
      ])
  

2. Core Rotation Function #

The full workflow involves centering coordinates, applying rotation, and interpolating to handle non-integer positions (since rotating arbitrary angles won't land points exactly on grid indices). Here's a reusable function:

import numpy as np
from scipy.interpolate import interpn

def rotate_3d_array(arr, theta, axis='x'):
    # Ensure input is a square 3D array
    n = arr.shape[0]
    assert arr.shape == (n, n, n), "Input must be an n×n×n 3D array"
    
    # Select the correct rotation matrix
    if axis == 'x':
        rot_mat = rotation_matrix_x(theta)
    elif axis == 'y':
        rot_mat = rotation_matrix_y(theta)
    elif axis == 'z':
        rot_mat = rotation_matrix_z(theta)
    else:
        raise ValueError("Axis must be 'x', 'y', or 'z'")
    
    # Generate and center the original coordinate grid
    coords = np.mgrid[0:n, 0:n, 0:n].reshape(3, -1)  # Flatten to (3, n³)
    coords -= (n - 1) / 2  # Shift origin to the center of the array
    
    # Apply rotation and shift back to original coordinate space
    rotated_coords = rot_mat @ coords
    rotated_coords += (n - 1) / 2
    
    # Interpolate values from original array to rotated positions
    original_grid = (np.arange(n), np.arange(n), np.arange(n))
    # Use 'nearest' for sharp pixel results, 'linear' for smoother gradients
    rotated_arr = interpn(original_grid, arr, rotated_coords.T, method='linear', bounds_error=False, fill_value=0)
    
    # Reshape back to n×n×n and return
    return rotated_arr.reshape(n, n, n)

3. Test It Out #

Let's use a simple example to verify:

# Create a 5x5x5 cube with a marker at the corner (0,0,0)
cube = np.zeros((5,5,5))
cube[0,0,0] = 1

# Rotate 45 degrees (π/4 radians) around the Y-axis
rotated_cube = rotate_3d_array(cube, np.pi/4, axis='y')

# Check where the marker moved (should be near the edge)
print(np.where(rotated_cube > 0.5))

4. Key Notes #

• Interpolation Choice: Use method='nearest' for binary arrays (like segmentation masks) to keep sharp edges, or 'linear' for continuous data (like volumetric images) for smoother rotations.
• Bounds Handling: The fill_value=0 parameter fills any points that rotate outside the original array with 0—you can adjust this to match your data (e.g., the mean value of the array).
• Rotation Direction: All matrices follow the right-hand rule—point your thumb along the axis, and your fingers curl in the direction of rotation. Reverse the angle (use -theta) if you need the opposite direction.

内容的提问来源于stack exchange，提问作者Minebomber