How to properly rotate 3D numpy array to uniformly sample sphere?

Question

I have a 3D numpy array with density values. I would like to rotate this array around the origin to uniformly sample the sphere as closely as possible. I found some answers here that allow me to get close (Evenly distributing n points on a sphere, Apply rotation defined by Euler angles to 3D image, in python).

However, when rotating the array it does not seem to correctly sample the sphere, and leaves about half of it out and lacks sampling at the poles. I can't seem to figure out what my issue is. I'm guessing it has something to do with conversion of spherical to Euler to rotation matrix. Any help would be appreciated.

import sys
import numpy as np
from numpy.testing import assert_allclose
from scipy.spatial.transform import Rotation
from scipy import ndimage
import matplotlib.pyplot as plt

def generate_phi_theta_uniform_sphere_by_spiral_method(num_pts):
    #taken from https://stackoverflow.com/a/44164075/2836338
    indices = np.arange(0, num_pts, dtype=float) + 0.5
    phi = np.arccos(1 - 2 * indices / num_pts)
    theta = np.pi * (1 + 5 ** 0.5) * indices
    return phi, theta

def rotate_density(rho, alpha, beta, gamma, order=1):
    #https://nbviewer.jupyter.org/gist/lhk/f05ee20b5a826e4c8b9bb3e528348688
    #https://stackoverflow.com/questions/59738230

    # create meshgrid
    dim = rho.shape

    ax = np.arange(dim[0])
    ay = np.arange(dim[1])
    az = np.arange(dim[2])
    coords = np.meshgrid(ax, ay, az)

    # stack the meshgrid to position vectors, center them around 0 by substracting dim/2
    xyz = np.vstack([coords[0].reshape(-1) - float(dim[0]) / 2,  # x coordinate, centered
                     coords[1].reshape(-1) - float(dim[1]) / 2,  # y coordinate, centered
                     coords[2].reshape(-1) - float(dim[2]) / 2])  # z coordinate, centered

    # create transformation matrix
    r = Rotation.from_euler('zyx', [alpha, beta, gamma], degrees=False)
    mat = r.as_matrix()

    # apply transformation
    transformed_xyz = np.dot(mat, xyz)

    # extract coordinates
    x = transformed_xyz[0, :] + float(dim[0]) / 2
    y = transformed_xyz[1, :] + float(dim[1]) / 2
    z = transformed_xyz[2, :] + float(dim[2]) / 2

    x = x.reshape((dim[1],dim[0],dim[2]))
    y = y.reshape((dim[1],dim[0],dim[2]))
    z = z.reshape((dim[1],dim[0],dim[2])) # reason for strange ordering: see next line

    # the coordinate system seems to be strange, it has to be ordered like this
    new_xyz = [y, x, z]

    # sample
    new_rho = ndimage.map_coordinates(rho, new_xyz, order=order)

    return new_rho

def spherical_to_euler(phi, theta):
    """Converts spherical coordinates (phi, theta) to Euler angles (zyx convention).

    Args:
        phi: Azimuth angle in radians (0 to 2*pi).
        theta: Coltitude angle in radians (0 to pi or extended range).

    Returns:
        alpha, beta, gamma: Euler angles (zyx convention) in radians.
    """

    phi = np.atleast_1d(phi)
    theta = np.atleast_1d(theta)

    # Check for theta close to pi (tolerance can be adjusted)
    tolerance = 1e-6
    theta_near_pi = np.abs(theta - np.pi) < tolerance

    # Handle negative theta values (wrap to positive range)
    theta = np.where(~theta_near_pi, np.mod(theta + np.pi, 2 * np.pi) - np.pi, theta)

    # Alpha (rotation around Z)
    alpha = phi

    # Beta (rotation around Y) with handling for theta at pi
    beta = np.pi / 2 - theta
    beta = np.where(beta < -np.pi, beta + 2 * np.pi, beta)  # Adjust for other negative beta

    # Set beta to pi/2 for theta close to pi (South Pole)
    beta = np.where(theta_near_pi, np.pi / 2, beta)

    # Gamma (rotation around X) - set to zero for ZYX convention
    gamma = np.zeros(phi.shape)

    return alpha, beta, gamma

def test_spherical_to_euler():
    # Test cases with various phi and theta values
    # each tuple formatted as (phi, theta, alpha_expected, beta_expected, gamma_expected)
    test_data = [
        # Standard case (ZYX convention)
        (np.pi/3, np.pi/4, np.pi/3, np.pi/2 - np.pi/4, 0),
        # Edge cases (theta=0, pi)
        (np.pi/3, 0, np.pi/3, np.pi/2, 0),
        (np.pi/3, np.pi, np.pi/3, np.pi/2, 0),
        # Edge cases (phi=0, 2*pi)
        (0, np.pi/4, 0, np.pi/2 - np.pi/4, 0),
        (2*np.pi, np.pi/4, 2*np.pi, np.pi/2 - np.pi/4, 0),
        # Additional test cases
        (np.pi/6, np.pi/2, np.pi/6, 0, 0),
        (5*np.pi/6, np.pi/3, 5*np.pi/6, np.pi/2 - np.pi/3, 0),
    ]

    for phi, theta, alpha_expected, beta_expected, gamma_expected in test_data:
        alpha, beta, gamma = spherical_to_euler(phi, theta)

        # Check results
        np.testing.assert_allclose(alpha, alpha_expected)
        np.testing.assert_allclose(beta, beta_expected)
        np.testing.assert_allclose(gamma, gamma_expected)

    print("Unit tests passed!")

test_spherical_to_euler()


num_pts = 1000 #number of samples on unit sphere
phi, theta = generate_phi_theta_uniform_sphere_by_spiral_method(num_pts)

#convert to cartesian coordinates for plotting for the spiral method
x_spiral, y_spiral, z_spiral = np.cos(theta) * np.sin(phi), np.sin(theta) * np.sin(phi), np.cos(phi)

#make a density map with only one nonzero pixel
n = 32
rho = np.zeros((n,n,n))
rho[n//4, n//4, n//4] = 1

x_ = np.linspace(-n/2.,n/2.,n)
x,y,z = np.meshgrid(x_,x_,x_,indexing='ij')

#convert spherical to euler
alpha, beta, gamma = spherical_to_euler(phi=phi, theta=theta)

#transform the array for each orientation
rho_rotated = np.copy(rho)
for i in range(num_pts):
    #since its only one voxel, to make plotting easy just add up all the arrays
    #to show all the voxels on one scatter plot easily
    rho_rotated += rotate_density(rho, alpha[i], beta[i], gamma[i], order=0)


fig = plt.figure(figsize=plt.figaspect(0.5))

#plot the points using the spiral method
ax0 = fig.add_subplot(1, 2, 1, projection='3d')
ax0.set_box_aspect((np.ptp(x_spiral), np.ptp(y_spiral), np.ptp(z_spiral)))
ax0.scatter(x_spiral, y_spiral, z_spiral)
ax0.title.set_text('Spiral Method')
ax0.view_init(elev=45, azim=30)

#plot the voxel points for the array rotation method
ax1 = fig.add_subplot(1, 2, 2, projection='3d')
ax1.set_box_aspect((np.ptp(x), np.ptp(y), np.ptp(z)))
ax1.scatter(x[rho_rotated>0], y[rho_rotated>0], z[rho_rotated>0])
ax1.set_xlim([x.min(),x.max()])
ax1.set_ylim([y.min(),y.max()])
ax1.set_zlim([z.min(),z.max()])
ax1.title.set_text('3D Array Rotation Method')
ax1.view_init(elev=45, azim=30)

plt.savefig("points_sampled_on_sphere_by_spiral_method_or_array_rotation.png",dpi=150)
plt.show()