Uniform randomly transform set of points around the origin in 3D-CodePudding

I have some point cloud data in a numpy array of shape (3, n) where each column is a single point in form [x, y, z]. I would like to apply a function to this data that randomly transforms it around [0, 0, 0] such that the cloud maintains its shape and every point maintains its distance from the origin, and every orientation that meets those criteria is equally likely in the output (so its uniform).

My current solution is pretty naive and it does not meet the uniform criteria. I found a function I call here rotate_3D which transforms a set of points around one axis given a theta value, and I just applied it 3 times with 3 different random values to rotate around the 3 axes. I can't imagine this is very efficient, and it does not result in a uniform result:

def rotate_3D(set1, set2, theta):
    cos_theta = math.cos(theta)
    sin_theta = math.sin(theta)
    set1_rotated = set1 * cos_theta - set2 * sin_theta
    set2_rotated = set2 * cos_theta   set1 * sin_theta
    return set1_rotated, set2_rotated

def rotate_full_set(matrix):
    theta_x = random.uniform(-math.pi, math.pi)
    theta_y = random.uniform(-math.pi, math.pi)
    theta_z = random.uniform(-math.pi, math.pi)
    x_positions = matrix[0]
    y_positions = matrix[1]
    z_positions = matrix[2]
    x_rotated, y_rotated = rotate_3D(x_positions, y_positions, theta_z)
    y_rotated, z_rotated = rotate_3D(y_rotated, z_positions, theta_x)
    z_rotated, x_rotated = rotate_3D(z_rotated, x_rotated, theta_y)

    return np.stack([x_rotated, y_rotated, z_rotated])

Do you know of a better method, or maybe just a way I can generate my theta values that results in a uniform distribution of results?

CodePudding user response：

You can generate random rotations with scipy.stats.special_ortho_group, which, as explained in the docstring, will "return a random rotation matrix, drawn from the Haar distribution (the only uniform distribution on SO(N)) with a determinant of 1." If reflections are allowed (so the possible determinants are 1 and -1), you can use scipy.stats.ortho_group.

For example,

In [12]: from scipy.stats import special_ortho_group

In [13]: R = special_ortho_group.rvs(3)  # Generate a random rotation from SO(3).

In [14]: R
Out[14]: 
array([[ 0.79450465,  0.05168791, -0.60505431],
       [ 0.05454748,  0.98626859,  0.15588086],
       [ 0.60480322, -0.15685226,  0.78077553]])

Multiply your data by R to compute the rotated set:

In [35]: rng = np.random.default_rng()

In [36]: pts = rng.integers(0, 3, size=(3, 5))

In [37]: pts
Out[37]: 
array([[2, 1, 2, 1, 0],
       [1, 1, 1, 1, 0],
       [1, 1, 0, 1, 2]])

In [38]: R @ pts
Out[38]: 
array([[ 1.03564289,  0.24113824,  1.6406972 ,  0.24113824, -1.21010863],
       [ 1.25124442,  1.19669693,  1.09536356,  1.19669693,  0.31176172],
       [ 1.83352971,  1.22872649,  1.05275418,  1.22872649,  1.56155107]])