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 (/a/44164075/2836338, 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 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 /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
    #/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):
    phi = np.atleast_1d(phi)
    theta = np.atleast_1d(theta)

    # Convert spherical coordinates to Euler angles (zyx)
    alpha = phi  # Alpha is phi (rotation around z)
    beta = np.pi / 2 - theta  # Beta from theta (rotation around y)
    gamma = np.zeros(phi.shape)  # Gamma is always 0 for uniform sampling on the unit sphere

    return alpha, beta, gamma

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