Skip to content

GitLab

Commit b5881ee2 authored by maming's avatar maming
Browse files

Initial commit

parents
Show whitespace changes
Inline Side-by-side
Showing with 2745 additions and 0 deletions
lie_learn/lie_learn/groups/SE2.py 0 → 100755
View file @ b5881ee2
import numpy as np
import spaces.Rn as Rn
parameterizations = ('rotation-translation', '2x3 matrix', '3x3 matrix')
def compose(g, h, parameterization=None, g_parameterization=None, h_parameterization=None, out_parameterization=None):
"""
Compose elements g, h in SE(2).
"""
if parameterization is not None:
g_parameterization = parameterization
h_parameterization = parameterization
out_parameterization = parameterization
g_mat = change_parameterization(g, p_from=g_parameterization, p_to='3x3 matrix')
h_mat = change_parameterization(h, p_from=h_parameterization, p_to='3x3 matrix')
gh_mat = np.einsum('...ij,...jk->...ik', g_mat, h_mat)
return change_parameterization(g=gh_mat, p_from='3x3 matrix', p_to=out_parameterization)
def invert(g, parameterization):
"""
Invert element g in SE(2), where g can have any supported parameterization.
"""
# Change to (theta, tau1, tau2) paramterization.
g_rt = change_parameterization(g, p_from=parameterization, p_to='rotation-translation')
g_inv_rt = np.empty_like(g_rt)
g_inv_rt[..., 0] = -g_rt[..., 0]
g_inv_rt[..., 1] = -(np.cos(-g_rt[..., 0]) * g_rt[..., 1] - np.sin(-g_rt[..., 0]) * g_rt[..., 2])
g_inv_rt[..., 2] = -(np.sin(-g_rt[..., 0]) * g_rt[..., 1] + np.cos(-g_rt[..., 0]) * g_rt[..., 2])
return change_parameterization(g=g_inv_rt, p_from='rotation-translation', p_to=parameterization)
def transform(g, g_parameterization, x, x_parameterization):
"""
Apply rotation g in SE(2) to points x.
"""
g_3x3 = change_parameterization(g, p_from=g_parameterization, p_to='3x3 matrix')
x_homvec = Rn.change_coordinates(x, n=2, p_from=x_parameterization, p_to='homogeneous')
#gx_homvec = g_3x3.dot(x_homvec)
gx_homvec = np.einsum('...ij,...j->...i', g_3x3, x_homvec)
return Rn.change_coordinates(gx_homvec, n=2, p_from='homogeneous', p_to=x_parameterization)
def change_parameterization(g, p_from, p_to):
g = np.array(g)
if p_from == p_to:
return g
if p_from == 'rotation-translation' and p_to == '2x3 matrix':
g_out = np.empty(g.shape[:-1] + (2, 3))
g_out[..., 0, 0] = np.cos(g[..., 0])
g_out[..., 0, 1] = -np.sin(g[..., 0])
g_out[..., 1, 0] = np.sin(g[..., 0])
g_out[..., 1, 1] = np.cos(g[..., 0])
g_out[..., 0, 2] = g[..., 1]
g_out[..., 1, 2] = g[..., 2]
return g_out
if p_from == 'rotation-translation' and p_to == '3x3 matrix':
g_out = np.empty(g.shape[:-1] + (3, 3))
g_out[..., 0, 0] = np.cos(g[..., 0])
g_out[..., 0, 1] = -np.sin(g[..., 0])
g_out[..., 0, 2] = g[..., 1]
g_out[..., 1, 0] = np.sin(g[..., 0])
g_out[..., 1, 1] = np.cos(g[..., 0])
g_out[..., 1, 2] = g[..., 2]
g_out[..., 2, 0] = 0.
g_out[..., 2, 1] = 0.
g_out[..., 2, 2] = 1.
return g_out
else:
raise ValueError('Not supported (yet):' + str(p_from) + ' to ' + str(p_to))
\ No newline at end of file
lie_learn/lie_learn/groups/SO2.py 0 → 100755
View file @ b5881ee2
import numpy as np
parameterizations = ('MAT', 'C', 'ANG')
def compose(g, h, parameterization='MAT', g_parameterization=None, h_parameterization=None, out_parameterization=None):
"""
Compose elements g, h in SO(2).
g and h can have the following parameterizations:
1: MAT 2x2 rotation matrix
2: C 1x1 complex exponential (z=exp(i theta))
"""
if parameterization is not None:
g_parameterization = parameterization
h_parameterization = parameterization
out_parameterization = parameterization
g_mat = change_parameterization(g, p_from=g_parameterization, p_to='MAT')
h_mat = change_parameterization(h, p_from=h_parameterization, p_to='MAT')
gh_mat = np.einsum('...ij,...jk->...ik', g_mat, h_mat)
return change_parameterization(g=gh_mat, p_from='MAT', p_to=out_parameterization)
def invert(g, parameterization):
"""
Invert element g in SO(2), where g can have any supported parameterization:
1: MAT 2x2 rotation matrix
2: C 1x1 complex exponential (z=exp(i theta))
"""
g_mat = change_parameterization(g, p_from=parameterization, p_to='MAT')
g_mat_T = g_mat.transpose(list(range(0, g_mat.ndim - 2)) + [g_mat.ndim - 1, g_mat.ndim - 2]) # Transpose last axes
return change_parameterization(g=g_mat_T, p_from='MAT', p_to=parameterization)
def transform(g, g_parameterization, x, x_parameterization):
"""
Apply rotation g in SO(2) to points x.
"""
#g_mat = change_parameterization(g_parameterization + 'toMAT', g, ichk=0)
#x_vec = change_coordinates(x_parameterization + 'toC', x)
#gx_vec = g_mat.dot(x_vec)
#return change_coordinates('Cto' + x_parameterization, gx_vec)
raise NotImplementedError('SO2 transform not implemented')
def change_parameterization(g, p_from, p_to):
"""
"""
if p_from == p_to:
return g
elif p_from == 'MAT' and p_to == 'C':
theta = np.arctan2(g[..., 1, 0], g[..., 0, 0])
return np.exp(1j * theta)
elif p_from == 'MAT' and p_to == 'ANG':
return np.arctan2(g[..., 1, 0], g[..., 0, 0])
elif p_from == 'C' and p_to == 'MAT':
theta = np.angle(g)
c = np.cos(theta)
s = np.sin(theta)
return np.array([[c, -s], [s, c]]).transpose(list(range(2, 2 + c.ndim)) + [0, 1])
elif p_from == 'C' and p_to == 'ANG':
return np.angle(g)
elif p_from == 'ANG' and p_to == 'MAT':
c = np.cos(g)
s = np.sin(g)
return np.array([[c, -s], [s, c]]).transpose(list(range(2, 2 + c.ndim)) + [0, 1])
elif p_from == 'ANG' and p_to == 'C':
return np.exp(1j * g)
else:
raise ValueError('Unsupported conversion:' + str(p_from) + ' to ' + str(p_to))
\ No newline at end of file
lie_learn/lie_learn/groups/SO3.pyx 0 → 100755
View file @ b5881ee2
# # Notes
# Rotation conventions consist of
# -> orientation of frame: (left- or right-handed) associate x,y,z with thumb, index, middle finger resp. of r or l hand
# -> active or passive / frame-fixed or body-fixed / alibi or alias
# This code uses a Right-hand rule
# (i.e. positive Euler angles indicate clockwise rotations when looking in the positive direction
# along an axis) in a right-handed coordinate frame (thumb, index, middle finger = x, y, z, respectively)
# Furthermore, this code uses a frame-rotation / vector-fixed / alias convention instead of
# the vector-rotation / frame-fixed / alibi convention.
# To obtain the other, simply invert the rotation operator.
##
# Todo:
# -Use quaternions in compose_rotation because they're more stable and cheaper to compose. (do check if this works..)
# -Make compose broadcast as a gufunc
# -change_parameterization could be further optimized
# -implement fix the input checking code in change_parameterizations
# -detect singularities and handle them correctly, so that the output is always acceptable
import numpy as np
cimport numpy as np
import lie_learn.spaces.rn as Rn
FLOAT_TYPE = np.float64
ctypedef np.float64_t FLOAT_TYPE_t
parameterizations = ('MAT', 'Q', 'EV',
'EA123', 'EA132', 'EA213', 'EA231', 'EA312', 'EA321', # Type-I Euler angles AKA Tait-Bryan angles
'EA121', 'EA131', 'EA212', 'EA232', 'EA313', 'EA323') # Type-II Euler angles
def compose(g, h, parameterization='MAT', g_parameterization=None, h_parameterization=None, out_parameterization=None):
"""
Compose (i.e. multiply) elements g, h in SO(3).
g and h can have any parameterization supported by our adaptation of the SpinCalc function:
1: Q Rotation Quaternions
2: EV Euler Vector and rotation angle (degrees)
3: MAT Orthogonal MAT Rotation Matrix
4: EAXXX Euler angles (12 possible sets) (degrees)
"""
if parameterization is not None:
g_parameterization = parameterization
h_parameterization = parameterization
out_parameterization = parameterization
g_mat = change_coordinates(g=g, p_from=g_parameterization, p_to='MAT', perform_checks=False)
h_mat = change_coordinates(g=h, p_from=h_parameterization, p_to='MAT', perform_checks=False)
gh_mat = np.einsum('...ij,...jk->...ik', g_mat, h_mat)
return change_coordinates(g=gh_mat, p_from='MAT', p_to=out_parameterization, perform_checks=False)
def invert(g, parameterization='MAT'):
"""
Invert element g in SO(3), where g can have any supported parameterization:
1: Q Rotation Quaternions
2: EV Euler Vector and rotation angle (degrees)
3: MAT Orthogonal MAT Rotation Matrix
4: EAXXX Euler angles (12 possible sets) (degrees)
"""
g_mat = change_coordinates(g=g, p_from=parameterization, p_to='MAT', perform_checks=False)
g_mat_T = g_mat.transpose(list(range(0, g_mat.ndim - 2)) + [g_mat.ndim - 1, g_mat.ndim - 2]) # Transpose last axes
return change_coordinates(g=g_mat_T, p_from='MAT', p_to=parameterization, perform_checks=False)
def transform_r3(g, x, g_parameterization='MAT', x_parameterization='C'):
"""
Apply rotation g in SO(3) to points x in R^3.
"""
g_mat = change_coordinates(g=g, p_from=g_parameterization, p_to='MAT', perform_checks=False)
x_vec = Rn.change_coordinates(x, n=3, p_from=x_parameterization, p_to='C')
#gx_vec = g_mat.dot(x_vec)
gx_vec = np.einsum('...ij,...j->...i', g_mat, x_vec)
return Rn.change_coordinates(gx_vec, n=3, p_from = 'C', p_to=x_parameterization)
def change_coordinates(g, p_from, p_to, units='rad', perform_checks=False):
"""
Change the coordinates of rotation operators g_in in SO(3).
Parameterizations:
MAT - 3x3 matrix which pre-multiplies by a coordinate frame vector to rotate it to the desired new frame.
EAXXX - [psi,theta,phi] (3,) vector list dictating to the first angle rotation (psi), the second (theta),
and third (phi). XXX stands for the axes of rotation, in order (e.g. XXX=131)
EV - [m1,m2,m3,MU] (4,) vector list dictating the components of euler rotation vector (original coordinate frame),
and the Euler rotation angle about that vector (MU)
Q - [q1,q2,q3,q4] (4,) row vector list defining quaternion of rotation.
q4 = np.cos(MU/2) where MU is Euler rotation angle
:param p_from: parameterization of the input g_in
:param p_to: parameterization to convert to
:param g: input transformations, shape s_input + s_p_from, where s_input is an arbitrary shape
and s_p_from is the shape a single input rotation in the p_from parameterization.
e.g. (3, 3) for rotation matrices, (4,) for quaternions and Euler vectors, and (3,) for Euler angles.
:param units: 'rad' or 'deg'
:param perform_checks: if True (default) this function will perform checks on the input
:return: array shape (s_input, s_p_to)
"""
g = np.asarray(g)
if perform_checks:
raise NotImplementedError('Need to update checks code to handle arrays g_in with many elements')
if p_from[:2] == 'EA': # if input type is Euler angles, determine the order of rotations
if units == 'deg':
# If input is in degrees convert it to radians.
# It will be converted back to degrees for output.
g = g * np.pi / 180.
rot_1_in = int(p_from[2])
rot_2_in = int(p_from[3])
rot_3_in = int(p_from[4])
# check that all orders are between 1 and 3
if rot_1_in < 1 or rot_2_in < 1 or rot_3_in < 1 or rot_1_in > 3 or rot_2_in > 3 or rot_3_in > 3:
raise ValueError('Invalid input Euler angle order type (conversion string)')
# check that no 2 consecutive orders are equal (invalid)
elif rot_1_in == rot_2_in or rot_2_in == rot_3_in:
raise ValueError('Invalid input Euler angle order type (conversion string)')
# check first input dimensions
if g.shape[-1] != 3:
raise ValueError('Input euler angle data vector should have last dimension of length 3')
# identify np.singularities
if rot_1_in == rot_3_in: # Type 2 rotation (first and third rotations about same axis)
if (g[1] <= 0) or (g[1] >= np.pi): # confirm second angle within range
raise ValueError('Second input Euler angle(s) outside 0 to 180 degree range')
elif (np.abs(g[1]) < 2) or np.abs((g[1]) > (np.pi - 0.035)): # check for np.singularity 88 deg
if perform_checks:
print('Warning: Input Euler angle rotation(s) near a singularity. ' +
'Second angle near 0 or 180 degrees.')
else: # Type 1 rotation (all rotations about each of three axes)
if np.abs(g[1]) >= np.pi / 2: # confirm second angle within range
raise ValueError('Second input Euler angle(s) outside -90 to 90 degree range')
elif np.abs(g[1]) > (np.pi / 2 - 0.035): # check for np.singularity
if perform_checks:
print('Warning: Input Euler angle(s) rotation near a singularity. ' +
'Second angle near -90 or 90 degrees.')
if p_from == 'MAT':
if g.shape != (3, 3):
raise ValueError('Matrix is not 3x3. It is:' + str(g.shape))
# Check if matrix is indeed special orthogonal
if not np.isclose(np.sum(np.abs(g.dot(g.T) - np.eye(3))), 0):
raise ValueError('g_in not orthogonal')
if not np.isclose(np.abs(np.linalg.det(g) - 1), 0.0):
raise ValueError('g_in not a proper rotation')
if p_from == 'EV': # Euler Vector Input Type
if g.shape[0] != 4: # or INPUT.shape[2]!=1: # check dimensions
raise ValueError('Input euler vector and rotation data matrix is not Nx4')
if units == 'deg':
mu = g[3] * np.pi / 180 # assign mu name for clarity
else:
mu = g[3]
# check that input m's constitute unit vector
if not np.isclose(np.sum(g ** 2), 1.0):
raise ValueError('Input Euler vector is not unit length')
if mu < 0 or mu > 2 * np.pi: # check if rotation about euler vector is between 0 and 360
print('Warning: Input euler rotation angle(s) not between 0 and 360 degrees')
if p_from == 'Q':
if g.shape[0] != 4:
raise ValueError('Input quaternion matrix is not 4xN')
if not np.isclose(np.abs(np.sqrt(g[0] ** 2 + g[1] ** 2 + g[2] ** 2 + g[3] ** 2)), 1.):
print('Warning: Input quaternion norm(s) deviate(s) from unity by more than tolerance')
if p_to[:2] == 'EA': # if output type is Euler angles, determine order of rotations
rot_1_out = int(p_to[2])
rot_2_out = int(p_to[3])
rot_3_out = int(p_to[4])
if rot_1_out < 1 or rot_2_out < 1 or rot_3_out < 1 or rot_1_out > 3 or rot_2_out > 3 or rot_3_out > 3:
raise ValueError('Invalid output Euler angle order type (conversion string).')
elif rot_1_out == rot_2_out or rot_2_out == rot_3_out:
raise ValueError('Invalid output Euler angle order type (conversion string).')
# Convert inputs to quaternions.
# The output will be calculated in the second portion of the code from these quaternions.
if p_from == 'MAT':
in_shape = g.shape[:-2]
g = g.reshape(-1, 3, 3)
q = rotation_matrix_to_quaternion(g)
elif p_from == 'EV': # Euler Vector Input Type
in_shape = g.shape[:-1]
g = g.reshape(-1, 4)
if units == 'deg':
mu = g[:, 3] * np.pi / 180
else:
mu = g[:, 3]
# Construct the quaternion:
half_mu = 0.5 * mu
s = np.sin(half_mu)
c = np.cos(half_mu)
q = np.asarray([g[:, 0] * s,
g[:, 1] * s,
g[:, 2] * s,
c]).T
elif p_from[:2] == 'EA':
in_shape = g.shape[:-1]
g = g.reshape(-1, 3)
if units == 'deg':
# convert from deg to radians if necessary
psi = g[:, 0] * np.pi / 180.
theta = g[:, 1] * np.pi / 180.
phi = g[:, 2] * np.pi / 180.
else:
psi = g[:, 0]
theta = g[:, 1]
phi = g[:, 2]
# Pre-calculate cosines and sines of the half-angles for conversion.
c1 = np.cos(0.5 * psi)
c2 = np.cos(0.5 * theta)
c3 = np.cos(0.5 * phi)
s1 = np.sin(0.5 * psi)
s2 = np.sin(0.5 * theta)
s3 = np.sin(0.5 * phi)
c13 = np.cos(0.5 * (psi + phi))
s13 = np.sin(0.5 * (psi + phi))
c1_3 = np.cos(0.5 * (psi - phi))
s1_3 = np.sin(0.5 * (psi - phi))
c3_1 = np.cos(0.5 * (phi - psi))
s3_1 = np.sin(0.5 * (phi - psi))
p_from_axis_order = int(p_from[2:])
if p_from_axis_order == 121:
q = np.asarray([c2 * s13, s2 * c1_3, s2 * s1_3, c2 * c13]).T
elif p_from_axis_order == 232:
q = np.asarray([s2 * s1_3, c2 * s13, s2 * c1_3, c2 * c13]).T
elif p_from_axis_order == 313:
q = np.asarray([s2 * c1_3, s2 * s1_3, c2 * s13, c2 * c13]).T
elif p_from_axis_order == 131:
q = np.asarray([c2 * s13, s2 * s3_1, s2 * c3_1, c2 * c13]).T
elif p_from_axis_order == 212:
q = np.asarray([s2 * c3_1, c2 * s13, s2 * s3_1, c2 * c13]).T
elif p_from_axis_order == 323:
q = np.asarray([s2 * s3_1, s2 * c3_1, c2 * s13, c2 * c13]).T
elif p_from_axis_order == 123:
q = np.asarray([s1 * c2 * c3 + c1 * s2 * s3, c1 * s2 * c3 - s1 * c2 * s3,
c1 * c2 * s3 + s1 * s2 * c3, c1 * c2 * c3 - s1 * s2 * s3]).T
elif p_from_axis_order == 231:
q = np.asarray([c1 * c2 * s3 + s1 * s2 * c3, s1 * c2 * c3 + c1 * s2 * s3,
c1 * s2 * c3 - s1 * c2 * s3, c1 * c2 * c3 - s1 * s2 * s3]).T
elif p_from_axis_order == 312:
q = np.asarray([c1 * s2 * c3 - s1 * c2 * s3, c1 * c2 * s3 + s1 * s2 * c3,
s1 * c2 * c3 + c1 * s2 * s3, c1 * c2 * c3 - s1 * s2 * s3]).T
elif p_from_axis_order == 132:
q = np.asarray([s1 * c2 * c3 - c1 * s2 * s3, c1 * c2 * s3 - s1 * s2 * c3,
c1 * s2 * c3 + s1 * c2 * s3, c1 * c2 * c3 + s1 * s2 * s3]).T
elif p_from_axis_order == 213:
q = np.asarray([c1 * s2 * c3 + s1 * c2 * s3, s1 * c2 * c3 - c1 * s2 * s3,
c1 * c2 * s3 - s1 * s2 * c3, c1 * c2 * c3 + s1 * s2 * s3]).T
elif p_from_axis_order == 321:
q = np.asarray([c1 * c2 * s3 - s1 * s2 * c3, c1 * s2 * c3 + s1 * c2 * s3,
s1 * c2 * c3 - c1 * s2 * s3, c1 * c2 * c3 + s1 * s2 * s3]).T
else:
raise ValueError('Invalid input Euler angle order type (conversion string)')
elif p_from == 'Q':
in_shape = g.shape[:-1]
g = g.reshape(-1, 4)
q = g
# Normalize quaternions in case of deviation from unity.
# User has already been warned of deviation.
q = q / np.sqrt((q ** 2).sum(axis=1))[:, None]
# Convert the quaternion that represents g_in to the desired output parameterization
if p_to == 'MAT':
output = np.asarray([q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2 + q[:, 3] ** 2,
2 * (q[:, 0] * q[:, 1] + q[:, 2] * q[:, 3]),
2 * (q[:, 0] * q[:, 2] - q[:, 1] * q[:, 3]),
2 * (q[:, 0] * q[:, 1] - q[:, 2] * q[:, 3]),
-q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2 + q[:, 3] ** 2,
2 * (q[:, 1] * q[:, 2] + q[:, 0] * q[:, 3]),
2 * (q[:, 0] * q[:, 2] + q[:, 1] * q[:, 3]),
2 * (q[:, 1] * q[:, 2] - q[:, 0] * q[:, 3]),
-q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2 + q[:, 3] ** 2])
output = output.reshape((3, 3, -1)).transpose(2, 0, 1)
output = output.reshape(in_shape + (3, 3))
elif p_to == 'EV':
mu = 2 * np.arctan2(np.sqrt((q[:, 0:3] ** 2).sum(axis=1)), q[:, 3])
N = q.shape[0]
output = np.empty((N, 4))
for i in range(N):
if np.sin(0.5 * mu[i]) != 0:
if units == 'deg':
output[i, :] = np.asarray([q[i, 0] / np.sin(0.5 * mu[i]),
q[i, 1] / np.sin(0.5 * mu[i]),
q[i, 2] / np.sin(0.5 * mu[i]),
mu[i] * 180. / np.pi])
else:
output[i, :] = np.asarray([q[i, 0] / np.sin(0.5 * mu[i]),
q[i, 1] / np.sin(0.5 * mu[i]),
q[i, 2] / np.sin(0.5 * mu[i]),
mu[i]])
else:
if units == 'deg':
output[i, :] = [1, 0, 0, mu[i] * 180. / np.pi]
else:
output[i, :] = [1, 0, 0, mu[i]]
output = output.reshape(in_shape + (4,))
elif p_to == 'Q':
output = q
output = output.reshape(in_shape + (4,))
elif p_to[:2] == 'EA':
p_to_axis_order = int(p_to[2:])
if p_to_axis_order == 121:
psi = np.arctan2(q[:, 0] * q[:, 1] + q[:, 2] * q[:, 3], q[:, 1] * q[:, 3] - q[:, 0] * q[:, 2])
theta = np.arccos(np.clip(q[:, 3] ** 2 + q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2, -1.0, 1.0))
phi = np.arctan2(q[:, 0] * q[:, 1] - q[:, 2] * q[:, 3],
q[:, 0] * q[:, 2] + q[:, 1] * q[:, 3])
euler_type = 2
elif p_to_axis_order == 232:
psi = np.arctan2(q[:, 0] * q[:, 3] + q[:, 1] * q[:, 2], q[:, 2] * q[:, 3] - q[:, 0] * q[:, 1])
theta = np.arccos(np.clip(q[:, 3] ** 2 - q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2, -1.0, 1.0))
phi = np.arctan2(q[:, 1] * q[:, 2] - q[:, 0] * q[:, 3], q[:, 0] * q[:, 1] + q[:, 2] * q[:, 3])
euler_type = 2
elif p_to_axis_order == 313:
psi = np.arctan2(q[:, 0] * q[:, 2] + q[:, 1] * q[:, 3], q[:, 0] * q[:, 3] - q[:, 1] * q[:, 2])
theta = np.arccos(np.clip(q[:, 3] ** 2 - q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2, -1.0, 1.0))
phi = np.arctan2(q[:, 0] * q[:, 2] - q[:, 1] * q[:, 3], q[:, 0] * q[:, 3] + q[:, 1] * q[:, 2])
euler_type = 2
elif p_to_axis_order == 131:
psi = np.arctan2(q[:, 0] * q[:, 2] - q[:, 1] * q[:, 3], q[:, 0] * q[:, 1] + q[:, 2] * q[:, 3])
theta = np.arccos(np.clip(q[:, 3] ** 2 + q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2, -1.0, 1.0))
phi = np.arctan2(q[:, 0] * q[:, 2] + q[:, 1] * q[:, 3], q[:, 2] * q[:, 3] - q[:, 0] * q[:, 1])
euler_type = 2
elif p_to_axis_order == 212:
psi = np.arctan2(q[:, 0] * q[:, 1] - q[:, 2] * q[:, 3], q[:, 0] * q[:, 3] + q[:, 1] * q[:, 2])
theta = np.arccos(np.clip(q[:, 3] ** 2 - q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2, -1.0, 1.0))
phi = np.arctan2(q[:, 0] * q[:, 1] + q[:, 2] * q[:, 3], q[:, 0] * q[:, 3] - q[:, 1] * q[:, 2])
euler_type = 2
elif p_to_axis_order == 323:
psi = np.arctan2(q[:, 1] * q[:, 2] - q[:, 0] * q[:, 3], q[:, 0] * q[:, 2] + q[:, 1] * q[:, 3])
theta = np.arccos(np.clip(q[:, 3] ** 2 - q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2, -1.0, 1.0))
phi = np.arctan2(q[:, 0] * q[:, 3] + q[:, 1] * q[:, 2], q[:, 1] * q[:, 3] - q[:, 0] * q[:, 2])
euler_type = 2
elif p_to_axis_order == 123:
psi = np.arctan2(2 * (q[:, 0] * q[:, 3] - q[:, 1] * q[:, 2]),
q[:, 3] ** 2 - q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2)
theta = np.arcsin(2 * (q[:, 0] * q[:, 2] + q[:, 1] * q[:, 3]))
phi = np.arctan2(2 * (q[:, 2] * q[:, 3] - q[:, 0] * q[:, 1]),
q[:, 3] ** 2 + q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2)
euler_type = 1
elif p_to_axis_order == 231:
psi = np.arctan2(2 * (q[:, 1] * q[:, 3] - q[:, 0] * q[:, 2]),
q[:, 3] ** 2 + q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2)
theta = np.arcsin(2 * (q[:, 0] * q[:, 1] + q[:, 2] * q[:, 3]))
phi = np.arctan2(2 * (q[:, 0] * q[:, 3] - q[:, 2] * q[:, 1]),
q[:, 3] ** 2 - q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2)
euler_type = 1
elif p_to_axis_order == 312:
psi = np.arctan2(2 * (q[:, 2] * q[:, 3] - q[:, 0] * q[:, 1]),
q[:, 3] ** 2 - q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2)
theta = np.arcsin(2 * (q[:, 0] * q[:, 3] + q[:, 1] * q[:, 2]))
phi = np.arctan2(2 * (q[:, 1] * q[:, 3] - q[:, 2] * q[:, 0]),
q[:, 3] ** 2 - q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2)
euler_type = 1
elif p_to_axis_order == 132:
psi = np.arctan2(2 * (q[:, 0] * q[:, 3] + q[:, 1] * q[:, 2]),
q[:, 3] ** 2 - q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2)
theta = np.arcsin(2 * (q[:, 2] * q[:, 3] - q[:, 0] * q[:, 1]))
phi = np.arctan2(2 * (q[:, 0] * q[:, 2] + q[:, 1] * q[:, 3]),
q[:, 3] ** 2 + q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2)
euler_type = 1
elif p_to_axis_order == 213:
psi = np.arctan2(2 * (q[:, 0] * q[:, 2] + q[:, 1] * q[:, 3]),
q[:, 3] ** 2 - q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2)
theta = np.arcsin(2 * (q[:, 0] * q[:, 3] - q[:, 1] * q[:, 2]))
phi = np.arctan2(2 * (q[:, 0] * q[:, 1] + q[:, 2] * q[:, 3]),
q[:, 3] ** 2 - q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2)
euler_type = 1
elif p_to_axis_order == 321:
psi = np.arctan2(2 * (q[:, 0] * q[:, 1] + q[:, 2] * q[:, 3]),
q[:, 3] ** 2 + q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2)
theta = np.arcsin(2 * (q[:, 1] * q[:, 3] - q[:, 0] * q[:, 2]))
phi = np.arctan2(2 * (q[:, 0] * q[:, 3] + q[:, 2] * q[:, 1]),
q[:, 3] ** 2 - q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2)
euler_type = 1
else:
raise ValueError('Invalid output Euler angle order type (p_to).')
if units == 'deg':
output = np.mod(np.asarray([psi, theta, phi]) * 180. / np.pi, 360).T
else:
output = np.mod(np.asarray([psi, theta, phi]), 2 * np.pi).T
output = output.reshape(in_shape + (3,))
if perform_checks:
if euler_type == 1:
sing_chk = np.abs(theta) > (np.pi / 2 - 0.0017) # 89.9 deg
#sing_chk=sing_chk[sing_chk>0]
if sing_chk:
print ('Exception: Input rotation %s s resides too close to' \
'Type 1 Euler singularity.\nType 1 ' \
'Euler singularity occurs when second ' \
'angle is -90 or 90 degrees.\nPlease choose ' \
'different output type.' % str(sing_chk))
elif euler_type == 2:
# TODO
print("Euler Type 2 Singularity Check not implemented")
#sing_chk=[find(np.abs(theta*180/np.pi)<0.1)find(np.abs(theta*180/np.pi-180)<0.1)find(np.abs(theta*180/np.pi-360))<0.1]
#sing_chk=sort(sing_chk(sing_chk>0))
#if size(sing_chk,1)>=1:
# print ('Exception: Input rotation ## s resides too close to Type 2 Euler np.singularity.\nType 2 Euler np.singularity occurs when second angle is 0 or 180 degrees.\nPlease choose different output type.',str(sing_chk(1,1)))
return output
cdef rotation_matrix_to_quaternion(np.ndarray[FLOAT_TYPE_t, ndim=3] g_in):
"""
Convert an array of 3x3 rotation matrices to an array of quaternions.
:param g_in: a numpy array of shape (N, 3, 3), where N is the number of rotation matrices.
:return: a numpy array of shape (N, 4) containing quaternions.
"""
cdef int N = g_in.shape[0]
cdef np.ndarray[FLOAT_TYPE_t, ndim=2] q = np.empty((N, 4))
cdef np.ndarray[FLOAT_TYPE_t, ndim=2] denom = np.empty_like(q)
denom[: ,0] = 0.5 * np.sqrt(np.abs(1 + g_in[:, 0, 0] - g_in[:, 1, 1] - g_in[:, 2, 2]))
denom[:, 1] = 0.5 * np.sqrt(np.abs(1 - g_in[:, 0, 0] + g_in[:, 1, 1] - g_in[:, 2, 2]))
denom[:, 2] = 0.5 * np.sqrt(np.abs(1 - g_in[:, 0, 0] - g_in[:, 1, 1] + g_in[:, 2, 2]))
denom[:, 3] = 0.5 * np.sqrt(np.abs(1 + g_in[:, 0, 0] + g_in[:, 1, 1] + g_in[:, 2, 2]))
case = np.argmax(denom, axis=1)
for i in range(N):
if case[i] == 0:
q[i, 0] = denom[i, 0]
q[i, 1] = (g_in[i, 0, 1] + g_in[i, 1, 0]) / (4 * q[i, 0])
q[i, 2] = (g_in[i, 0, 2] + g_in[i, 2, 0]) / (4 * q[i, 0])
q[i, 3] = (g_in[i, 1, 2] - g_in[i, 2, 1]) / (4 * q[i, 0])
elif case[i] == 1:
q[i, 1] = denom[i, 1]
q[i, 0] = (g_in[i, 0, 1] + g_in[i, 1, 0]) / (4 * q[i, 1])
q[i, 2] = (g_in[i, 1, 2] + g_in[i, 2, 1]) / (4 * q[i, 1])
q[i, 3] = (g_in[i, 2, 0] - g_in[i, 0, 2]) / (4 * q[i, 1])
elif case[i] == 2:
q[i, 2] = denom[i, 2]
q[i, 0] = (g_in[i, 0, 2] + g_in[i, 2, 0]) / (4 * q[i, 2])
q[i, 1] = (g_in[i, 1, 2] + g_in[i, 2, 1]) / (4 * q[i, 2])
q[i, 3] = (g_in[i, 0, 1] - g_in[i, 1, 0]) / (4 * q[i, 2])
elif case[i] == 3:
q[i, 3] = denom[i, 3]
q[i, 0] = (g_in[i, 1, 2] - g_in[i, 2, 1]) / (4 * q[i, 3])
q[i, 1] = (g_in[i, 2, 0] - g_in[i, 0, 2]) / (4 * q[i, 3])
q[i, 2] = (g_in[i, 0, 1] - g_in[i, 1, 0]) / (4 * q[i, 3])
return q
# The following was found in licence.txt for spincalc.m,
# from on which the change_parameterization function in this file is derived:
#
# Copyright (c) 2011, John Fuller
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in
# the documentation and/or other materials provided with the distribution
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
\ No newline at end of file
lie_learn/lie_learn/groups/SO3_tests.py 0 → 100755
View file @ b5881ee2
import numpy as np
from lie_learn.groups.SO3 import *
def test_change_parameterization():
def is_equal(R1, R2, p):
if p == 'Q':
# Quaternions are only defined up to a sign, so check each row, what sign we need
for i in range(R1.shape[0]):
if not (np.allclose(R1[i, ...], R2[i, ...]) or np.allclose(R1[i, ...], -R2[i, ...])):
return False
return True
elif p == 'EV':
# Euler vector (x,y,z,theta) == (-x,-y,-z,-theta mod 2pi)
for i in range(R1.shape[0]):
R2i = np.array([-R2[i, 0], -R2[i, 1], -R2[i, 2], (-R2[i, 3]) % (2 * np.pi)])
if not (np.allclose(R1[i, ...], R2[i, ...]) or np.allclose(R1[i, ], R2i)):
return False
return True
else:
return np.allclose(R1, R2)
for p1 in parameterizations:
for p2 in parameterizations:
# Create two random rotations in 313 Euler angles
R1_EA313 = (np.random.rand(3) * np.array([2 * np.pi, np.pi, 2 * np.pi]))[np.newaxis, :]
R2_EA313 = (np.random.rand(3) * np.array([2 * np.pi, np.pi, 2 * np.pi]))[np.newaxis, :]
R_EA313 = np.r_[R1_EA313, R2_EA313]
R1_p1 = change_coordinates(p_from='EA313', p_to=p1, g=R1_EA313)
R1_p2 = change_coordinates(p_from='EA313', p_to=p2, g=R1_EA313)
R2_p1 = change_coordinates(p_from='EA313', p_to=p1, g=R2_EA313)
R2_p2 = change_coordinates(p_from='EA313', p_to=p2, g=R2_EA313)
R_p1 = change_coordinates(p_from='EA313', p_to=p1, g=R_EA313)
R_p2 = change_coordinates(p_from='EA313', p_to=p2, g=R_EA313)
R1_p2_from_R1_p1 = change_coordinates(p_from=p1, p_to=p2, g=R1_p1)
R1_p1_from_R1_p2 = change_coordinates(p_from=p2, p_to=p1, g=R1_p2)
R2_p2_from_R2_p1 = change_coordinates(p_from=p1, p_to=p2, g=R2_p1)
R2_p1_from_R2_p2 = change_coordinates(p_from=p2, p_to=p1, g=R2_p2)
R_p2_from_R_p1 = change_coordinates(p_from=p1, p_to=p2, g=R_p1)
R_p1_from_R_p2 = change_coordinates(p_from=p2, p_to=p1, g=R_p2)
assert is_equal(R1_p1_from_R1_p2, R1_p1, p1), (
p1 + ' to ' + p2 + ' | R1_p1: ' + str(R1_p1) + ' | R1_p2: ' + str(R1_p2) + ' | R1_p1_from_R1_p2: ' +
str(R1_p1_from_R1_p2))
assert is_equal(R2_p1_from_R2_p2, R2_p1, p1), (
p1 + ' to ' + p2 + ' | R2_p1: ' + str(R2_p1) + ' | R2_p2: ' + str(R2_p2) + ' | R2_p1_from_R2_p2: ' +
str(R2_p1_from_R2_p2))
assert is_equal(R_p1_from_R_p2, R_p1, p1), (
p1 + ' to ' + p2 + ' | R_p1: ' + str(R_p1) + ' | R_p2: ' + str(R_p2) + ' | R_p1_from_R_p2: ' +
str(R_p1_from_R_p2))
assert is_equal(R1_p2_from_R1_p1, R1_p2, p2), (
p1 + ' to ' + p2 + ' | R1_p1: ' + str(R1_p1) + ' | R1_p2: ' + str(R1_p2) + ' | R1_p2_from_R1_p1: ' +
str(R1_p2_from_R1_p1))
assert is_equal(R2_p2_from_R2_p1, R2_p2, p2), (
p1 + ' to ' + p2 + ' | R2_p1: ' + str(R2_p1) + ' | R2_p2: ' + str(R2_p2) + ' | R2_p2_from_R2_p1: ' +
str(R2_p2_from_R2_p1))
assert is_equal(R_p2_from_R_p1, R_p2, p2), (
p1 + ' to ' + p2 + ' | R_p1: ' + str(R_p1) + ' | R_p2: ' + str(R_p2) + ' | R_p2_from_R_p1: ' +
str(R_p2_from_R_p1))
def test_invert():
for p in parameterizations:
R_EA = np.random.rand(4, 5, 6, 3) * np.array([2 * np.pi, np.pi, 2 * np.pi])[None, None, None, :]
R_p = change_coordinates(R_EA, p_from='EA313', p_to=p)
R_p_inv = invert(R_p, parameterization=p)
e = compose(R_p, R_p_inv, parameterization=p)
eM = change_coordinates(e, p_from=p, p_to='MAT')
assert np.isclose(np.sum(eM - np.eye(3)), 0.0), 'not the identity: ' + eM
\ No newline at end of file
lie_learn/lie_learn/probability/HarmonicDensity.py 0 → 100755
View file @ b5881ee2
class HarmonicDensity(object):
pass
lie_learn/lie_learn/probability/S1HarmonicDensity.py 0 → 100755
View file @ b5881ee2
import numpy as np
from scipy.fftpack import rfft, irfft
from scipy.optimize import fmin_l_bfgs_b
"""
samples = vonmises.rvs(kappa=1.0, size=10000)
J = 10; alpha = 10; n = (2 * J + 1) * alpha
even = np.arange(0, 2 * J, 2); odd = np.arange(1, 2 * J, 2)
empirical_moments = np.zeros(2 * J)
empirical_moments[even] = np.mean(np.cos(np.arange(1, J + 1)[np.newaxis, :] * samples[:, np.newaxis]), axis=0)
empirical_moments[odd] = np.mean(np.sin(np.arange(1, J + 1)[np.newaxis, :] * samples[:, np.newaxis]), axis=0)
def logp_and_grad(eta):
# Compute moments:
negative_energy = irfft(np.hstack([[0], eta]), n=n) * (n / 2.)
unnormalized_moments = rfft(np.exp(negative_energy)) / (n / 2.) * np.pi
Z = unnormalized_moments[0]
moments = unnormalized_moments[1:eta.size + 1] / Z
# Compute gradients and log-prob:
grad_logp = empirical_moments - moments
logp = eta.dot(empirical_moments) - np.log(Z)
return -logp, -grad_logp
opt_eta, opt_neg_logp, info = fmin_l_bfgs_b(logp_and_grad, x0=np.zeros(2 * J), iprint=0, factr=1e7, pgtol=1e-5)
print info['task']
print 'Optimum log-likelihood:', -opt_neg_logp
print 'Optimal parameters:', np.round(opt_eta, 2)
"""
class S2HarmonicDensity():
def __init__(self, L_max, oversampling_factor=2):
self.L_max = L_max
self.L_max_os = self.L_max * oversampling_factor
self.even = np.arange(0, 2 * L_max, 2)
self.odd = np.arange(1, 2 * L_max, 2)
self.even_os = np.arange(0, 2 * self.L_max_os, 2)
self.odd_os = np.arange(1, 2 * self.L_max_os, 2)
def negative_energy(self, x, eta):
"""
:param x:
:param eta:
:return:
"""
pass #return eta.dot(sh(self.ls[:, np.newaxis], self.ms[:, np.newaxis],
# x[:, 0][np.newaxis, :], x[:, 1][np.newaxis, :],
# field='real', normalization='quantum', condon_shortley=True))
def moments(self, eta):
"""
:param eta:
:return:
"""
pass
def empirical_moments(self, X, average=True):
"""
Compute the empirical moments of the sample x
:param x: dataset shape (N, 2) for 2 spherical coordinates (theta, phi) per point
:return: the moments 1/N sum_i=1^N T(x_i)
"""
pass
def grad_log_p(self, eta, empirical_moments):
"""
:param eta:
:param M:
:return:
"""
pass
def log_p_and_grad(self, eta, empirical_moments):
"""
"""
pass
def mle_lbfgs(self, empirical_moments, eta_init=None, SigmaInv=None, verbose=True):
# Move to base-class?
if eta_init is None:
eta = np.zeros((self.L_max + 1) ** 2 - 1)
else:
eta = eta_init.copy()
if SigmaInv is None: # No regularization
def objective_and_grad(eta):
logp, grad = self.log_p_and_grad(eta, empirical_moments)
return -logp, -grad
else:
lnZ_prior = 0.5 * SigmaInv.size * np.log(2 * np.pi) - 0.5 * np.sum(np.log(SigmaInv))
def objective_and_grad(eta):
logp, grad = self.log_p_and_grad(eta, empirical_moments)
SigmaInv_eta = SigmaInv * eta
logp += -0.5 * eta.dot(SigmaInv_eta) - lnZ_prior
grad += -SigmaInv_eta
return -logp, -grad
opt_eta, opt_neg_logp, info = fmin_l_bfgs_b(objective_and_grad, x0=eta, iprint=int(verbose) - 1,
factr=1e7, # moderate accuracy
#factr=1e12, # low accuracy
pgtol=1e-5) # norm of proj. grad. to stop iteration at
if verbose:
print('Maximum log prob:', -opt_neg_logp)
print('Optimization info:', info)
# Finally, compute Z:
_, lnZ = self.moments(opt_eta)
return opt_eta, lnZ
def _moment_numerical_integration(self, eta, l, m):
"""
Compute the (l,m)-moment of the density with natural parameter eta using slow numerical integration.
The output of this function should be equal to the *unnormalized* moment as it comes out of the FFT
(without dividing by Z).
:param eta:
:param l:
:param m:
:return:
"""
pass
#f = lambda theta, phi: (np.exp(self.negative_energy(np.array([[theta, phi]]), eta))
# * sh(l, m, theta, phi,
# field='real', normalization='quantum',
# condon_shortley=True))
#return S2.integrate(f)
lie_learn/lie_learn/probability/S2HarmonicDensity.py 0 → 100755
View file @ b5881ee2
import numpy as np
from scipy.optimize import fmin_cg, fmin_l_bfgs_b
from ..spectral.S2FFT_NFFT import S2FFT_NFFT
from ..spaces import S2
from ..representations.SO3.spherical_harmonics import sh
# TODO: add unit-test to check moments agains numerical integration (have done this in terminal)
class S2HarmonicDensity():
def __init__(self, L_max, oversampling_factor=2, fft=None):
# Compute the maximum degree from the length of eta
# sum_l=0^L (2 l + 1) = (L + 1)^2
# so
# L = sqrt(eta.size) - 1
#if (np.sqrt(eta.shape[-1] + 1) != np.sqrt(eta.shape[-1] + 1).astype(int)).any():
# raise ValueError('Incorrect eta: last dimension must be a square.')
#self.L_max = int(np.sqrt(eta.shape[-1] + 1) - 1)
self.L_max = L_max
self.L_max_os = self.L_max * oversampling_factor
# Create arrays containing the (l,m) coordinates
l = [[l] * (2 * l + 1) for l in range(1, self.L_max + 1)]
self.ls = np.array([ll for sublist in l for ll in sublist]) # 1, 1, 1, 2, 2, 2, 2, 2, ...
l_oversampled = [[l] * (2 * l + 1) for l in range(1, self.L_max_os + 1)]
self.ls_oversampled = np.array([ll for sublist in l_oversampled for ll in sublist])
m = [list(range(-l, l + 1)) for l in range(1, self.L_max + 1)]
self.ms = np.array([mm for sublist in m for mm in sublist]) # -1, 0, 1, -2, -1, 0, 1, 2, ...
m_oversampled = [list(range(-l, l + 1)) for l in range(1, self.L_max_os + 1)]
self.ms_oversampled = np.array([mm for sublist in m_oversampled for mm in sublist])
if fft is None:
# Setup a spherical grid and corresponding quadrature weights
convention = 'Clenshaw-Curtis'
#convention = 'Gauss-Legendre'
x = S2.meshgrid(b=self.L_max_os, grid_type=convention)
w = S2.quadrature_weights(b=self.L_max_os, grid_type=convention)
self.fft = S2FFT_NFFT(L_max=self.L_max_os, x=x, w=w)
else:
if fft.L_max < self.L_max_os:
raise ValueError('fft.L_max must be larger than or equal to L_max * oversampling_factor')
self.fft = fft
def negative_energy(self, x, eta):
"""
:param x:
:param eta:
:return:
"""
x = np.atleast_2d(x)
return eta.dot(sh(self.ls[:, np.newaxis], self.ms[:, np.newaxis],
x[:, 0][np.newaxis, :], x[:, 1][np.newaxis, :],
field='real', normalization='quantum', condon_shortley=True))
def sufficient_statistics(self, x):
x = np.atleast_2d(x)
return sh(self.ls[:, np.newaxis], self.ms[:, np.newaxis],
x[:, 0][np.newaxis, :], x[:, 1][np.newaxis, :],
field='real', normalization='quantum', condon_shortley=True)
def moments(self, eta):
"""
:param eta:
:return:
"""
#
eta_os = np.zeros((self.L_max_os + 1) ** 2)
eta_os[1:eta.size + 1] = eta
neg_e = self.fft.synthesize(eta_os)
#unnormalized_moments = self.fft.analyze(np.exp(neg_e))
#return unnormalized_moments[1:] / unnormalized_moments[0], unnormalized_moments[0]
maximum = np.max(neg_e)
unnormalized_moments = self.fft.analyze(np.exp(neg_e - maximum))
#log_unnormalized_moments = np.log(unnormalized_moments + 0j)
#moments = np.exp(log_unnormalized_moments - log_unnormalized_moments[0]).real
#Z = np.exp(log_raw_moments[0] + maximum).real
#lnZ = (log_unnormalized_moments[0] + maximum).real
unnormalized_moments[0] *= np.sqrt(4 * np.pi)
moments = unnormalized_moments / unnormalized_moments[0]
#print np.sum(np.abs(m2 - moments))
lnZ = np.log(unnormalized_moments[0]) + maximum
#print lnZ, lnZ2, lnZ - lnZ2
return moments[1:(self.L_max + 1) ** 2], lnZ
def moments_numint(self, eta):
moments = np.zeros((self.L_max + 1) ** 2)
f = lambda th, ph: np.exp(self.negative_energy([th, ph], eta))
moments[0] = S2.integrate(f, normalize=False)
for l in range(1, self.L_max + 1):
for m in range(-l, l + 1):
print('integrating', l, m)
f = lambda th, ph: np.exp(self.negative_energy([th, ph], eta)) * sh(l, m, th, ph,
field='real', normalization='quantum', condon_shortley=True)
moments[l ** 2 + l + m] = S2.integrate(f, normalize=False)
return moments[1:] / moments[0], moments[0]
def empirical_moments(self, X, average=True):
"""
Compute the empirical moments of the sample x
:param x: dataset shape (N, 2) for 2 spherical coordinates (theta, phi) per point
:return: the moments 1/N sum_i=1^N T(x_i)
"""
# TODO: this can be done potentially more efficiently by computing T(0,0) (the suff. stats. of the north pole),
# and then transforming by D(theta, phi, 0) or something similar. This matrix vector-multiplication can be
# done efficiently by the Pinchon-Hoggan method. (or asymptotically even faster using other methods)
T = sh(self.ls[np.newaxis, :], self.ms[np.newaxis, :],
X[:, 0][:, np.newaxis], X[:, 1][:, np.newaxis],
field='real', normalization='quantum', condon_shortley=True)
if average:
return T.mean(axis=0)
else:
return T
def grad_log_p(self, eta, empirical_moments):
"""
:param eta:
:param M:
:return:
"""
moments, _ = self.moments(eta)
return empirical_moments - moments
def log_p_and_grad(self, eta, empirical_moments):
"""
Compute the gradient of the log probability of the density given by eta,
evaluated at a sample of data summarized by the empirical moments.
The log-prob is:
ln prod_i=1^N p(x_i | eta)
=
sum_i=1^N eta^T T(x_i) - ln Z_eta
=
N (eta^T T_bar - ln Z_eta)
where T_bar = 1/N sum_i=1^N T(x_i) are the empirical moments, as computed by self.empirical_moments(X).
In this function we work with the *average* log-prob, i.e. leaving out the factor N from the log-prob formula.
The gradient is (leaving out the factor of N)
T_bar - E_eta[T(x)]
where E_eta[T(x)] are the moments of p(x|eta), as computed by self.moments(eta).
:param eta: the natural parameters of the distribution
:param empirical_moments: the average sufficient statistics, as computed by self.empirical_moments(X)
:return: the gradient of the average log-prob with respect to eta, and the average log prob itself.
"""
moments, lnZ = self.moments(eta)
grad_logp = empirical_moments - moments
logp = eta.dot(empirical_moments) - lnZ
return logp, grad_logp
def mle_sgd(self, empirical_moments, eta_init=None, learning_rate=0.1, max_iter=1000, verbose=True):
"""
:param X:
:return:
"""
if eta_init is None:
eta = np.zeros((self.L_max + 1) ** 2 - 1)
else:
eta = eta_init.copy()
for i in range(max_iter):
log_p, grad_log_p = self.log_p_and_grad(eta, empirical_moments)
eta += learning_rate * grad_log_p
if verbose:
print('log-prob:', log_p)
# Finally, compute Z:
_, lnZ = self.moments(eta)
return eta, lnZ
def mle_lbfgs(self, empirical_moments, eta_init=None, SigmaInv=None, verbose=True):
if eta_init is None:
eta = np.zeros((self.L_max + 1) ** 2 - 1)
else:
eta = eta_init.copy()
if SigmaInv is None: # No regularization
def objective_and_grad(eta):
logp, grad = self.log_p_and_grad(eta, empirical_moments)
return -logp, -grad
else:
#lnZ_prior = 0.5 * SigmaInv.size * np.log(2 * np.pi) - 0.5 * np.sum(np.log(SigmaInv))
def objective_and_grad(eta):
logp, grad = self.log_p_and_grad(eta, empirical_moments)
SigmaInv_eta = SigmaInv * eta
logp += -0.5 * eta.dot(SigmaInv_eta) # - lnZ_prior
grad += -SigmaInv_eta
return -logp, -grad
opt_eta, opt_neg_logp, info = fmin_l_bfgs_b(objective_and_grad, x0=eta, iprint=int(verbose) - 1,
#factr=1e7, # moderate accuracy
factr=1e12, # low accuracy
pgtol=1e-4,
maxiter=1000) # norm of proj. grad. to stop iteration at
if verbose:
print('Maximum log prob:', -opt_neg_logp)
print('Optimization info:', info['warnflag'], info['task'], np.mean(info['grad']))
# Finally, compute Z:
_, lnZ = self.moments(opt_eta)
return opt_eta, lnZ
def mle_cg(self, empirical_moments, eta_init=None, verbose=True):
if eta_init is None:
eta = np.zeros((self.L_max + 1) ** 2 - 1)
else:
eta = eta_init.copy()
def objective(eta):
logp, _ = self.log_p_and_grad(eta, empirical_moments)
return -logp
def grad(eta):
_, grad = self.log_p_and_grad(eta, empirical_moments)
return -grad
eta_min, logp_min, fun_calls, grad_calls, warnflag = fmin_cg(f=objective, fprime=grad, x0=eta,
full_output=True)
if verbose:
print('min log p:', logp_min)
print('fun_calls:', fun_calls)
print('grad_calls:', grad_calls)
print('warnflag:', warnflag)
#print 'allvecs:', allvecs
# Finally, compute Z:
_, lnZ = self.moments(eta_min)
return eta_min, lnZ
def _moment_numerical_integration(self, eta, l, m):
"""
Compute the (l,m)-moment of the density with natural parameter eta using slow numerical integration.
The output of this function should be equal to the *unnormalized* moment as it comes out of the FFT
(without dividing by Z).
:param eta:
:param l:
:param m:
:return:
"""
f = lambda theta, phi: (np.exp(self.negative_energy(np.array([[theta, phi]]), eta))
* sh(l, m, theta, phi,
field='real', normalization='quantum',
condon_shortley=True))
return S2.integrate(f)
lie_learn/lie_learn/probability/SO3HarmonicDensity.py 0 → 100755
View file @ b5881ee2
from .HarmonicDensity import HarmonicDensity
class SO3HarmonicDensity(HarmonicDensity):
def __init__(self):
pass
lie_learn/lie_learn/representations/SO3/clebsch_gordan_numerical.py 0 → 100755
View file @ b5881ee2
"""
Compute the Clebsch-Gordan coefficients of SO(3) numerically.
These coefficients specify how a product of Wigner D functions can be re-expressed as a linear combination of
Wigner D functions.
Wikipedia gives the formula:
D^l_mn(a,b,c) D^l'_m'n'(a,b,c) = sum_{L=|l-l'|^(l+l') sum_M=-L^L sum_N=-L^L <lml'm'|LM> <lnl'n'|LN> D^L_MN(a,b,c)
where D are the Wigner D functions and <....|..> are CG coefficients.
For our computations related to the representations of SO(3) on the projective plane, we are most interested in
the case l=l'=1
D^1_mn(a,b,c) D^1_m'n'(a,b,c) = sum_{L=0^2 sum_M=-L^L sum_N=-L^L <lml'm'|LM> <lnl'n'|LN> D^L_MN(a,b,c)
Since we typically employ the Pinchon-Hoggan basis of real spherical harmonics, we cannot use the formulas for the
CG coefficients that are given in the standard references. We could re-do the analysis to find these coefficients
for the real basis, but here we employ a simple numerical approach.
We view the products coefficients <1m1m'|LM><1n1n'|LN> as unknowns C(m, m', n, n', L, M, N)
We can randomly sample a large number of euler angles g_i = (alpha_i, beta_i, gamma_i), and solve the linear
system for C.
Looking at this system, it appears that the coefficients are ratios of integers or (square?) roots, such as
1/2, 1/3, 1/(2 sqrt(3)), etc.
This is reminiscent of the numbers Pinchon & Hoggan encounter in their J matrix, so that is something to look into.
I've saved dtThe results of the CG computation for l=1 to clebsch_gordan_l1.npy
"""
import numpy as np
from pinchon_hoggan import *
def compute_CG_3D(m1, n1, m2, n2, N=1000):
l = 1
l_min = 0
l_max = 2
num_coefs = sum([(2 * j + 1) ** 2 for j in range(l_min, l_max + 1)])
g = np.random.rand(3, N) * np.array([[2 * np.pi], [np.pi], [2 * np.pi]])
D1 = SO3_irrep(g, 1)[l + m1, l + n1, :]
D2 = SO3_irrep(g, 1)[l + m2, l + n2, :]
target = D1 * D2
A = np.zeros((N, num_coefs))
for i in range(N):
Ds = np.concatenate([SO3_irrep(g[:, i][:, None], j).flatten() for j in range(l_min, l_max + 1)])
A[i, :] = Ds
return A, target, np.linalg.pinv(A).dot(target)
def compute_CG_matrix(N=1000):
CG = np.zeros((1 + 3 * 3 + 5 * 5, 3, 3, 3, 3))
l = 1
for m1 in range(-l, l + 1):
for n1 in range(-l, l + 1):
for m2 in range(-l, l + 1):
for n2 in range(-l, l + 1):
print(m1, n1, m2, n2)
_, _, w = compute_CG_3D(m1, n1, m2, n2, N)
CG[:, l + m1, l + n1, l + m2, l + n2] = w
return CG
if __name__ == '__main__':
CG = compute_CG_matrix(1000)
CG_exact = np.zeros_like(CG)
uniques = [0., 1. / 2., -1. / 2.,
1. / 3., -1. / 3.,
1. / 6., 2. / 3.,
1. / (2 * np.sqrt(3)), -1. / (2 * np.sqrt(3)),
1. / np.sqrt(3), -1. / np.sqrt(3)]
print('Hypothetical exact uniques:')
print(np.sort(uniques))
print('Numerically obtained uniques (rounded to 5 decimals)')
print(np.unique(np.round(CG, 5)))
for value in uniques:
inds = np.nonzero(np.isclose(CG, value))
CG_exact[inds] = value
print('Absolute error between exact and numerical:', np.sum(np.abs(CG_exact - CG)))
\ No newline at end of file
lie_learn/lie_learn/representations/SO3/indexing.py 0 → 100755
View file @ b5881ee2
import numpy as np
def flat_ind_so3(l, m, n):
"""
The SO3 spectrum consists of matrices f_hat^l of size (2l+1, 2l+1) for l=0, ..., L_max.
If we flatten these matrices and stack them, we get a big vector where the element f_hat^l_mn has a certain
flat index. This function computes that index.
The number of elements up to and including order L is
N_L = sum_{l=0}^L (2l+1)^2 = 1/3 (2 L + 1) (2 L + 3) (L + 1)
Element (l, m, n) has N_l elements before it from previous blocks, and in addition several elements in the current
block. The number of elements in the current block, before (m, n) is determined as follows.
First we associate with indices m and n (running from -l to l) with their zero-based index:
m' = m + l
n' = n + l
The linear index of this pair (m', n') is
i = m' * w + n'
where w is the width of the matrix, i.e. w = 2l + 1
The final index of (l, m, n) is N_L + i
:param l, m, n: spectral indices
:return: index of (l, m, n) in flat vector
"""
assert np.abs(m) <= l
assert np.abs(n) <= l
if l == 0:
return 0 # The N_L formula only works for l > 0, so we special case this
L = l - 1
N_L = ((2 * L + 1) * (2 * L + 3) * (L + 1)) // 3
i = (m + l) * (2 * l + 1) + (n + l)
return N_L + i
def flat_ind_zp_so3(l, m, n, b):
"""
The SO3 spectrum consists of matrices f_hat^l of size (2l+1, 2l+1) for l=0, ..., L_max = b - 1.
These can be stored in a zero-padded array A of shape (b, 2b, 2b) with axes l, m, n with zero padding around
the center of the last two axes. If we flatten this array A we get a vector v of size 4b^3.
This function gives the flat index in this array v corresponding to element (l, m, n)
The zero-based 3D index of (l, m, n) in A is (l, b + m, b + n).
The corresponding flat index is i = l * 4b^2 + (b + m) * 2b + b + n
:param l, m, n: spectral indices
:return: index of (l, m, n) in flat zero-padded vector
"""
return l * 4 * (b ** 2) + (b + m) * 2 * b + b + n
def list_to_flat(f_hat_list):
"""
A function on the SO(3) spectrum can be represented as:
1. a list f_hat of matrices f_hat[l] of size (2l+1, 2l+1)
2. a flat vector which is the concatenation of the flattened matrices
3. a zero-padded tensor with axes l, m, n.
This function converts 1 to 2.
:param f_hat: a list of matrices
:return: a flat vector
"""
return np.hstack([a.flat for a in f_hat_list])
def num_spectral_coeffs_up_to_order(b):
"""
The SO(3) spectrum consists of matrices of size (2l+1, 2l+1) for l=0, ..., b - 1.
This function computes the number of elements in a spectrum up to (but excluding) b - 1.
The number of elements up to and including order L is
N_L = sum_{l=0}^L (2l+1)^2 = 1/3 (2 L + 1) (2 L + 3) (L + 1)
:param b: bandwidth
:return: the number of spectral coefficients
"""
L_max = b - 1
assert L_max >= 0
return ((2 * L_max + 1) * (2 * L_max + 3) * (L_max + 1)) // 3
def flat_to_list(f_hat_flat, b):
"""
A function on the SO(3) spectrum can be represented as:
1. a list f_hat of matrices f_hat[l] of size (2l+1, 2l+1)
2. a flat vector which is the concatenation of the flattened matrices
3. a zero-padded tensor with axes l, m, n.
This function converts 2 to 1.
:param f_hat: a flat vector
:return: a list of matrices
"""
f_hat_list = []
start = 0
for l in range(b):
f_hat_list.append(f_hat_flat[start:start + (2 * l + 1) ** 2].reshape(2 * l + 1, 2 * l + 1))
start += (2 * l + 1) ** 2
return f_hat_list
lie_learn/lie_learn/representations/SO3/irrep_bases.pyx 0 → 100755
View file @ b5881ee2
"""
There are a number of different bases for the irreducible representations of SO(3),
each of which results in a different form for the irrep matrices.
This file contains routines that produce change-of-basis matrices
to take you from one basis to the others.
Recall that all irreducible representations of SO(3) appear in the
decomposition of the regular representations on well-behaved functions
f: S^2 -> C or f : S^2 -> R
from the sphere S^2 to the real or complex numbers.
The regular representation is defined by left translation:
(T(g) f)(h) = f(g^{-1} h)
The most common basis for the irreducible representation of weight l are some
form of *complex* spherical harmonics (CSH) Y_l^m, for -l <= m <= l.
For real functions, one can use real spherical harmonics (RSH) S_l^m,
which have the same indexing scheme and are related to the CSH
by a unitary change of basis.
For both CSH and RSH, there are a number of normalization conventions,
as described in spherical_harmonics.py and in [1]. However, these differ
by either
1) a constant scale factor of sqrt(4 pi), or
2) a scale factor (-1)^m, which is the same for +m and -m.
Since the RSH S_l^m is obtained by a linear combination of complex Y_l^m and Y_l^{-m} (see [1]),
the process of changing normalization and that of changing CSH to RSH commute (we can pull out the scale/phase factor).
Since the CSH-RSH change of basis is a unitary transformation, the change of basis maps each kind of CSH to a kind of
RSH that has the same normalization properties.
When changing the normalization, the change-of-basis matrix need not be unitary.
In particular, all changes in normalization, except quantum <--> seismology, lead to non-unitary matrices.
Besides normalization, the harmonics can be rearanged in different orders than m=-l,...,l
This is useful because the Pinchon-Hoggan J matrix assumes a block structure in a certain ordering.
For each normalization convention, we have the following bases:
- Complex centered (cc): Y^{-l}, ..., Y^{l}
- Real centered (rc): S^{-l}, ..., S^{l}
- Real block Pinchon-Hoggan (rb): this basis is aligned with the subspaces
E_xyz,k (etc.) described by Pinchon & Hoggan, and is obtained by a reordering of the RSH.
In this basis, the Pinchon-Hoggan J matrix has a block structure.
References:
[1] http://en.wikipedia.org/wiki/Spherical_harmonics#Conventions
[2] Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes.
"""
import numpy as np
cimport numpy as np
import collections
from scipy.linalg import block_diag
INT_TYPE = np.int64
ctypedef np.int64_t INT_TYPE_t
FLOAT_TYPE = np.float64
ctypedef np.float64_t FLOAT_TYPE_t
COMPLEX_TYPE = np.complex128
ctypedef np.complex128_t COMPLEX_TYPE_t
def change_of_basis_matrix(l, frm=('complex', 'seismology', 'centered', 'cs'), to=('real', 'quantum', 'centered', 'cs')):
"""
Compute change-of-basis matrix that takes the 'frm' basis to the 'to' basis.
Each basis is identified by:
1) A field (real or complex)
2) A normalization / phase convention ('seismology', 'quantum', 'nfft', or 'geodesy')
3) An ordering convention ('centered', 'block')
4) Whether to use Condon-Shortley phase (-1)^m for m > 0 ('cs', 'nocs')
Let B = change_of_basis_matrix(l, frm, to).
Then if Y is a vector in the frm basis, B.dot(Y) represents the same vector in the to basis.
:param l: the weight (non-negative integer) of the irreducible representation, or an iterable of weights.
:param frm: a 3-tuple (field, normalization, ordering) indicating the input basis.
:param to: a 3-tuple (field, normalization, ordering) indicating the output basis.
:return: a (2 * l + 1, 2 * l + 1) change of basis matrix.
"""
from_field, from_normalization, from_ordering, from_cs = frm
to_field, to_normalization, to_ordering, to_cs = to
if isinstance(l, collections.abc.Iterable):
blocks = [change_of_basis_matrix(li, frm, to)
for li in l]
return block_diag(*blocks)
# First, bring us to the centered basis:
if from_ordering == 'block':
B = _c2b(l).T
elif from_ordering == 'centered':
B = np.eye(2 * l + 1)
else:
raise ValueError('Invalid from_order: ' + str(from_ordering))
# Make sure we're using CS-phase (this should work for both real and complex bases)
if from_cs == 'nocs':
m = np.arange(-l, l + 1)
B = ((-1.) ** (m * (m > 0)))[:, None] * B
elif from_cs != 'cs':
raise ValueError('Invalid from_cs: ' + str(from_cs))
# If needed, change complex to real or real to complex
# (we know how to do that in the centered, CS-phase bases)
if from_field != to_field:
if from_field == 'complex' and to_field == 'real':
B = _cc2rc(l).dot(B)
elif from_field == 'real' and to_field == 'complex':
B = _cc2rc(l).conj().T.dot(B)
else:
raise ValueError('Invalid field:' + str(from_field) + ', ' + str(to_field))
# If needed, change the normalization:
if from_normalization != to_normalization:
# First, change normalization to quantum
if from_normalization == 'seismology':
B = _seismology2quantum(l, full_matrix=False)[:, None] * B
elif from_normalization == 'geodesy':
B = _geodesy2quantum(l, full_matrix=False)[:, None] * B
elif from_normalization == 'nfft':
B = _nfft2quantum(l, full_matrix=False)[:, None] * B
elif from_normalization != 'quantum':
raise ValueError('Invalud from_normalization:' + str(from_normalization))
# We're now in quantum normalization, change to output normalization
if to_normalization == 'seismology':
B = (1. / _seismology2quantum(l, full_matrix=False))[:, None] * B
elif to_normalization == 'geodesy':
B = (1. / _geodesy2quantum(l, full_matrix=False))[:, None] * B
elif to_normalization == 'nfft':
B = (1. / _nfft2quantum(l, full_matrix=False))[:, None] * B
elif to_normalization != 'quantum':
raise ValueError('Invalid to_normalization:' + str(to_normalization))
#if from_field != to_field:
# if from_field == 'complex' and to_field == 'real':
# B = cc2rc(l).dot(B)
# elif from_field == 'real' and to_field == 'complex':
# #B = cc2rc(l).conj().T.dot(B)
# pass
# else:
# raise ValueError('Invalid field:' + str(from_field) + ', ' + str(to_field))
#if to_field == 'real':
# B = cc2rc(l).dot(B)
#elif to_field != 'complex':
# raise ValueError('Invalid to_field: ' + str(to_field))
# Set the correct CS phase
if to_cs == 'nocs':
# We're in CS phase now, so cancel it:
m = np.arange(-l, l + 1)
B = ((-1.) ** (m * (m > 0)))[:, None] * B
elif to_cs != 'cs':
raise ValueError('Invalid to_cs: ' + str(to_cs))
# If needed, change the order from centered:
if to_ordering == 'block':
B = _c2b(l).dot(B)
elif to_ordering != 'centered':
raise ValueError('Invalid to_ordering:' + str(to_ordering))
return B
#TODO: make sure that change_of_basis_function accepts matrices, where each row is a vector to be changed of basis.
def change_of_basis_function(l, frm=('complex', 'seismology', 'centered', 'cs'),
to=('real', 'quantum', 'centered', 'cs')):
"""
Return a function that will compute the change-of-basis that takes the 'frm' basis to the 'to' basis.
Each basis is identified by:
1) A field (real or complex)
2) A normalization / phase convention ('seismology', 'quantum', or 'geodesy')
3) An ordering convention ('centered', 'block')
4) Whether to use Condon-Shortley phase (-1)^m for m > 0 ('cs', 'nocs')
:param l: the weight (non-negative integer) of the irreducible representation, or an iterable of weights.
:param frm: a 3-tuple (field, normalization, ordering) indicating the input basis.
:param to: a 3-tuple (field, normalization, ordering) indicating the output basis.
:return:
"""
from_field, from_normalization, from_ordering, from_cs = frm
to_field, to_normalization, to_ordering, to_cs = to
if not isinstance(l, np.ndarray): # collections.Iterable):
l = np.atleast_1d(np.array(l))
# First, bring us to the centered basis:
if from_ordering == 'block':
f1 = _b2c_func(l)
elif from_ordering == 'centered':
f1 = lambda x: x
else:
raise ValueError('Invalid from_order: ' + str(from_ordering))
ms = np.zeros(np.sum(2 * l + 1), dtype=INT_TYPE)
ls = np.zeros(np.sum(2 * l + 1), dtype=INT_TYPE)
i = 0
for ll in l:
for mm in range(-ll, ll + 1):
ms[i] = mm
ls[i] = ll
i += 1
# Make sure we're using CS-phase (this should work for both real and complex bases)
if from_cs == 'nocs':
p = ((-1.) ** (ms * (ms > 0)))
f2 = lambda x: f1(x) * p
elif from_cs == 'cs':
f2 = f1
else: # elif from_cs != 'cs':
raise ValueError('Invalid from_cs: ' + str(from_cs))
# If needed, change complex to real or real to complex
# (we know how to do that in the centered, CS-phase bases)
if from_field != to_field:
if from_field == 'complex' and to_field == 'real':
#B = _cc2rc(l).dot(B)
#pos_m = m > 0
#neg_m = m < 0
#zero_m = m == 0
#f3 = lambda x: r(f2(x) * _cc2rc_func(x, m), 3)
f3 = lambda x: _cc2rc_func(f2(x), ms)
elif from_field == 'real' and to_field == 'complex':
f3 = lambda x: _rc2cc_func(f2(x), ms)
#raise NotImplementedError('Real to complex not implemented yet')
else:
raise ValueError('Invalid field:' + str(from_field) + ', ' + str(to_field))
else:
f3 = f2
# If needed, change the normalization:
if from_normalization != to_normalization:
# First, change normalization to quantum
if from_normalization == 'seismology':
f4 = lambda x: f3(x) * _seismology2quantum(l, full_matrix=False)
elif from_normalization == 'geodesy':
f4 = lambda x: f3(x) * _geodesy2quantum(l, full_matrix=False)
elif from_normalization == 'nfft':
f4 = lambda x: f3(x) * _nfft2quantum(l, full_matrix=False)
elif from_normalization == 'quantum':
f4 = f3
else: # elif from_normalization != 'quantum':
raise ValueError('Invalud from_normalization:' + str(from_normalization))
# We're now in quantum normalization, change to output normalization
if to_normalization == 'seismology':
f5 = lambda x: f4(x) / _seismology2quantum(l, full_matrix=False)
elif to_normalization == 'geodesy':
f5 = lambda x: f4(x) / _geodesy2quantum(l, full_matrix=False)
elif to_normalization == 'nfft':
f5 = lambda x: f4(x) / _nfft2quantum(l, full_matrix=False)
elif to_normalization == 'quantum':
f5 = f4
else: # elif to_normalization != 'quantum':
raise ValueError('Invalid to_normalization:' + str(to_normalization))
else:
f5 = f3
# Set the correct CS phase
if to_cs == 'nocs':
# We're in CS phase now, so cancel it:
#m = np.arange(-l, l + 1)
#B = ((-1.) ** (m * (m > 0)))[:, None] * B
p = ((-1.) ** (ms * (ms > 0)))
f6 = lambda x: f5(x) * p
elif to_cs == 'cs':
f6 = f5
elif to_cs != 'cs':
raise ValueError('Invalid to_cs: ' + str(to_cs))
# If needed, change the order from centered:
if to_ordering == 'block':
#B = _c2b(l).dot(B)
#raise NotImplementedError('Block basis not supported yet')
f7 = lambda x: _c2b_func(l)(f6(x))
elif to_ordering == 'centered':
f7 = f6
else:
raise ValueError('Invalid to_ordering:' + str(to_ordering))
return f7
def _cc2rc(l):
"""
Compute change of basis matrix from the complex centered (cc) basis
to the real centered (rc) basis.
Let Y be a vector of complex spherical harmonics:
Y = (Y^{-l}, ..., Y^0, ..., Y^l)^T
Let S be a vector of real spherical harmonics as defined on the SH wiki page:
S = (S^{-l}, ..., S^0, ..., S^l)^T
Let B = cc2rc(l)
Then S = B.dot(Y)
B is a complex unitary matrix.
Formula taken from:
http://en.wikipedia.org/wiki/Spherical_harmonics#Real_form_2
"""
B = np.zeros((2 * l + 1, 2 * l + 1), dtype=complex)
for m in range(-l, l + 1):
for n in range(-l, l + 1):
row_ind = m + l
col_ind = n + l
if m == 0 and n == 0:
B[row_ind, col_ind] = np.sqrt(2)
if m > 0 and m == n:
B[row_ind, col_ind] = (-1.) ** m
elif m > 0 and m == -n:
B[row_ind, col_ind] = 1.
elif m < 0 and m == n:
B[row_ind, col_ind] = 1j
elif m < 0 and m == -n:
B[row_ind, col_ind] = -1j * ((-1.) ** m)
return (1.0 / np.sqrt(2)) * B
def _cc2rc_func(np.ndarray[COMPLEX_TYPE_t, ndim=1] x,
np.ndarray[INT_TYPE_t, ndim=1] m_arr):
"""
Compute change of basis from the complex centered (cc) basis
to the real centered (rc) basis.
Let Y be a vector of complex spherical harmonics:
Y = (Y^{-l}, ..., Y^0, ..., Y^l)^T
Let S be a vector of real spherical harmonics as defined on the SH wiki page:
S = (S^{-l}, ..., S^0, ..., S^l)^T
Let B = cc2rc(l)
Then S = B.dot(Y)
B is a complex unitary matrix.
Formula taken from:
http://en.wikipedia.org/wiki/Spherical_harmonics#Real_form_2
"""
cdef int i = 0
cdef np.ndarray[FLOAT_TYPE_t, ndim=1] x_out = np.empty(x.size)
cdef double sq2 = np.sqrt(2)
cdef double isq2 = 1. / sq2
for i in range(m_arr.size):
m = m_arr[i]
if m > 0:
x_out[i] = ((-1.) ** m * x[i] + x[i - 2 * m]).real * isq2
elif m < 0:
x_out[i] = (1j * x[i] - 1j * ((-1.) ** m) * x[i - 2 * m]).real * isq2
else:
x_out[i] = x[i].real
return x_out
def _rc2cc_func(np.ndarray[FLOAT_TYPE_t, ndim=1] x,
np.ndarray[INT_TYPE_t, ndim=1] m_arr):
"""
Compute change of basis from the real centered (rc) basis
to the complex centered (cc) basis.
Formula taken from:
http://en.wikipedia.org/wiki/Spherical_harmonics#Real_form_2
"""
cdef int i = 0
cdef np.ndarray[COMPLEX_TYPE_t, ndim=1] x_out = np.empty(x.size, dtype=COMPLEX_TYPE)
cdef double sq2 = np.sqrt(2)
cdef double isq2 = 1. / sq2
for i in range(m_arr.size):
m = m_arr[i]
if m > 0:
x_out[i] = ((-1.) ** m * x[i - 2 * m] * 1j + (-1.) ** m * x[i]) * isq2
elif m < 0:
x_out[i] = (-1j * x[i] + x[i - 2 * m]) * isq2
else:
x_out[i] = x[i]
return x_out
def _c2b(l, full_matrix=True):
"""
Compute change of basis matrix from the centered basis to
the Pinchon-Hoggan block basis, in which the Pinchon-Hoggan J matrices
are brought in block form.
Let B = c2b(l)
then B.dot(J_l).dot(B.T) is in block form with 4 blocks,
as described by PH.
"""
k = int(l) // 2
if l % 2 == 0:
# Permutation as defined by Pinchon-Hoggan for 1-based indices,
# and l = 2 k
sigma = np.array([2 * i for i in range(1, 2 * k + 1)]
+ [2 * i - 1 for i in range(1, 2 * k + 2)])
else:
# Permutation as defined by Pinchon-Hoggan for 1-based indices,
# and l = 2 k + 1
sigma = np.array([2 * i for i in range(1, 2 * k + 2)]
+ [2 * i - 1 for i in range(1, 2 * k + 3)])
if full_matrix:
# From permutation array sigma, create a permutation matrix B:
B = np.zeros((2 * l + 1, 2 * l + 1))
B[np.arange(2 * l + 1), sigma - 1] = 1.
return B
else:
return sigma
def _c2b_func(l):
"""
:param l:
:return:
"""
sigma = np.hstack([_c2b(li, full_matrix=False) - 1 for li in l])
i_begin = 0
for li in l:
sigma[i_begin:i_begin + 2 * li + 1] += i_begin
i_begin += 2 * li + 1
f = lambda x: x[sigma]
return f
def _b2c_func(l):
sigma = np.hstack([_c2b(li, full_matrix=False) - 1 for li in l])
i_begin = 0
for li in l:
sigma[i_begin:i_begin + 2 * li + 1] += i_begin
i_begin += 2 * li + 1
sigma_inv = np.argsort(sigma)
f = lambda x: x[sigma_inv]
return f
def _seismology2quantum(l, full_matrix=False):
"""
:param l:
:param full_matrix:
:return:
"""
if isinstance(l, collections.Iterable):
diags = [_seismology2quantum(li, full_matrix=False) for li in l]
diagonal = np.hstack(diags)
if full_matrix:
return np.diag(diagonal)
else:
return diagonal
diagonal = (-np.ones(2 * l + 1)) ** np.arange(-l, l + 1)
if full_matrix:
return np.diag(diagonal)
else:
return diagonal
def _geodesy2quantum(l, full_matrix=False):
if isinstance(l, collections.Iterable):
diags = [_geodesy2quantum(li, full_matrix=False) for li in l]
diagonal = np.hstack(diags)
if full_matrix:
return np.diag(diagonal)
else:
return diagonal
diagonal = (-np.ones(2 * l + 1)) ** np.arange(-l, l + 1)
diagonal /= np.sqrt(4 * np.pi)
if full_matrix:
return np.diag(diagonal)
else:
return diagonal
def _nfft2quantum(l, full_matrix=False):
if isinstance(l, collections.Iterable):
diags = [_nfft2quantum(li, full_matrix=False) for li in l]
diagonal = np.hstack(diags)
if full_matrix:
return np.diag(diagonal)
else:
return diagonal
diagonal = np.ones(2 * l + 1) * np.sqrt((2 * l + 1) / (4. * np.pi))
# nfft only has (-1)^m phase factor for positive m; quantum has both pos and neg, so add phase factor to neg inds:
# -> this is now done using a CS setting
#m = np.arange(-l, l + 1)
#diagonal *= ((-1) ** (m * (m < 0)))
m = np.arange(-l, l + 1)
diagonal *= (-1.) ** m
if full_matrix:
return np.diag(diagonal)
else:
return diagonal
lie_learn/lie_learn/representations/SO3/pinchon_hoggan/download.py 0 → 100755
View file @ b5881ee2
import os
import numpy as np
import requests
def download(url):
base = os.path.basename(url)
path = os.path.join(os.path.dirname(__file__), base)
if not os.path.isfile(path):
with open(path, 'wb') as f:
f.write(requests.get(url).content)
return np.load(path, encoding='latin1', allow_pickle=True)
lie_learn/lie_learn/representations/SO3/pinchon_hoggan/pinchon_hoggan.pyx 0 → 100755
View file @ b5881ee2
"""
Code for rotating spherical harmonics expansions by the method described in:
Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes
D. Pinchon, P. E. Hoggan
All functions in this file assume that we're working in the basis of real, quantum-normalized, Pinchon-Hoggan block
spherical harmonics
This is NOT a user-facing API, so the interface of this file may change.
"""
import numpy as np
cimport numpy as np
cimport cython
from lie_learn.broadcasting import generalized_broadcast
# Load the J-matrices, which are stored in the same folder as this file
#import os
#Jb = np.load(os.path.join(os.path.dirname(__file__), 'J_block_0-478.npy'), allow_pickle=True)
FLOAT_TYPE = np.float64
ctypedef np.float64_t FLOAT_TYPE_t
INT_TYPE = np.int
ctypedef np.int64_t INT_TYPE_t
def apply_rotation_block(g, X, irreps, c2b, J_block, l_max, X_out=None):
X, g = generalized_broadcast([X, g])
out_shape = X.shape
X = X.reshape(-1, X.shape[-1]).copy()
g = g.reshape(-1, g.shape[-1]).copy()
if X_out is None:
X_out = np.empty_like(X)
X_out = X_out.reshape(-1, X.shape[-1])
X_temp = np.empty_like(X_out)
apply_z_rotation_block(g[:, 2], X, irreps, c2b, l_max, X_out=X_out)
apply_J_block(X_out, J_block, X_out=X_temp)
apply_z_rotation_block(g[:, 1], X_temp, irreps, c2b, l_max, X_out=X_out)
apply_J_block(X_out, J_block, X_out=X_temp)
apply_z_rotation_block(g[:, 0], X_temp, irreps, c2b, l_max, X_out=X_out)
return X_out.reshape(out_shape)
@cython.wraparound(False)
@cython.nonecheck(False)
@cython.boundscheck(False)
cdef apply_z_rotation_block(np.ndarray[FLOAT_TYPE_t, ndim=1] angles,
np.ndarray[FLOAT_TYPE_t, ndim=2] X,
np.ndarray[INT_TYPE_t, ndim=1] irreps,
np.ndarray[INT_TYPE_t, ndim=1] c2b,
int l_max,
np.ndarray[FLOAT_TYPE_t, ndim=2] X_out):
"""
Apply the rotation about the z-axis by angle angles[i] to the vector
X[i, :] for all i. The vectors in X are assumed to be in the basis of real
block spherical harmonics, corresponding to the irreps.
In the *centered* basis, the z-axis rotation matrix (called X(angle) in the P&H paper) has a special
form with cosines of different frequencies on the diagonal and sines on the
anti-diagonal. This matrix is very sparse, so constructing it explicitly is
very inefficient. This function applies a batch of such 'cross' matrices,
represented implicitly by the corresponding angles, to a batch of vectors,
without explicitly constructing the matrices.
To do this in the block basis, we use a permutation array c2b that takes an index in the centered basis
and returns the index for the block basis (which is a permutation of the centered basis).
Args:
angles: matrix of angles, shape (num_data, num_angles)
irreps: a list of irrep weights (integers)
c2b: an array of length dim, where c2b[i] is the index of centere basis vector i in the block basis
l_max: an integer that is equal to np.max(irreps)
X: matrix to be rotated, shape (num_data, dim)
X_out: matrix where output will be stored, shape (dim, num_data, num_angles)
Returns:
X_out: matrix of shape (dim, num_data, num_angles), such that the
vector X_out[:, i, j] is the rotation of X[:,i] by angles[i, j].
"""
cdef int irrep_ind # Index into the IRREPS2 matrix
cdef int irrep # The irrep weight
cdef int center # The index of the center element (l=0) of current irrep
cdef int offset # The offset of the current coordinate from the center
cdef int abs_offset # The absolute value of the offset
cdef int offset_sign # The sign of the offset
#cdef int l_max = np.max(irreps)
cdef int Xs0 = X.shape[0]
cdef int Xs1 = X.shape[1]
#cdef int Cs2
cdef int vec
cdef int coord
cdef int angle_ind
cdef int coord1_block
cdef int coord2_block
cdef np.ndarray[FLOAT_TYPE_t, ndim=2] C = np.cos(angles[None, :] * np.arange(l_max + 1)[:, None])
cdef np.ndarray[FLOAT_TYPE_t, ndim=2] S = np.sin(angles[None, :] * np.arange(l_max + 1)[:, None])
#Cs2 = C.shape[2]
# Check that the irrep dimensions sum to the right dimensionality
#assert (2 * irreps + 1).sum() == X.shape[0], "Invalid irrep dimensions"
#assert angles.shape[0] == X.shape[1]
#assert X_out.shape[0] == X.shape[0], "Invalid shape for X_out"
#assert X_out.shape[1] == X.shape[1], "Invalid shape for X_out"
#assert X_out.shape[2] == angles.shape[1], "Invalid shape for X_out"
for vec in range(Xs0): # Xs1):
irrep_ind = 0 # Start at the first irrep
irrep = irreps[irrep_ind]
center = irrep
for coord in range(Xs1): # Xs0): # X.shape[0]):
offset = coord - center
if offset > irrep: # Finished with the current irrep?
irrep_ind += 1 # Go to the next irrep
irrep = irreps[irrep_ind] # Get its weight from the list
center = coord + irrep # Compute the new center
offset = -irrep # equivalent to offset=coord-center;
# Compute the absolute value and sign of the offset
abs_offset = abs(offset)
if offset >= 0:
offset_sign = 1
else:
offset_sign = -1
coord1_block = c2b[coord]
coord2_block = c2b[center - offset]
# Compute the value of the transformed coordinate
# Note: we're always adding *two* values, even when offset=0 and hence there is only
# one non-zero element in that row. This is not a problem because S2(0, vec) is always 0.
#X_out[coord1_block, vec] = C[abs_offset, vec] * X[coord1_block, vec] \
# - offset_sign * S[abs_offset, vec] * X[coord2_block, vec]
X_out[vec, coord1_block] = C[abs_offset, vec] * X[vec, coord1_block] \
- offset_sign * S[abs_offset, vec] * X[vec, coord2_block]
return X_out
@cython.wraparound(False)
@cython.nonecheck(False)
@cython.boundscheck(False)
cdef apply_J_block(np.ndarray[FLOAT_TYPE_t, ndim=2] X,
list J_block,
np.ndarray[FLOAT_TYPE_t, ndim=2] X_out):
"""
Multiply the Pinchon-Hoggan J matrix by a matrix X.
This function uses the J matrix in the Pinchon-Hoggan block-basis. In this basis, the J-matrix is a block matrix
with 4 blocks. This function performs this block-multiplication efficiently.
:param X: numpy array of shape (dim, N) where dim is the total dimension of the representation and N is the number
of vectors to be multiplied by J.
:param J_block: list of list of precomputed J matrices in block form.
:return:
"""
cdef int l
cdef int k
cdef int rep_begin = 0
cdef int b1s
cdef int b1e
cdef int b2s
cdef int b2e
cdef int b3s
cdef int b3e
cdef int b4s
cdef int b4e
cdef int li
# Loop over irreps
for li in range(len(J_block)):
Jl = J_block[li]
k = Jl[0].shape[0]
l = (Jl[0].shape[0] + 2 * Jl[1].shape[0] + Jl[2].shape[0]) // 2
# Determine block begin and end indices
if l % 2 == 0:
# blocks have dimension k, k, k, k+1
b1s = rep_begin
b1e = rep_begin + k
b2s = rep_begin + k
b2e = rep_begin + 2 * k
b3s = rep_begin + 2 * k
b3e = rep_begin + 3 * k
b4s = rep_begin + 3 * k
b4e = rep_begin + 4 * k + 1
#b1 = Jbl[0:k, 0:k]
#b2 = Jbl[k:2 * k, 2 * k:3 * k]
##b3 = Jbl[2 * k: 3 * k, k:2 * k]
#b4 = Jbl[3 * k:, 3 * k:]
else:
# blocks have dimension k, k+1, k+1, k+1
b1s = rep_begin
b1e = rep_begin + k
b2s = rep_begin + k
b2e = rep_begin + 2 * k + 1
b3s = rep_begin + 2 * k + 1
b3e = rep_begin + 3 * k + 2
b4s = rep_begin + 3 * k + 2
b4e = rep_begin + 4 * k + 3
#b1 = Jbl[0:k, 0:k]
#b2 = Jbl[k:2 * k + 1, 2 * k + 1:3 * k + 2]
##b3 = Jbl[2 * k + 1:3 * k + 2, k:2 * k + 1]
#b4 = Jbl[3 * k + 2:, 3 * k + 2:]
# Multiply each block:
X_out[:, b1s:b1e] = np.dot(Jl[0], X[:, b1s:b1e].T).T
X_out[:, b2s:b2e] = np.dot(Jl[1], X[:, b3s:b3e].T).T
X_out[:, b3s:b3e] = np.dot(Jl[1].T, X[:, b2s:b2e].T).T
X_out[:, b4s:b4e] = np.dot(Jl[2], X[:, b4s:b4e].T).T
rep_begin += 2 * l + 1
return X_out
def make_c2b(irreps):
# Centered to block basis
# Maybe put this in separate file irrep_bases?
c2b = np.empty((2 * irreps + 1).sum(), dtype=INT_TYPE)
irrep_begin = 0
for l in irreps:
k = int(l) / 2
if l % 2 == 0:
# Permutation as defined by Pinchon-Hoggan for 1-based indices,
# and l = 2 k
sigma = np.array([2 * i for i in range(1, 2 * k + 1)]
+ [2 * i - 1 for i in range(1, 2 * k + 2)])
else:
# Permutation as defined by Pinchon-Hoggan for 1-based indices,
# and l = 2 k + 1
sigma = np.array([2 * i for i in range(1, 2 * k + 2)]
+ [2 * i - 1 for i in range(1, 2 * k + 3)])
sigma_inv = np.arange(0, 2 * l + 1)[np.argsort(sigma)]
c2b[irrep_begin:irrep_begin + 2 * l + 1] = sigma_inv + irrep_begin
irrep_begin += 2 * l + 1
return c2b
\ No newline at end of file
lie_learn/lie_learn/representations/SO3/pinchon_hoggan/pinchon_hoggan_dense.py 0 → 100755
View file @ b5881ee2
import os
import numpy as np
from scipy.linalg import block_diag
# This code is not very optimized,
# and can never become very efficient because it cannot exploit the sparsity of the J matrix.
# Load the J-matrices, which are stored in the same folder as this file
from .download import download
# J matrices come from this paper
# Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes
# Didier Pinchon1 and Philip E Hoggan2
# https://iopscience.iop.org/article/10.1088/1751-8113/40/7/011/
# Jd = download('https://github.com/AMLab-Amsterdam/lie_learn/releases/download/v1.0/J_dense_0-278.npy')
base = 'J_dense_0-150.npy'
path = os.path.join(os.path.dirname(__file__), base)
Jd = np.load(path, allow_pickle=True)
def SO3_irreps(g, irreps):
global Jd
# First, compute sinusoids at all required frequencies, i.e.
# cos(n x) for n=0, ..., max(irreps)
# sin(n x) for n=-max(irreps), ..., max(irreps)
# where x ranges over the three parameters of SO(3).
# In theory, it may be faster to evaluate cos(x) once and then use
# Chebyshev polynomials to obtain cos(n*x), but in practice this appears
# to be slower than just evaluating cos(n*x).
dim = np.sum(2 * np.array(irreps) + 1)
T = np.empty((dim, dim, g.shape[1]))
for i in range(g.shape[1]):
T[:, :, i] = block_diag(*[rot_mat(g[0, i], g[1, i], g[2, i], l, Jd[l]) for l in irreps])
return T
def SO3_irrep(g, l):
global Jd
g = np.atleast_2d(g)
T = np.empty((2 * l + 1, 2 * l + 1, g.shape[1]))
for i in range(g.shape[1]):
T[:, :, i] = rot_mat(g[0, i], g[1, i], g[2, i], l, Jd[l])
return T # np.squeeze(T)
def z_rot_mat(angle, l):
"""
Create the matrix representation of a z-axis rotation by the given angle,
in the irrep l of dimension 2 * l + 1, in the basis of real centered
spherical harmonics (RC basis in rep_bases.py).
Note: this function is easy to use, but inefficient: only the entries
on the diagonal and anti-diagonal are non-zero, so explicitly constructing
this matrix is unnecessary.
"""
M = np.zeros((2 * l + 1, 2 * l + 1))
inds = np.arange(0, 2 * l + 1, 1)
reversed_inds = np.arange(2 * l, -1, -1)
frequencies = np.arange(l, -l - 1, -1)
M[inds, reversed_inds] = np.sin(frequencies * angle)
M[inds, inds] = np.cos(frequencies * angle)
return M
def rot_mat(alpha, beta, gamma, l, J):
"""
Compute the representation matrix of a rotation by ZYZ-Euler
angles (alpha, beta, gamma) in representation l in the basis
of real spherical harmonics.
The result is the same as the wignerD_mat function by Johann Goetz,
when the sign of alpha and gamma is flipped.
The forementioned function is here:
https://sites.google.com/site/theodoregoetz/notes/wignerdfunction
"""
Xa = z_rot_mat(alpha, l)
Xb = z_rot_mat(beta, l)
Xc = z_rot_mat(gamma, l)
return Xa.dot(J).dot(Xb).dot(J).dot(Xc)
def derivative_z_rot_mat(angle, l):
M = np.zeros((2 * l + 1, 2 * l + 1))
inds = np.arange(0, 2 * l + 1, 1)
reversed_inds = np.arange(2 * l, -1, -1)
frequencies = np.arange(l, -l - 1, -1)
M[inds, reversed_inds] = np.cos(frequencies * angle) * frequencies
M[inds, inds] = -np.sin(frequencies * angle) * frequencies
return M
def derivative_rot_mat(alpha, beta, gamma, l, J):
Xa = z_rot_mat(alpha, l)
Xb = z_rot_mat(beta, l)
Xc = z_rot_mat(gamma, l)
dXa_da = derivative_z_rot_mat(alpha, l)
dXb_db = derivative_z_rot_mat(beta, l)
dXc_dc = derivative_z_rot_mat(gamma, l)
dDda = dXa_da.dot(J).dot(Xb).dot(J).dot(Xc)
dDdb = Xa.dot(J).dot(dXb_db).dot(J).dot(Xc)
dDdc = Xa.dot(J).dot(Xb).dot(J).dot(dXc_dc)
return dDda, dDdb, dDdc
lie_learn/lie_learn/representations/SO3/pinchon_hoggan/pinchon_hoggan_parsing.py 0 → 100755
View file @ b5881ee2
import os
import numpy as np
import scipy.sparse as sp
from ..irrep_bases import change_of_basis_function, change_of_basis_matrix
def convert_all(num_mat_J_folder='/Users/user/Projects/LieLearn/SO3/shrot/legacy/pinchon2/Maple/NumMatJ/',
out_folder='./', l_min=0, l_max=None):
"""
Convert all numMatJ-XXX.dat files in a given folder to produce the following python pickle files:
1) J_dense_0-[l_max].npy:
2) J_sparse_0-[l_max].npy:
3) J_block_0-[l_max].npy
The numMatJ-XXX.dat files are the output of mkNumMatJ.mw
:param num_mat_J_folder:
:return:
"""
J_dense = []
J_sparse = []
J_block = []
if l_max is None:
files = [f for f in os.listdir(num_mat_J_folder) if '.dat' in f]
nums = [int(f[len("numMatJ-"):-len(".dat")]) for f in files]
l_max = np.max(nums)
print('Maximum l found:', l_max)
for l in range(l_min, l_max + 1):
print('Parsing l', l)
filename = os.path.join(num_mat_J_folder, 'numMatJ-' + str(l) + '.dat')
# Obtain the J matrix as a dense numpy array
Jd = parse_J_single_l(filename)
J_dense.append(Jd)
J_sparse.append(sp.csr_matrix(Jd))
print('Saving dense matrices...')
np.save(os.path.join(out_folder, 'J_dense_' + str(l_min) + '-' + str(l_max)), J_dense)
print('Saving sparse matrices...')
np.save(os.path.join(out_folder, 'J_sparse_' + str(l_min) + '-' + str(l_max)), J_sparse)
del J_sparse
print('Converting to block basis...')
J_block = make_block_J(J_dense)
print('Saving block matrices...')
np.save(os.path.join(out_folder, 'J_block_' + str(l_min) + '-' + str(l_max)), J_block)
def parse_J(file):
"""
Loads the numMatJ.dat files provided by Pinchon & Hoggan into numpy matrices.
I saved the result as a compressed pickle, so this function will
only be needed when J matrices for larger l are needed.
"""
f = open(file)
lmax = int(f.readline().split(' ')[1])
matj = [np.zeros((2 * l + 1, 2 * l + 1)) for l in range(lmax)]
# Fill in matrix l=0 and l=1
matj[0][0, 0] = 1.0
matj[1][0, 1] = -1.0
matj[1][1, 0] = -1.0
matj[1][2, 2] = 1.0
# Read out matrices >= 2
for l in range(2, lmax):
# Read and discard l-value:
lval = int(f.readline().split(' ')[1])
assert lval == l
k = l/2
k1 = k
for i in range(k):
for j in range(i):
x = float(f.readline())
matj[l][2*i+1, 2*j+1] = x
matj[l][2*j+1, 2*i+1] = x
x = float(f.readline())
matj[l][2*i+1,2*i+1] = x
if l != 2*k1:
k += 1
for i in range(k):
for j in range(k):
x = float(f.readline())
matj[l][2*k1+2*j+1, 2*i] = x
matj[l][2*i, 2*k1+2*j+1] = x
if l == 2*k1:
k += 1
else:
k1 += 1
for i in range(k):
for j in range(i):
x = float(f.readline())
matj[l][2*k1+2*i, 2*k1+2*j] = x
matj[l][2*k1+2*j, 2*k1+2*i] = x
x = float(f.readline())
matj[l][2*k1+2*i, 2*k1+2*i] = x
f.close()
return matj
def parse_J_single_l(file):
"""
Parse a single numMatJ-XXX.dat file produced by Pinchon & Hoggan's maple script mkNumMatJ.mw
:param file:
:return:
"""
f = open(file)
lval = int(f.readline().split(' ')[1])
matj = np.zeros((2 * lval + 1, 2 * lval + 1))
k = lval / 2
k1 = k
for i in range(k):
for j in range(i):
x = float(f.readline())
matj[2*i+1, 2*j+1] = x
matj[2*j+1, 2*i+1] = x
x = float(f.readline())
matj[2*i+1,2*i+1] = x
if lval != 2 * k1:
k += 1
for i in range(k):
for j in range(k):
x = float(f.readline())
matj[2*k1+2*j+1, 2*i] = x
matj[2*i, 2*k1+2*j+1] = x
if lval == 2 * k1:
k += 1
else:
k1 += 1
for i in range(k):
for j in range(i):
x = float(f.readline())
matj[2*k1+2*i, 2*k1+2*j] = x
matj[2*k1+2*j, 2*k1+2*i] = x
x = float(f.readline())
matj[2*k1+2*i, 2*k1+2*i] = x
f.close()
return matj
def make_block_J(Jd):
"""
Convert a list of J matrices 0 to N, to block form.
We change the basis on each J matrix (which is assumed to be in the real, quantum-normalized, centered basis,
so that it is in the real, quantum-normalized, block basis.
Then, we extract the blocks. There are 4 blocks for each irrep l, but the middle two are transposes of each other,
so we store only 3 blocks. The outer two blocks are symmetric, but this is not exploited.
:param Jd:
:return:
"""
Jb = []
for l in range(len(Jd)):
print('Converting to block matrix. (', l, 'of', len(Jd), ')')
#Bl = c2b(l)
#Jbl = Bl.dot(Jd[l]).dot(Bl.T)
#c2b = change_of_basis_function(l,
# frm=('real', 'quantum', 'centered', 'cs'),
# to=('real', 'quantum', 'block', 'cs'))
#Jbl = c2b(c2b(Jd[l]).T).T
Bl = change_of_basis_matrix(l,
frm=('real', 'quantum', 'centered', 'cs'),
to=('real', 'quantum', 'block', 'cs'))
Jbl = Bl.dot(Jd[l]).dot(Bl.T)
k = l // 2
if l % 2 == 0:
# blocks have dimension k, k, k, k+1
b1 = Jbl[0:k, 0:k]
b2 = Jbl[k:2 * k, 2 * k:3 * k]
#b3 = Jbl[2 * k: 3 * k, k:2 * k]
b4 = Jbl[3 * k:, 3 * k:]
else:
# blocks have dimension k, k+1, k+1, k+1
b1 = Jbl[0:k, 0:k]
b2 = Jbl[k:2 * k + 1, 2 * k + 1:3 * k + 2]
#b3 = Jbl[2 * k + 1:3 * k + 2, k:2 * k + 1]
b4 = Jbl[3 * k + 2:, 3 * k + 2:]
Jb.append([b1, b2, b4])
return Jb
\ No newline at end of file
lie_learn/lie_learn/representations/SO3/spherical_harmonics.py 0 → 100755
View file @ b5881ee2
import numpy as np
from scipy.special import sph_harm, lpmv
try:
from scipy.misc import factorial
except:
from scipy.special import factorial
def sh(l, m, theta, phi, field='real', normalization='quantum', condon_shortley=True):
if field == 'real':
return rsh(l, m, theta, phi, normalization, condon_shortley)
elif field == 'complex':
return csh(l, m, theta, phi, normalization, condon_shortley)
else:
raise ValueError('Unknown field: ' + str(field))
def sh_squared_norm(l, normalization='quantum', normalized_haar=True):
"""
Compute the squared norm of the spherical harmonics.
The squared norm of a function on the sphere is defined as
|f|^2 = int_S^2 |f(x)|^2 dx
where dx is a Haar measure.
:param l: for some normalization conventions, the norm of a spherical harmonic Y^l_m depends on the degree l
:param normalization: normalization convention for the spherical harmonic
:param normalized_haar: whether to use the Haar measure da db sinb or the normalized Haar measure da db sinb / 4pi
:return: the squared norm of the spherical harmonic with respect to given measure
"""
if normalization == 'quantum' or normalization == 'seismology':
# The quantum and seismology spherical harmonics are normalized with respect to the Haar measure
# dmu(theta, phi) = dtheta sin(theta) dphi
sqnorm = 1.
elif normalization == 'geodesy':
# The geodesy spherical harmonics are normalized with respect to the *normalized* Haar measure
# dmu(theta, phi) = dtheta sin(theta) dphi / 4pi
sqnorm = 4 * np.pi
elif normalization == 'nfft':
sqnorm = 4 * np.pi / (2 * l + 1)
else:
raise ValueError('Unknown normalization')
if normalized_haar:
return sqnorm / (4 * np.pi)
else:
return sqnorm
def block_sh_ph(L_max, theta, phi):
"""
Compute all spherical harmonics up to and including degree L_max, for angles theta and phi.
This function is currently rather hacky, but the method used here is very fast and stable, compared
to builtin scipy functions.
:param L_max:
:param theta:
:param phi:
:return:
"""
from .pinchon_hoggan.pinchon_hoggan import apply_rotation_block, make_c2b
from .irrep_bases import change_of_basis_function
irreps = np.arange(L_max + 1)
ls = [[ls] * (2 * ls + 1) for ls in irreps]
ls = np.array([ll for sublist in ls for ll in sublist]) # 0, 1, 1, 1, 2, 2, 2, 2, 2, ...
ms = [list(range(-ls, ls + 1)) for ls in irreps]
ms = np.array([mm for sublist in ms for mm in sublist]) # 0, -1, 0, 1, -2, -1, 0, 1, 2, ...
# Get a vector Y that selects the 0-frequency component from each irrep in the centered basis
# If D is a Wigner D matrix, then D Y is the center column of D, which is equal to the spherical harmonics.
Y = (ms == 0).astype(float)
# Change to / from the block basis (since the rotation code works in that basis)
c2b = change_of_basis_function(irreps,
frm=('real', 'quantum', 'centered', 'cs'),
to=('real', 'quantum', 'block', 'cs'))
b2c = change_of_basis_function(irreps,
frm=('real', 'quantum', 'block', 'cs'),
to=('real', 'quantum', 'centered', 'cs'))
Yb = c2b(Y)
# Rotate Yb:
c2b = make_c2b(irreps)
import os
J_block = np.load(os.path.join(os.path.dirname(__file__), 'pinchon_hoggan', 'J_block_0-278.npy'), allow_pickle=True)
J_block = list(J_block[irreps])
g = np.zeros((theta.size, 3))
g[:, 0] = phi
g[:, 1] = theta
TYb = apply_rotation_block(g=g, X=Yb[np.newaxis, :],
irreps=irreps, c2b=c2b,
J_block=J_block, l_max=np.max(irreps))
print(Yb.shape, TYb.shape)
# Change back to centered basis
TYc = b2c(TYb.T).T # b2c doesn't work properly for matrices, so do a transpose hack
print(TYc.shape)
# Somehow, the SH obtained so far are equal to real, nfft, cs spherical harmonics
# Change to real quantum centered cs
c = change_of_basis_function(irreps,
frm=('real', 'nfft', 'centered', 'cs'),
to=('real', 'quantum', 'centered', 'cs'))
TYc2 = c(TYc)
print(TYc2.shape)
return TYc2
def rsh(l, m, theta, phi, normalization='quantum', condon_shortley=True):
"""
Compute the real spherical harmonic (RSH) S_l^m(theta, phi).
The RSH are obtained from Complex Spherical Harmonics (CSH) as follows:
if m < 0:
S_l^m = i / sqrt(2) * (Y_l^m - (-1)^m Y_l^{-m})
if m == 0:
S_l^m = Y_l^0
if m > 0:
S_l^m = 1 / sqrt(2) * (Y_l^{-m} + (-1)^m Y_l^m)
(see [1])
Various normalizations for the CSH exist, see the CSH() function. Since the CSH->RSH change of basis is unitary,
the orthogonality and normalization properties of the RSH are the same as those of the CSH from which they were
obtained. Furthermore, the operation of changing normalization and that of changeing field
(complex->real or vice-versa) commute, because the ratio c_m of normalization constants are always the same for
m and -m (to see this that this implies commutativity, substitute Y_l^m * c_m for Y_l^m in the above formula).
Pinchon & Hoggan [2] define a different change of basis for CSH -> RSH, but they also use an unusual definition
of CSH. To obtain RSH as defined by Pinchon-Hoggan, use this function with normalization='quantum'.
References:
[1] http://en.wikipedia.org/wiki/Spherical_harmonics#Real_form
[2] Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes.
:param l: non-negative integer; the degree of the CSH.
:param m: integer, -l <= m <= l; the order of the CSH.
:param theta: the colatitude / polar angle,
ranging from 0 (North Pole, (X,Y,Z)=(0,0,1)) to pi (South Pole, (X,Y,Z)=(0,0,-1)).
:param phi: the longitude / azimuthal angle, ranging from 0 to 2 pi.
:param normalization: how to normalize the RSH:
'seismology', 'quantum', 'geodesy'.
these are immediately passed to the CSH functions, and since the change of basis
from CSH to RSH is unitary, the orthogonality and normalization properties are unchanged.
:return: the value of the real spherical harmonic S^l_m(theta, phi)
"""
l, m, theta, phi = np.broadcast_arrays(l, m, theta, phi)
# Get the CSH for m and -m, using Condon-Shortley phase (regardless of whhether CS is requested or not)
# The reason is that the code that changes from CSH to RSH assumes CS phase.
a = csh(l=l, m=m, theta=theta, phi=phi, normalization=normalization, condon_shortley=True)
b = csh(l=l, m=-m, theta=theta, phi=phi, normalization=normalization, condon_shortley=True)
#if m > 0:
# y = np.array((b + ((-1.)**m) * a).real / np.sqrt(2.))
#elif m < 0:
# y = np.array((1j * a - 1j * ((-1.)**(-m)) * b).real / np.sqrt(2.))
#else:
# # For m == 0, the complex spherical harmonics are already real
# y = np.array(a.real)
y = ((m > 0) * np.array((b + ((-1.)**m) * a).real / np.sqrt(2.))
+ (m < 0) * np.array((1j * a - 1j * ((-1.)**(-m)) * b).real / np.sqrt(2.))
+ (m == 0) * np.array(a.real))
if condon_shortley:
return y
else:
# Cancel the CS phase of y (i.e. multiply by -1 when m is both odd and greater than 0)
return y * ((-1.) ** (m * (m > 0)))
def csh(l, m, theta, phi, normalization='quantum', condon_shortley=True):
"""
Compute Complex Spherical Harmonics (CSH) Y_l^m(theta, phi).
Unlike the scipy.special.sph_harm function, we use the common convention that
theta is the polar angle (0 to pi) and phi is the azimuthal angle (0 to 2pi).
The spherical harmonic 'backbone' is:
Y_l^m(theta, phi) = P_l^m(cos(theta)) exp(i m phi)
where P_l^m is the associated Legendre function as defined in the scipy library (scipy.special.sph_harm).
Various normalization factors can be multiplied with this function.
-> seismology: sqrt( ((2 l + 1) * (l - m)!) / (4 pi * (l + m)!) )
-> quantum: (-1)^2 sqrt( ((2 l + 1) * (l - m)!) / (4 pi * (l + m)!) )
-> unnormalized: 1
-> geodesy: sqrt( ((2 l + 1) * (l - m)!) / (l + m)! )
-> nfft: sqrt( (l - m)! / (l + m)! )
The 'quantum' and 'seismology' CSH are normalized so that
<Y_l^m, Y_l'^m'>
=
int_S^2 Y_l^m(theta, phi) Y_l'^m'* dOmega
=
delta(l, l') delta(m, m')
where dOmega is the volume element for the sphere S^2:
dOmega = sin(theta) dtheta dphi
The 'geodesy' convention have unit power, meaning the norm is equal to the surface area of the unit sphere (4 pi)
<Y_l^m, Y_l'^m'> = 4pi delta(l, l') delta(m, m')
So these are orthonormal with respect to the *normalized* Haar measure sin(theta) dtheta dphi / 4pi
On each of these normalizations, one can optionally include a Condon-Shortley phase factor:
(-1)^m (if m > 0)
1 (otherwise)
Note that this is the definition of Condon-Shortley according to wikipedia [1], but other sources call a
phase factor of (-1)^m a Condon-Shortley phase (without mentioning the condition m > 0).
References:
[1] http://en.wikipedia.org/wiki/Spherical_harmonics#Conventions
:param l: non-negative integer; the degree of the CSH.
:param m: integer, -l <= m <= l; the order of the CSH.
:param theta: the colatitude / polar angle,
ranging from 0 (North Pole, (X,Y,Z)=(0,0,1)) to pi (South Pole, (X,Y,Z)=(0,0,-1)).
:param phi: the longitude / azimuthal angle, ranging from 0 to 2 pi.
:param normalization: how to normalize the CSH:
'seismology', 'quantum', 'geodesy', 'unnormalized', 'nfft'.
:return: the value of the complex spherical harmonic Y^l_m(theta, phi)
"""
# NOTE: it seems like in the current version of scipy.special, sph_harm no longer accepts keyword arguments,
# so I'm removing them. I hope the order of args hasn't changed
if normalization == 'quantum':
# y = ((-1.) ** m) * sph_harm(m, l, theta=phi, phi=theta)
y = ((-1.) ** m) * sph_harm(m, l, phi, theta)
elif normalization == 'seismology':
# y = sph_harm(m, l, theta=phi, phi=theta)
y = sph_harm(m, l, phi, theta)
elif normalization == 'geodesy':
# y = np.sqrt(4 * np.pi) * sph_harm(m, l, theta=phi, phi=theta)
y = np.sqrt(4 * np.pi) * sph_harm(m, l, phi, theta)
elif normalization == 'unnormalized':
# y = sph_harm(m, l, theta=phi, phi=theta) / np.sqrt((2 * l + 1) * factorial(l - m) /
# (4 * np.pi * factorial(l + m)))
y = sph_harm(m, l, phi, theta) / np.sqrt((2 * l + 1) * factorial(l - m) /
(4 * np.pi * factorial(l + m)))
elif normalization == 'nfft':
# y = sph_harm(m, l, theta=phi, phi=theta) / np.sqrt((2 * l + 1) / (4 * np.pi))
y = sph_harm(m, l, phi, theta) / np.sqrt((2 * l + 1) / (4 * np.pi))
else:
raise ValueError('Unknown normalization convention:' + str(normalization))
if condon_shortley:
# The sph_harm function already includes CS phase
return y
else:
# Cancel the CS phase in sph_harm (i.e. multiply by -1 when m is both odd and greater than 0)
return y * ((-1.) ** (m * (m > 0)))
# For testing only:
def _naive_csh_unnormalized(l, m, theta, phi):
"""
Compute unnormalized SH
"""
return lpmv(m, l, np.cos(theta)) * np.exp(1j * m * phi)
def _naive_csh_quantum(l, m, theta, phi):
"""
Compute orthonormalized spherical harmonics in a naive way.
"""
return (((-1.) ** m) * lpmv(m, l, np.cos(theta)) * np.exp(1j * m * phi) *
np.sqrt(((2 * l + 1) * factorial(l - m))
/
(4 * np.pi * factorial(l + m))))
def _naive_csh_seismology(l, m, theta, phi):
"""
Compute the spherical harmonics according to the seismology convention, in a naive way.
This appears to be equal to the sph_harm function in scipy.special.
"""
return (lpmv(m, l, np.cos(theta)) * np.exp(1j * m * phi) *
np.sqrt(((2 * l + 1) * factorial(l - m))
/
(4 * np.pi * factorial(l + m))))
def _naive_csh_ph(l, m, theta, phi):
"""
CSH as defined by Pinchon-Hoggan. Same as wikipedia's quantum-normalized SH = naive_Y_quantum()
"""
if l == 0 and m == 0:
return 1. / np.sqrt(4 * np.pi)
else:
phase = ((1j) ** (m + np.abs(m)))
normalizer = np.sqrt(((2 * l + 1.) * factorial(l - np.abs(m)))
/
(4 * np.pi * factorial(l + np.abs(m))))
P = lpmv(np.abs(m), l, np.cos(theta))
e = np.exp(1j * m * phi)
return phase * normalizer * P * e
def _naive_rsh_ph(l, m, theta, phi):
if m == 0:
return np.sqrt((2 * l + 1.) / (4 * np.pi)) * lpmv(m, l, np.cos(theta))
elif m < 0:
return np.sqrt(((2 * l + 1.) * factorial(l + m)) /
(2 * np.pi * factorial(l - m))) * lpmv(-m, l, np.cos(theta)) * np.sin(-m * phi)
elif m > 0:
return np.sqrt(((2 * l + 1.) * factorial(l - m)) /
(2 * np.pi * factorial(l + m))) * lpmv(m, l, np.cos(theta)) * np.cos(m * phi)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment