def test_Wigner_D_element_symmetries(Rs, ell_max):
LMpM = sf.LMpM_range(0, ell_max)
# D_{mp,m}(R) = (-1)^{mp+m} \bar{D}_{-mp,-m}(R)
MpPM = np.array([
mp + m for ell in range(ell_max + 1) for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1)
])
LmMpmM = np.array([[ell, -mp, -m] for ell in range(ell_max + 1)
for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1)])
print()
for R in Rs:
print("\t", R)
assert np.allclose(sf.Wigner_D_element(R, LMpM), (-1.)**MpPM *
np.conjugate(sf.Wigner_D_element(R, LmMpmM)),
atol=ell_max**2 * precision_Wigner_D_element,
rtol=ell_max**2 * precision_Wigner_D_element)
# D is a unitary matrix, so its conjugate transpose is its
# inverse. D(R) should equal the matrix inverse of D(R^{-1}).
# So: D_{mp,m}(R) = \bar{D}_{m,mp}(\bar{R})
LMMp = np.array([[ell, m, mp] for ell in range(ell_max + 1)
for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1)])
for R in Rs:
print("\t", R)
assert np.allclose(sf.Wigner_D_element(R, LMpM),
np.conjugate(sf.Wigner_D_element(R.inverse(),
LMMp)),
atol=ell_max**4 * precision_Wigner_D_element,
rtol=ell_max**4 * precision_Wigner_D_element)
def test_Wigner_D_element_underflow(Rs, ell_max):
# NOTE: This is a delicate test, which depends on the result underflowing exactly when expected.
# In particular, it should underflow to 0.0 when |mp+m|>32, but should never undeflow to 0.0
# when |mp+m|<32. So it's not the end of the world if this test fails, but it does mean that
# the underflow properties have changed, so it might be worth a look.
assert sf.ell_max >= 15 # Test can't work if this has been set lower
eps = 1.e-10
ell_max = max(sf.ell_max, ell_max)
# Test |Ra|=1e-10
LMpM = np.array([[ell, mp, m] for ell in range(ell_max + 1)
for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1) if abs(mp + m) > 32])
R = np.quaternion(eps, 1, 0, 0).normalized()
assert np.all(sf.Wigner_D_element(R, LMpM) == 0j)
LMpM = np.array([[ell, mp, m] for ell in range(ell_max + 1)
for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1) if abs(mp + m) < 32])
R = np.quaternion(eps, 1, 0, 0).normalized()
assert np.all(sf.Wigner_D_element(R, LMpM) != 0j)
# Test |Rb|=1e-10
LMpM = np.array([[ell, mp, m] for ell in range(ell_max + 1)
for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1) if abs(m - mp) > 32])
R = np.quaternion(1, eps, 0, 0).normalized()
assert np.all(sf.Wigner_D_element(R, LMpM) == 0j)
LMpM = np.array([[ell, mp, m] for ell in range(ell_max + 1)
for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1) if abs(m - mp) < 32])
R = np.quaternion(1, eps, 0, 0).normalized()
assert np.all(sf.Wigner_D_element(R, LMpM) != 0j)
def test_Wigner_D_element_overflow(Rs, ell_max):
assert sf.ell_max >= 15 # Test can't work if this has been set lower
ell_max = max(sf.ell_max, ell_max)
LMpM = sf.LMpM_range(0, ell_max)
# Test |Ra|=1e-10
R = np.quaternion(1.e-10, 1, 0, 0).normalized()
assert np.all(np.isfinite(sf.Wigner_D_element(R, LMpM)))
# Test |Rb|=1e-10
R = np.quaternion(1, 1.e-10, 0, 0).normalized()
assert np.all(np.isfinite(sf.Wigner_D_element(R, LMpM)))
def test_Wigner_D_signatures(Rs):
"""There are two ways to call the WignerD function: with an array of Rs, or with an array of (ell,mp,m) values.
This test ensures that the results are the same in both cases."""
# from spherical_functions.WignerD import _Wigner_D_elements
ell_max = 6
ell_mp_m = sf.LMpM_range(0, ell_max)
Ds1 = np.zeros((Rs.size, ell_mp_m.shape[0]), dtype=np.complex)
Ds2 = np.zeros_like(Ds1)
for i, R in enumerate(Rs):
Ds1[i, :] = sf.Wigner_D_element(R, ell_mp_m)
for i, (ell, mp, m) in enumerate(ell_mp_m):
Ds2[:, i] = sf.Wigner_D_element(Rs, ell, mp, m)
assert np.allclose(Ds1, Ds2, rtol=3e-15, atol=3e-15)
def test_Wigner_D_element_values(special_angles, ell_max):
LMpM = sf.LMpM_range_half_integer(0, ell_max // 2)
# Compare with more explicit forms given in Euler angles
print("")
for alpha in special_angles:
print("\talpha={0}".format(
alpha)) # Need to show some progress to Travis
for beta in special_angles:
print("\t\tbeta={0}".format(beta))
for gamma in special_angles:
a = np.conjugate(
np.array([
slow_Wigner_D_element(alpha, beta, gamma, ell, mp, m)
for ell, mp, m in LMpM
]))
b = sf.Wigner_D_element(
quaternion.from_euler_angles(alpha, beta, gamma), LMpM)
# if not np.allclose(a, b,
# atol=ell_max ** 6 * precision_Wigner_D_element,
# rtol=ell_max ** 6 * precision_Wigner_D_element):
# for i in range(min(a.shape[0], 100)):
# print(LMpM[i], "\t", abs(a[i]-b[i]), "\t\t", a[i], "\t", b[i])
assert np.allclose(
a,
b,
atol=ell_max**6 * precision_Wigner_D_element,
rtol=ell_max**6 * precision_Wigner_D_element)
def swsh(self, s, l, m, theta, phi):
# get the rotor related for Wigner matrix
# Note: We need to specify the (-theta,-phi direction to match the convenction used in HAD)
R_tp = quaternion.from_spherical_coords(-theta, -phi)
W = sf.Wigner_D_element(R_tp, l, m, -s)
#print(W)
return ((-1)**(s)) * math.sqrt((2 * l + 1) / (4 * math.pi)) * W
def swsh(self,s,l,m,theta,phi):
# get the rotor related for Wigner matrix
R_tp = quaternion.from_spherical_coords(theta, phi)
W=sf.Wigner_D_element(R_tp,l,m,-s)
#print(W)
return ((-1)**s ) *math.sqrt((2*l+1)/(4*math.pi)) * W
def test_Wigner_D_matrix_inverse(Rs, ell_max):
# Ensure that the matrix of the inverse rotation is the inverse of
# the matrix of the rotation
print()
for i, R in enumerate(Rs):
print("\t{0} of {1}: {2}".format(i+1, len(Rs), R))
for twoell in range(2*ell_max + 1):
LMpM = np.array([[twoell/2, twomp/2, twom/2]
for twomp in range(-twoell, twoell+1, 2)
for twom in range(-twoell, twoell+1, 2)])
D1 = sf.Wigner_D_element(R.a, R.b, LMpM)
D2 = sf.Wigner_D_element(R.a.conjugate(), -R.b, LMpM)
D1 = D1.reshape((twoell + 1, twoell + 1))
D2 = D2.reshape((twoell + 1, twoell + 1))
assert np.allclose(D1.dot(D2), np.identity(twoell + 1),
atol=ell_max ** 4 * precision_Wigner_D_element,
rtol=ell_max ** 4 * precision_Wigner_D_element)
def rotate_tensor(tensor, angles, inverse=False):
""" Rotate a complex tensor.
Parameters
----------
tensor: dict
complex rank-2 tensor to rotate; the tensor is expected to be complete
i.e. no entries should be missing
angles: np.ndarray (3,)
euler angles: alpha, beta, gamma
inverse: bool,
{False: rotate vector, True: rotate CS}
Returns
---------
dict
Rotated version of tensor
Info
----
Remember that in nncs and elfcs alignment, inverse = True should be used
"""
if not isinstance(tensor['0,0,0'], np.complex128) and not isinstance(tensor['0,0,0'], np.complex64)\
and not type(tensor['0,0,0']) == complex:
raise Exception('tensor has to be complex')
R = {}
n_max, l_max = get_max(tensor)
for l in range(1, l_max):
# if not '0,{},0'.format(l) in tensor:
# break
R[l] = sf.Wigner_D_element(*angles,
np.array([l
])).reshape(2 * l + 1, 2 * l + 1)
if inverse:
R[l] = R[l].conj().T
# R[l] = R[l].conj()
tensor_rotated = {}
for n in range(n_max):
# if not '{},0,0'.format(n) in tensor:
# break
tensor_rotated['{},0,0'.format(n)] = tensor['{},0,0'.format(n)]
for l in range(1, l_max):
# if not '0,{},0'.format(l) in tensor:
# break
t = []
for m in range(-l, l + 1):
t.append(tensor['{},{},{}'.format(n, l, m)])
t = np.array(t)
t_rotated = R[l].dot(t)
for m in range(-l, l + 1):
tensor_rotated['{},{},{}'.format(n, l, m)] = t_rotated[l + m]
return tensor_rotated
def test_Wigner_D_input_types(Rs, special_angles, ell_max):
LMpM = sf.LMpM_range(0, ell_max // 2)
# Compare with more explicit forms given in Euler angles
print("")
for alpha in special_angles:
print("\talpha={0}".format(alpha)) # Need to show some progress to Travis
for beta in special_angles:
for gamma in special_angles:
a = sf.Wigner_D_element(alpha, beta, gamma, LMpM)
b = sf.Wigner_D_element(quaternion.from_euler_angles(alpha, beta, gamma), LMpM)
assert np.allclose(a, b,
atol=ell_max ** 6 * precision_Wigner_D_element,
rtol=ell_max ** 6 * precision_Wigner_D_element)
for R in Rs:
a = sf.Wigner_D_element(R, LMpM)
b = sf.Wigner_D_element(R.a, R.b, LMpM)
assert np.allclose(a, b,
atol=ell_max ** 6 * precision_Wigner_D_element,
rtol=ell_max ** 6 * precision_Wigner_D_element)
def test_SWSH_WignerD_expression(special_angles, ell_max):
for iota in special_angles:
for phi in special_angles:
for ell in range(ell_max + 1):
for s in range(-ell, ell + 1):
R = quaternion.from_euler_angles(phi, iota, 0)
LM = np.array([[ell, m] for m in range(-ell, ell + 1)])
Y = sf.SWSH(R, s, LM)
LMS = np.array([[ell, m, -s] for m in range(-ell, ell + 1)])
D = (-1.) ** (s) * math.sqrt((2 * ell + 1) / (4 * np.pi)) * sf.Wigner_D_element(R.a, R.b, LMS)
assert np.allclose(Y, D, atol=ell ** 6 * precision_SWSH, rtol=ell ** 6 * precision_SWSH)
def test_Wigner_D_element_roundoff(Rs, ell_max):
LMpM = sf.LMpM_range(0, ell_max)
# Test rotations with |Ra|<1e-15
expected = [((-1.) ** ell if mp == -m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in
range(-ell, ell + 1)]
assert np.allclose(sf.Wigner_D_element(quaternion.x, LMpM), expected,
atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element)
expected = [((-1.) ** (ell + m) if mp == -m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1)
for m in range(-ell, ell + 1)]
assert np.allclose(sf.Wigner_D_element(quaternion.y, LMpM), expected,
atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element)
for theta in np.linspace(0, 2 * np.pi):
expected = [((-1.) ** (ell + m) * (np.cos(theta) + 1j * np.sin(theta)) ** (2 * m) if mp == -m else 0.0)
for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)]
assert np.allclose(sf.Wigner_D_element(np.cos(theta) * quaternion.y + np.sin(theta) * quaternion.x, LMpM),
expected,
atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element)
# Test rotations with |Rb|<1e-15
expected = [(1.0 if mp == m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in
range(-ell, ell + 1)]
assert np.allclose(sf.Wigner_D_element(quaternion.one, LMpM), expected,
atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element)
expected = [((-1.) ** m if mp == m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in
range(-ell, ell + 1)]
assert np.allclose(sf.Wigner_D_element(quaternion.z, LMpM), expected,
atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element)
for theta in np.linspace(0, 2 * np.pi):
expected = [((np.cos(theta) + 1j * np.sin(theta)) ** (2 * m) if mp == m else 0.0)
for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)]
assert np.allclose(sf.Wigner_D_element(np.cos(theta) * quaternion.one + np.sin(theta) * quaternion.z, LMpM),
expected,
atol=ell_max * precision_Wigner_D_element, 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 rotate_vector(vec, angles, inverse = False):
""" Rotate a real vector (euclidean order: xyz) with euler angles
inverse = False: rotate vector
inverse = True: rotate CS"""
vec = vec[:,[1,2,0]]
T_inv = np.conj(T.T)
D = sf.Wigner_D_element(*angles,np.array([1])).reshape(3,3)
if inverse:
D = D.conj().T
# D = D.conj()
R = T.dot(D.dot(T_inv))
# assert np.allclose(R.conj(), R)
vec= np.einsum('ij,kj -> ki', R, vec)
return vec[:,[2,0,1]].real
def WignerD_mm(l, quat):
"""
Calculates Wigner D matrix (as an numpy (2*l+1,2*l+1)-shaped array)
for order l, and a rotation given by quaternion quat.
This represents the rotation of spherical vector basis
TODO doc
"""
if use_moble_quaternion:
indices = np.array([[l, i, j] for i in range(-l, l + 1)
for j in range(-l, l + 1)])
Delems = sf.Wigner_D_element(quat,
indices).reshape(2 * l + 1, 2 * l + 1)
return Delems
else:
Delems = np.zeros((2 * l + 1, 2 * l + 1), dtype=complex)
for i in range(-l, l + 1):
for j in range(-l, l + 1):
Delems[i, j] = quat.wignerDelem(l, i, j)
return Delems
def rotate_vector(vec, angles, inverse=False):
""" Rotate a real vector (euclidean order: xyz) with euler angles
Parameters
--------
vec: np.ndarray (?, 3)
vector(s) to rotate, note that if more than one vector provided,
ever vector is rotated by the same angles.
angles: np.ndarray (3)
euler angles: alpha, beta, gamma.
inverse: bool
{False: rotate vector, True: rotate coordinate system}
Returns
-------
np.ndarray
rotated vector(s)
"""
if vec.ndim == 1 and len(vec) == 3:
vec = vec.reshape(1, 3)
vec = vec[:, [1, 2, 0]]
T_inv = np.conj(T.T)
D = sf.Wigner_D_element(*angles, np.array([1])).reshape(3, 3)
if inverse:
D = D.conj().T
# D = D.conj()
R = T.dot(D.dot(T_inv))
# assert np.allclose(R.conj(), R)
vec = np.einsum('ij,kj -> ki', R, vec)
return vec[:, [2, 0, 1]].real
def Wigner_D_element_Cached(pRot, P, Rp, R):
"""Wrapper for WignerD caching with functools.lru_cache"""
return sf.Wigner_D_element(pRot, P, Rp, R)
def wDcalc(Lrange=[0, 1], Nangs=None, eAngs=None, R=None, XFlag=True):
'''
Calculate set of Wigner D functions D(l,m,mp,R) on a grid.
Parameters
----------
Lrange : list, optional, default [0, 1]
Range of L to calculate parameters for.
If len(Lrange) == 2 assumed to be of form [Lmin, Lmax], otherwise list is used directly.
For a given l, all (m, mp) combinations are calculated.
Options for setting angles (use one only):
Nangs : int, optional, default None
If passed, use this to define Euler angles sampled.
Ranges will be set as (theta, phi, chi) = (0:pi, 0:pi/2, 0:pi) in Nangs steps.
eAngs : np.array, optional, default None
If passed, use this to define Euler angles sampled.
Array of angles, [theta,phi,chi], in radians
R : np.array, optional, default None
If passed, use this to define Euler angles sampled.
Array of quaternions, as given by quaternion.from_euler_angles(eAngs).
XFlag : bool, optional, default True
Flag for output. If true, output is Xarray. If false, np.arrays
Outputs
-------
- if XFlag -
wDX
Xarray, dims (lmmp,Euler)
- else -
wD, R, lmmp
np.arrays of values, dims (lmmp,Euler), plus list of angles and lmmp sets.
Methods
-------
Uses Moble's spherical_functions package for wigner D function.
https://github.com/moble/spherical_functions
Moble's quaternion package for angles and conversions.
https://github.com/moble/quaternion
Examples
--------
>>> wDX1 = wDcalc(eAngs = np.array([0,0,0]))
>>> wDX2 = wDcalc(Nangs = 10)
'''
# Set QNs for calculation, (l,m,mp)
if len(Lrange) == 2:
Ls = np.arange(Lrange[0], Lrange[1] + 1)
else:
Ls = Lrange
QNs = []
for l in Ls:
for m in np.arange(-l, l + 1):
for mp in np.arange(-l, l + 1):
QNs.append([l, m, mp])
QNs = np.array(QNs)
# Set angles - either input as a range, a set or as quaternions
if Nangs is not None:
# Set a range of Eugler angles for testing
pRot = np.linspace(0, np.pi, Nangs)
tRot = np.linspace(0, np.pi / 2, Nangs)
cRot = np.linspace(0, np.pi, Nangs)
eAngs = np.array([
pRot,
tRot,
cRot,
]).T
if eAngs is not None:
if eAngs.shape[
-1] != 3: # Check dims, should be (N X 3) for quaternion... but transpose for pd.MultiIndex
eAngs = eAngs.T
if R is None:
# Convert to quaternions
R = quaternion.from_euler_angles(eAngs)
# Calculate WignerDs
# sf.Wigner_D_element is vectorised for QN OR angles
# Here loop over QNs for a set of angles R
wD = []
lmmp = []
for n in np.arange(0, QNs.shape[0]):
lmmp.append(QNs[n, :])
wD.append(sf.Wigner_D_element(R, QNs[n, 0], QNs[n, 1], QNs[n, 2]))
# Return values as Xarray or np.arrays
if XFlag:
# Put into Xarray
#TODO: this will currently fail for a single set of QNs.
QNs = pd.MultiIndex.from_arrays(np.asarray(lmmp).T,
names=['lp', 'mu', 'mu0'])
if eAngs.size == 3: # Ugh, special case for only one set of angles.
Euler = pd.MultiIndex.from_arrays(
[[eAngs[0]], [eAngs[1]], [eAngs[2]]], names=['P', 'T', 'C'])
wDX = xr.DataArray(np.asarray(wD), coords=[('QN', QNs)])
wDX = wDX.expand_dims({'Euler': Euler})
else:
Euler = pd.MultiIndex.from_arrays(eAngs.T, names=['P', 'T', 'C'])
wDX = xr.DataArray(np.asarray(wD),
coords=[('QN', QNs), ('Euler', Euler)])
return wDX
else:
return wD, R, np.asarray(lmmp).T
for m in range(-l, l + 1)])
nw = len(indices)
aves = np.zeros(nw, dtype='complex128')
print("\033[1;36mIncluding %d Wigner coefficients\033[0m" % nw)
# Average order parameters
for i in range(n_alpha):
for j in range(n_theta):
for k in range(n_phi):
alpha = alpha_grid[i]
theta = theta_grid[j]
phi = phi_grid[k]
Ds = sf.Wigner_D_element(phi, theta, alpha, indices)
aves += np.conj(Ds) * psi[i, j, k] * np.sin(
theta) * d_theta * d_phi * d_alpha
print("\033[1;34mAveraged over %d out of %d iso-alpha surfaces\033[0m" %
(i + 1, n_alpha))
print("\033[1;32mCompleted averaging process\033[0m")
# Wigner decomposition
for i in range(n_alpha):
for j in range(n_theta):
for k in range(n_phi):
alpha = alpha_grid[i]
theta = theta_grid[j]
phi = phi_grid[k]
sys.path.insert(0, '../mlgw_v2')
import GW_generator as gen
g = gen.GW_generator(1)
for i in range(10000):
alpha, beta, gamma = np.array([np.random.uniform(0, 2 * np.pi)]), np.array(
[np.random.uniform(0, 2 * np.pi)
]), np.array([np.random.uniform(0, 2 * np.pi)])
ell = np.random.randint(5) + 2
mp, m = [np.random.randint(ell + 1)], [np.random.randint(ell + 1)]
mp[0] *= (np.random.randint(2) * 2 - 1)
m[0] *= (np.random.randint(2) * 2 - 1)
wig_package = sf.Wigner_D_element(alpha[0], beta[0], gamma[0], ell, mp[0],
m[0])
wig_mine = g._GW_generator__get_Wigner_D_matrix(ell, mp, m, alpha, beta,
gamma)
all_close = np.allclose(wig_package, np.conj(wig_mine), rtol=1e-5)
print("Iteration: ", i)
print("l,m,mp ", ell, mp, m)
print("\t", all_close, wig_mine, wig_package)
assert all_close
#print("package wigner ", wig_package)
#print("my wigner ", wig_mine)
def mfpad(dataIn, thres = 1e-2, inds = {'Type':'L','it':1}, res = 50, R = None, p = 0):
"""
Parameters
----------
dataIn : Xarray
Contains set(s) of matrix elements to use, as output by epsproc.readMatEle().
thres : float, optional, default 1e-2
Threshold value for matrix elements to use in calculation.
ind : dictionary, optional.
Used for sub-selection of matrix elements from Xarrays.
Default set for length gauage, single it component only, inds = {'Type':'L','it':'1'}.
res : int, optional, default 50
Resolution for output (theta,phi) grids.
R : list of Euler angles or quaternions, optional.
Define LF > MF polarization geometry/rotations.
For default case (R = None), 3 geometries are calculated, corresponding to z-pol, x-pol and y-pol cases.
Defined by Euler angles (p,t,c) = [0 0 0] for z-pol, [0 pi/2 0] for x-pol, [pi/2 pi/2 0] for y-pol.
p : int, optional.
Defines LF polarization state, p = -1...1, default p = 0 (linearly pol light along z-axis).
TODO: add summation over p for multiple pol states in LF.
Returns
-------
Ta
Xarray (theta, phi, E, Sym) of MFPADs, summed over (l,m)
Tlm
Xarray (theta, phi, E, Sym, lm) of MFPAD components, expanded over all (l,m)
"""
# Define reduced data from selection over all data
daRed = matEleSelector(dataIn, thres = 1e-2, inds = inds)
# Generate spherical harmonics
Lmax = daRed.l.max()
YlmX = sphCalc(Lmax, res = res)
# Reindex to match data (should happen automagically, but not always!)
# YlmXre = YlmX.reindex_like(daRed)
# Set rotation angles for LF > MF
if R is None:
# Set (x,y,z) projection terms only
# Nangs = 10
# pRot = np.linspace(0,180,Nangs)
# tRot = np.linspace(0,90,Nangs)
# cRot = np.linspace(0,180,Nangs)
# eAngs = np.array([pRot, tRot, cRot,])*np.pi/180
# Convert to quaternions
# R = quaternion.from_euler_angles(pRot*np.pi/180, tRot*np.pi/180, cRot*np.pi/180)
# Eugler angles for rotation of LF->MF, set as [0 0 0] for z-pol, [0 pi/2 0] for x-pol, [pi/2 pi/2 0] for y-pol
pRot = [0, 0, np.pi/2]
tRot = [0, np.pi/2, np.pi/2]
cRot = [0, 0, 0]
eAngs = np.array([pRot, tRot, cRot]) # List form to use later
Euler = pd.MultiIndex.from_arrays(eAngs, names = ['P','T','C'])
# Convert to quaternions
R = quaternion.from_euler_angles(pRot, tRot, cRot)
#**************** Calculate MFPADs
Tlm = []
Ta = []
# Loop over pol geoms R
for n, Rcalc in enumerate(R):
T = []
# Loop over mu terms and multiply
for mu in np.arange(-1,2):
# Set by element replacement (preserves whole structure)
# daTemp = daRed.copy() # Set explicit copy for rotation.
# daTemp.loc[{'mu':mu}].values = daTemp.loc[{'mu':mu}].values * sf.Wigner_D_element(Rcalc, 1, mu, 0).conj()
# Issues with reindexing to extra coords at the moment, so reindex and multiply for specific mu only
# daTemp = daTemp.sel({'mu':mu})
# YlmXre = YlmX.reindex_like(daTemp)
# T.append(YlmXre.conj() * daTemp) # Output full (l,m,mu) expansion
# Set by looping and selection
daTemp = daRed.sel({'mu':mu}) * sf.Wigner_D_element(Rcalc, 1, mu, 0).conj()
YlmXre = YlmX.reindex_like(daTemp)
T.append(YlmXre.conj() * daTemp) # Output full (l,m,mu) expansion
# Concat & sum over symmetries
Ts = xr.combine_nested([T[0], T[1], T[2]], concat_dim=['LM'])
# Add dims - currently set for Euler angles only.
# Can't seem to add mutiindex as a single element, so set dummy coord here and replace below.
Ts = Ts.expand_dims({'Euler':[n]}) # Set as index
# Ts = Ts.expand_dims({'p':[eAngs[0,n]], 't':[eAngs[1,n]], 'c':[eAngs[2,n]]})
Tlm.append(Ts)
Ta.append(Ts.sum(dim = 'LM'))
TlmX = xr.combine_nested(Tlm, concat_dim=['Euler'])
TaX = xr.combine_nested(Ta, concat_dim=['Euler'])
# Assign Euler angles to dummy dim
TlmX = TlmX.assign_coords(Euler = Euler)
TaX = TaX.assign_coords(Euler = Euler)
return TaX, TlmX # , Ta, Tlm # For debug also return lists
