def _rotate_decomposition_basis_by_series(data, R_basis, ell_min, ell_max, D):
"""Rotate data by a different rotor at each point in time
`D` is just a workspace, which holds the Wigner D matrices.
During the summation, it is also used as temporary storage to hold
the results for each item of data, where the first row in the
matrix is overwritten with the new sums.
"""
for i_t in xrange(data.shape[0]):
sf._Wigner_D_matrices(R_basis[i_t, 0], R_basis[i_t, 1], ell_min,
ell_max, D)
for ell in xrange(ell_min, ell_max + 1):
i_data = ell**2 - ell_min**2
i_D = sf._linear_matrix_offset(ell, ell_min)
for i_m in xrange(2 * ell + 1):
new_data_mp = 0j
for i_mp in xrange(2 * ell + 1):
new_data_mp += data[i_t, i_data +
i_mp] * D[i_D + i_m +
(2 * ell + 1) * i_mp]
D[i_D + i_m] = new_data_mp
for i_m in xrange(2 * ell + 1):
data[i_t, i_data + i_m] = D[i_D + i_m]
def rotate_decomposition_basis(W, R_basis):
"""Rotate a Waveform in place
This function takes a Waveform object `W` and either a quaternion
or array of quaternions `R_basis`. It applies that rotation to
the decomposition basis of the modes in the Waveform. The change
in basis is also recorded in the Waveform's `frame` data.
For more information on the analytical details, see
http://moble.github.io/spherical_functions/SWSHs.html#rotating-swshs
"""
# This will be used in the jitted functions below to store the
# Wigner D matrix at each time step
D = np.empty((sf.WignerD._total_size_D_matrices(W.ell_min, W.ell_max),), dtype=complex)
if isinstance(R_basis, (list, np.ndarray)) and len(R_basis) == 1:
R_basis = R_basis[0]
if isinstance(R_basis, (list, np.ndarray)):
if isinstance(R_basis, np.ndarray) and R_basis.ndim != 1:
raise ValueError("Input dimension mismatch. R_basis.shape={1}".format(R_basis.shape))
if W.n_times != len(R_basis):
raise ValueError(
"Input dimension mismatch. (W.n_times={}) != (len(R_basis)={})".format(W.n_times, len(R_basis))
)
_rotate_decomposition_basis_by_series(W.data, quaternion.as_spinor_array(R_basis), W.ell_min, W.ell_max, D)
# Update the frame data, using right-multiplication
if W.frame.size:
if W.frame.shape[0] == 1:
# Numpy can't currently multiply one element times an array
W.frame = np.array([W.frame * R for R in R_basis])
else:
W.frame = W.frame * R_basis
else:
W.frame = np.copy(R_basis)
# We can't just use an `else` here because we need to process the
# case where the input was an iterable of length 1, which we've
# now changed to just a single quaternion.
if isinstance(R_basis, np.quaternion):
sf._Wigner_D_matrices(R_basis.a, R_basis.b, W.ell_min, W.ell_max, D)
tmp = np.empty((2 * W.ell_max + 1,), dtype=complex)
_rotate_decomposition_basis_by_constant(W.data, W.ell_min, W.ell_max, D, tmp)
# Update the frame data, using right-multiplication
if W.frame.size:
W.frame = W.frame * R_basis
else:
W.frame = np.array([R_basis])
opts = np.get_printoptions()
np.set_printoptions(threshold=6)
W.__history_depth__ -= 1
W._append_history(f"{W}.rotate_decomposition_basis({R_basis})")
np.set_printoptions(**opts)
return W
def test_Wigner_D_matrices_negative_argument(Rs, ell_max):
# For integer ell, D(R)=D(-R)
LMpM = sf.LMpM_range(0, ell_max)
a = np.empty((LMpM.shape[0], ), dtype=complex)
b = np.empty((LMpM.shape[0], ), dtype=complex)
for R in Rs:
sf._Wigner_D_matrices(R.a, R.b, 0, ell_max, a)
sf._Wigner_D_matrices(-R.a, -R.b, 0, ell_max, b)
assert np.allclose(a, b, rtol=ell_max * precision_Wigner_D_element)
def test_Wigner_D_matrices_negative_argument(Rs, ell_max):
# For integer ell, D(R)=D(-R)
LMpM = sf.LMpM_range(0, ell_max)
a = np.empty((LMpM.shape[0],), dtype=complex)
b = np.empty((LMpM.shape[0],), dtype=complex)
for R in Rs:
sf._Wigner_D_matrices(R.a, R.b, 0, ell_max, a)
sf._Wigner_D_matrices(-R.a, -R.b, 0, ell_max, b)
assert np.allclose(a, b, rtol=ell_max * precision_Wigner_D_element)
def test_Wigner_D_matrix(Rs, ell_max):
for l_min in [0, 1, 2, ell_max // 2, ell_max - 1]:
print("")
for l_max in range(l_min + 1, ell_max + 1):
print("\tWorking on (l_min,l_max)=({0},{1})".format(l_min, l_max))
LMpM = sf.LMpM_range(l_min, l_max)
for R in Rs:
elements = sf.Wigner_D_element(R, LMpM)
matrix = np.empty(LMpM.shape[0], dtype=complex)
sf._Wigner_D_matrices(R.a, R.b, l_min, l_max, matrix)
assert np.allclose(elements, matrix,
atol=1e3 * l_max * ell_max * precision_Wigner_D_element,
rtol=1e3 * l_max * ell_max * precision_Wigner_D_element)
def test_Wigner_D_matrix(Rs, ell_max):
for l_min in [0, 1, 2, ell_max // 2, ell_max - 1]:
print("")
for l_max in range(l_min + 1, ell_max + 1):
print("\tWorking on (l_min,l_max)=({0},{1})".format(l_min, l_max))
LMpM = sf.LMpM_range(l_min, l_max)
for R in Rs:
elements = sf.Wigner_D_element(R, LMpM)
matrix = np.empty(LMpM.shape[0], dtype=complex)
sf._Wigner_D_matrices(R.a, R.b, l_min, l_max, matrix)
assert np.allclose(elements, matrix,
atol=1e3 * l_max * ell_max * precision_Wigner_D_element,
rtol=1e3 * l_max * ell_max * precision_Wigner_D_element)
def rotate_decomposition_basis(self, R_basis):
"""Rotate a Waveform (not in place).
This function takes a Waveform object `self` and either a quaternion
or array of quaternions `R_basis`. It applies that rotation to
the decomposition basis of the modes in the Waveform. The change
in basis is also recorded in the Waveform's `frame` data.
For more information on the analytical details, see
http://moble.github.io/spherical_functions/SWSHs.html#rotating-swshs
"""
# This will be used in the jitted functions below to store the
# Wigner D matrix at each time step
D = np.empty((spherical_functions.WignerD._total_size_D_matrices(
self.ell_min, self.ell_max), ),
dtype=complex)
rotated_data = self.copy()
if isinstance(R_basis, (list, np.ndarray)) and len(R_basis) == 1:
R_basis = R_basis[0]
if isinstance(R_basis, (list, np.ndarray)):
if isinstance(R_basis, np.ndarray) and R_basis.ndim != 1:
raise ValueError(
f"Input dimension mismatch. R_basis.shape={R_basis.shape}"
)
if self.shape[0] != len(R_basis):
raise ValueError(
"Input dimension mismatch. (self.shape[0]={self.shape[0]}) != (len(R_basis)={len(R_basis)})"
)
_rotate_decomposition_basis_by_series(
rotated_data, quaternion.as_spinor_array(R_basis),
self.ell_min, self.ell_max, D)
# We can't just use an `else` here because we need to process the
# case where the input was an iterable of length 1, which we've
# now changed to just a single quaternion.
if isinstance(R_basis, np.quaternion):
spherical_functions._Wigner_D_matrices(R_basis.a, R_basis.b,
self.ell_min, self.ell_max,
D)
tmp = np.empty((2 * self.ell_max + 1, ), dtype=complex)
_rotate_decomposition_basis_by_constant(rotated_data, self.ell_min,
self.ell_max, D, tmp)
return rotated_data
def test_Wigner_D_matrices_representation_property(Rs, ell_max):
# Test the representation property for special and random angles
# Can't try half-integers because Wigner_D_matrices doesn't accept them
ell_max = min(8, ell_max)
LMpM = sf.LMpM_range(0, ell_max)
print("")
D1 = np.empty((LMpM.shape[0], ), dtype=complex)
D2 = np.empty((LMpM.shape[0], ), dtype=complex)
D12 = np.empty((LMpM.shape[0], ), dtype=complex)
for i, R1 in enumerate(Rs):
print("\t{0} of {1}: R1 = {2}".format(i + 1, len(Rs), R1))
for j, R2 in enumerate(Rs[i:]):
# print("\t\t{0} of {1}: R2 = {2}".format(j+1, len(Rs), R2))
R12 = R1 * R2
sf._Wigner_D_matrices(R1.a, R1.b, 0, ell_max, D1)
sf._Wigner_D_matrices(R2.a, R2.b, 0, ell_max, D2)
sf._Wigner_D_matrices(R12.a, R12.b, 0, ell_max, D12)
M12 = np.array([
np.sum([
D1[sf.LMpM_index(ell, mp, mpp, 0)] *
D2[sf.LMpM_index(ell, mpp, m, 0)]
for mpp in range(-ell, ell + 1)
]) for ell in range(ell_max + 1)
for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)
])
assert np.allclose(M12,
D12,
atol=ell_max * precision_Wigner_D_element)
def test_Wigner_D_matrices_representation_property(Rs, ell_max):
# Test the representation property for special and random angles
# Can't try half-integers because Wigner_D_matrices doesn't accept them
ell_max = min(8, ell_max)
LMpM = sf.LMpM_range(0, ell_max)
print("")
D1 = np.empty((LMpM.shape[0],), dtype=complex)
D2 = np.empty((LMpM.shape[0],), dtype=complex)
D12 = np.empty((LMpM.shape[0],), dtype=complex)
for i, R1 in enumerate(Rs):
print("\t{0} of {1}: R1 = {2}".format(i+1, len(Rs), R1))
for j, R2 in enumerate(Rs[i:]):
# print("\t\t{0} of {1}: R2 = {2}".format(j+1, len(Rs), R2))
R12 = R1 * R2
sf._Wigner_D_matrices(R1.a, R1.b, 0, ell_max, D1)
sf._Wigner_D_matrices(R2.a, R2.b, 0, ell_max, D2)
sf._Wigner_D_matrices(R12.a, R12.b, 0, ell_max, D12)
M12 = np.array([np.sum([D1[sf.LMpM_index(ell, mp, mpp, 0)] * D2[sf.LMpM_index(ell, mpp, m, 0)]
for mpp in range(-ell, ell + 1)])
for ell in range(ell_max + 1)
for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1)])
assert np.allclose(M12, D12, atol=ell_max * precision_Wigner_D_element)