def ethbar_NP(modes, spin_weight, ell_min=0):
"""Spin-lowering \bar{eth} operator as defined by Newman and Penrose
N.B.: For our purposes, eth_GHP is the same as eth_NP/sqrt(2).
This is the complex conjugate of the `eth_NP` operator. See that function's docstring for more information.
Parameters
----------
modes : 1-d complex array
This array contains the modes starting from ell=ell_min, and continuing in standard order (as in sf.LM_range).
spin_weight : int
Spin weight of the input field. The \bar{eth} operators lower the spin weight by 1.
ell_min : int, optional
Smallest ell value present in input `modes`. Defaults to 0.
Returns
-------
1-d complex array
The output has the same size as the input `modes`, and corresponds to the same (ell,m) values. Note, however,
that these modes have spin weight less than the input by 1.
"""
ell_max = int(sqrt(len(modes) + LM_total_size(0, ell_min - 1))) - 1
ethbar_modes = np.array(modes)
i_mode = 0
for ell in xrange(ell_min, ell_max + 1):
factor = (0.0 if ell < abs(spin_weight - 1) else -sqrt((ell + spin_weight) * (ell - spin_weight + 1.)))
for m in xrange(-ell, ell + 1):
ethbar_modes[i_mode] *= factor
i_mode += 1
return ethbar_modes
def _LVector(data1, data2, lm, Lvec):
"""Helper function for the LVector function"""
# Big, bad, ugly, obvious way to do the calculation
# =================================================
# L+ = Lx + i Ly Lx = (L+ + L-) / 2
# L- = Lx - i Ly Ly = -i (L+ - L-) / 2
for i_mode in xrange(lm.shape[0]):
L = lm[i_mode, 0]
M = lm[i_mode, 1]
for i_time in xrange(data1.shape[0]):
# Compute first in (+,-,z) basis
Lp = (np.conjugate(data1[i_time, i_mode + 1]) *
data2[i_time, i_mode] * ladder(L, M) if M + 1 <= L else 0.0 +
0.0j)
Lm = (np.conjugate(data1[i_time, i_mode - 1]) *
data2[i_time, i_mode] *
ladder(L, -M) if M - 1 >= -L else 0.0 + 0.0j)
Lz = np.conjugate(data1[i_time, i_mode]) * data2[i_time,
i_mode] * M
# Convert into (x,y,z) basis
Lvec[i_time, 0] += 0.5 * (Lp + Lm)
Lvec[i_time, 1] += -0.5j * (Lp - Lm)
Lvec[i_time, 2] += Lz
return
def _LM_range(ell_min, ell_max, LM):
i = 0
for ell in xrange(ell_min, ell_max + 1):
for m in xrange(-ell, ell + 1):
LM[i, 0] = ell
LM[i, 1] = m
i += 1
def fd_indefinite_integral(f, t):
Sfdt = np.empty_like(f)
Sfdt[0] = 0.0
for i in xrange(1, len(t)):
for j in xrange(f.shape[1]):
Sfdt[i, j] = Sfdt[i - 1, j] + (f[i, j] + f[i - 1, j]) * ((t[i] - t[i - 1]) / 2.0)
return Sfdt
def complex_array_abs(c, s):
for i in xrange(len(s)):
s[i] = 0.0
for j in xrange(c.shape[1]):
s[i] += c[i, j].real ** 2 + c[i, j].imag ** 2
s[i] = np.sqrt(s[i])
return
def _LdtVector(data, datadot, lm, Ldt):
"""Helper function for the LdtVector function"""
# Big, bad, ugly, obvious way to do the calculation
# =================================================
# L+ = Lx + i Ly Lx = (L+ + L-) / 2
# L- = Lx - i Ly Ly = -i (L+ - L-) / 2
for i_mode in xrange(lm.shape[0]):
L = lm[i_mode, 0]
M = lm[i_mode, 1]
for i_time in xrange(data.shape[0]):
# Compute first in (+,-,z) basis
Lp = (np.conjugate(data[i_time, i_mode + 1]) *
datadot[i_time, i_mode] *
ladder(L, M) if M + 1 <= L else 0.0 + 0.0j)
Lm = (np.conjugate(data[i_time, i_mode - 1]) *
datadot[i_time, i_mode] *
ladder(L, -M) if M - 1 >= -L else 0.0 + 0.0j)
Lz = np.conjugate(data[i_time, i_mode]) * datadot[i_time,
i_mode] * M
# Convert into (x,y,z) basis
Ldt[i_time, 0] += 0.5 * (Lp.imag + Lm.imag)
Ldt[i_time, 1] += -0.5 * (Lp.real - Lm.real)
Ldt[i_time, 2] += Lz.imag
return
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 _LM_range(ell_min, ell_max, LM):
i = 0
for ell in xrange(ell_min, ell_max + 1):
for m in xrange(-ell, ell + 1):
LM[i, 0] = ell
LM[i, 1] = m
i += 1
def indefinite_integral(f, t):
Sfdt = np.empty_like(f)
Sfdt[0] = 0.0
for i in xrange(1, len(t)):
for j in xrange(f.shape[1]):
Sfdt[i, j] = Sfdt[i - 1, j] + (f[i, j] + f[i - 1, j]) * ((t[i] - t[i - 1]) / 2.0)
return Sfdt
def definite_integral(f, t):
Sfdt = np.empty_like(f)
for i in xrange(len(Sfdt)):
Sfdt[i] = 0.0
for i in xrange(1, f.shape[0]):
for j in xrange(f.shape[1]):
Sfdt[j] += (f[i, j] + f[i - 1, j]) * ((t[i] - t[i - 1]) / 2.0)
return Sfdt
def _LMpM_range_half_integer(twoell_min, twoell_max, LMpM):
i = 0
for twoell in xrange(twoell_min, twoell_max + 1):
for twomp in xrange(-twoell, twoell + 1, 2):
for twom in xrange(-twoell, twoell + 1, 2):
LMpM[i, 0] = twoell / 2
LMpM[i, 1] = twomp / 2
LMpM[i, 2] = twom / 2
i += 1
def _LMpM_range_half_integer(twoell_min, twoell_max, LMpM):
i = 0
for twoell in xrange(twoell_min, twoell_max + 1):
for twomp in xrange(-twoell, twoell + 1, 2):
for twom in xrange(-twoell, twoell + 1, 2):
LMpM[i, 0] = twoell / 2
LMpM[i, 1] = twomp / 2
LMpM[i, 2] = twom / 2
i += 1
def _LLDominantEigenvector(dpa, dpa_i):
"""Jitted helper function for LLDominantEigenvector"""
# Make the initial direction closer to RoughInitialEllDirection than not
if (dpa_i[0] * dpa[0, 0] + dpa_i[1] * dpa[0, 1] +
dpa_i[2] * dpa[0, 2]) < 0.:
dpa[0, 0] *= -1
dpa[0, 1] *= -1
dpa[0, 2] *= -1
# Now, go through and make the vectors reasonably continuous.
if dpa.shape[0] > 0:
LastNorm = sqrt(dpa[0, 0]**2 + dpa[0, 1]**2 + dpa[0, 2]**2)
for i in xrange(1, dpa.shape[0]):
Norm = dpa[i, 0]**2 + dpa[i, 1]**2 + dpa[i, 2]**2
dNorm = (dpa[i, 0] - dpa[i - 1, 0])**2 + (
dpa[i, 1] - dpa[i - 1, 1])**2 + (dpa[i, 2] - dpa[i - 1, 2])**2
if dNorm > Norm:
dpa[i, 0] *= -1
dpa[i, 1] *= -1
dpa[i, 2] *= -1
# While we're here, let's just normalize that last one
if LastNorm != 0.0 and LastNorm != 1.0:
dpa[i - 1, 0] /= LastNorm
dpa[i - 1, 1] /= LastNorm
dpa[i - 1, 2] /= LastNorm
LastNorm = sqrt(Norm)
if LastNorm != 0.0 and LastNorm != 1.0:
i = dpa.shape[0] - 1
dpa[i, 0] /= LastNorm
dpa[i, 1] /= LastNorm
dpa[i, 2] /= LastNorm
return
def Wigner3j(j_1, j_2, j_3, m_1, m_2, m_3):
"""Calculate the Wigner 3j symbol `Wigner3j(j_1,j_2,j_3,m_1,m_2,m_3)`
This function is copied with minor modification from
sympy.physics.Wigner, as written by Jens Rasch.
The inputs must be integers. (Half integer arguments are
sacrificed so that we can use numba.) Nonzero return quantities
only occur when the `j`s obey the triangle inequality (any two
must add up to be as big as or bigger than the third).
Examples
========
>>> from spherical_functions import Wigner3j
>>> Wigner3j(2, 6, 4, 0, 0, 0)
0.186989398002
>>> Wigner3j(2, 6, 4, 0, 0, 1)
0
"""
if (m_1 + m_2 + m_3 != 0):
return 0
if ((abs(m_1) > j_1) or (abs(m_2) > j_2) or (abs(m_3) > j_3)):
return 0
prefid = (1 if (j_1 - j_2 - m_3) % 2 == 0 else -1)
m_3 = -m_3
a1 = j_1 + j_2 - j_3
if (a1 < 0):
return 0
a2 = j_1 - j_2 + j_3
if (a2 < 0):
return 0
a3 = -j_1 + j_2 + j_3
if (a3 < 0):
return 0
argsqrt = (factorials[j_1 + j_2 - j_3] * factorials[j_1 - j_2 + j_3] *
factorials[-j_1 + j_2 + j_3] * factorials[j_1 - m_1] *
factorials[j_1 + m_1] * factorials[j_2 - m_2] *
factorials[j_2 + m_2] * factorials[j_3 - m_3] *
factorials[j_3 + m_3]) / factorials[j_1 + j_2 + j_3 + 1]
ressqrt = sqrt(argsqrt)
imin = max(-j_3 + j_1 + m_2, max(-j_3 + j_2 - m_1, 0))
imax = min(j_2 + m_2, min(j_1 - m_1, j_1 + j_2 - j_3))
sumres = 0.0
for ii in xrange(imin, imax + 1):
den = (factorials[ii] * factorials[ii + j_3 - j_1 - m_2] *
factorials[j_2 + m_2 - ii] * factorials[j_1 - ii - m_1] *
factorials[ii + j_3 - j_2 + m_1] *
factorials[j_1 + j_2 - j_3 - ii])
if (ii % 2 == 0):
sumres = sumres + 1.0 / den
else:
sumres = sumres - 1.0 / den
return ressqrt * sumres * prefid
def index_is_monotonic(y):
length = y.size
monotonic = np.ones_like(y, dtype=np.bool_)
direction = y[-1] - y[0]
if direction > 0.0:
max_value = y[0]
for i in xrange(1, length):
if y[i] <= max_value:
monotonic[i] = False
else:
max_value = y[i]
else:
min_value = y[0]
for i in xrange(1, length):
if y[i] >= min_value:
monotonic[i] = False
else:
min_value = y[i]
return monotonic
def eth_GHP(modes, spin_weight, ell_min=0):
"""Spin-raising eth operator as defined by Geroch-Held-Penrose
N.B.: For our purposes, eth_GHP is the same as eth_NP/sqrt(2).
This operator is defined in J. Math. Phys. 14, 874 (1973) <http://link.aip.org/link/?JMP/14/874/1>. The eth
operator was originally defined by Newman and Penrose in J. Math. Phys. 7, 863 (1966)
<http://link.aip.org/link/?JMP/7/863/1>, and discussed further by Goldberg et al., J. Math. Phys. 8, 2155 (1967)
<http://link.aip.org/link/?JMP/8/2155/1>. In the case of a spacelike 2-surface (which is all we treat in this
function), the GHP eth operator reduces almost to the original. By comparing GHP's Eq. (3.8) to Eq. (2.13) of
Goldberg et al., we can see the difference. As GHP say below their Eq. (3.11) that "for complete agreement in
the case of a general 2-surface metric we would have (strictly speaking) to multiply up our eth operator by a
factor sqrt(2)."
Parameters
----------
modes : 1-d complex array
This array contains the modes starting from ell=ell_min, and continuing in standard order (as in sf.LM_range).
spin_weight : int
Spin weight of the input field. The eth operators raise the spin weight by 1.
ell_min : int, optional
Smallest ell value present in input `modes`. Defaults to 0.
Returns
-------
1-d complex array
The output has the same size as the input `modes`, and corresponds to the same (ell,m) values. Note, however,
that these modes have spin weight greater than the input by 1.
"""
ell_max = int(sqrt(len(modes) + LM_total_size(0, ell_min - 1))) - 1
eth_modes = np.array(modes)
i_mode = 0
for ell in xrange(ell_min, ell_max + 1):
factor = (0.0 if ell < abs(spin_weight + 1) else sqrt(
(ell - spin_weight) * (ell + spin_weight + 1.) / 2.))
for m in xrange(-ell, ell + 1):
eth_modes[i_mode] *= factor
i_mode += 1
return eth_modes
def eth_GHP(modes, spin_weight, ell_min=0):
"""Spin-raising eth operator as defined by Geroch-Held-Penrose
N.B.: For our purposes, eth_GHP is the same as eth_NP/sqrt(2).
This operator is defined in J. Math. Phys. 14, 874 (1973) <http://link.aip.org/link/?JMP/14/874/1>. The eth
operator was originally defined by Newman and Penrose in J. Math. Phys. 7, 863 (1966)
<http://link.aip.org/link/?JMP/7/863/1>, and discussed further by Goldberg et al., J. Math. Phys. 8, 2155 (1967)
<http://link.aip.org/link/?JMP/8/2155/1>. In the case of a spacelike 2-surface (which is all we treat in this
function), the GHP eth operator reduces almost to the original. By comparing GHP's Eq. (3.8) to Eq. (2.13) of
Goldberg et al., we can see the difference. As GHP say below their Eq. (3.11) that "for complete agreement in
the case of a general 2-surface metric we would have (strictly speaking) to multiply up our eth operator by a
factor sqrt(2)."
Parameters
----------
modes : 1-d complex array
This array contains the modes starting from ell=ell_min, and continuing in standard order (as in sf.LM_range).
spin_weight : int
Spin weight of the input field. The eth operators raise the spin weight by 1.
ell_min : int, optional
Smallest ell value present in input `modes`. Defaults to 0.
Returns
-------
1-d complex array
The output has the same size as the input `modes`, and corresponds to the same (ell,m) values. Note, however,
that these modes have spin weight greater than the input by 1.
"""
ell_max = int(sqrt(len(modes) + LM_total_size(0, ell_min - 1))) - 1
eth_modes = np.array(modes)
i_mode = 0
for ell in xrange(ell_min, ell_max + 1):
factor = (0.0 if ell < abs(spin_weight + 1) else sqrt((ell - spin_weight) * (ell + spin_weight + 1.) / 2.))
for m in xrange(-ell, ell + 1):
eth_modes[i_mode] *= factor
i_mode += 1
return eth_modes
def ethbar_inverse_NP(modes, spin_weight, ell_min=0):
"""Inverse of the spin-lowering \bar{eth} operator as defined by Newman and Penrose
This function acts as a (partial) inverse or integral of the `ethbar_NP` operator. (See that function's
docstring for more information, including the calling signature.) In particular, composing the two functions in
either order should give (almost) the identity function. Essentially, this is an integral of the ethbar
function; so this is only a partial inverse because constants of integration could be added in some cases. The
difference between this function and `eth_NP` is mostly in the normalization.
"""
ell_max = int(sqrt(len(modes) + LM_total_size(0, ell_min - 1))) - 1
ethbar_inverse_modes = np.array(modes)
i_mode = 0
for ell in xrange(ell_min, ell_max + 1):
term = (ell + spin_weight + 1.) * (ell - spin_weight)
if term > 0.0:
factor = -sqrt(term)
for m in xrange(-ell, ell + 1):
ethbar_inverse_modes[i_mode] /= factor
i_mode += 1
else:
i_mode += 2*ell+1
return ethbar_inverse_modes
def _rotate_decomposition_basis_by_constant(data, ell_min, ell_max, D, tmp):
"""Rotate data by the same rotor at each point in time
`D` is the Wigner D matrix for all the ell values.
`tmp` is just a workspace used as temporary storage to hold the
results for each item of data during the sum.
"""
for i_t in xrange(data.shape[0]):
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):
tmp[i_m] = 0j
for i_mp in xrange(2 * ell + 1):
for i_m in xrange(2 * ell + 1):
tmp[i_m] += data[i_t, i_data +
i_mp] * D[i_D +
(2 * ell + 1) * i_mp + i_m]
for i_m in xrange(2 * ell + 1):
data[i_t, i_data + i_m] = tmp[i_m]
def definite_integral(f, t):
Sfdt = np.zeros_like(f)
for i in xrange(1, f.shape[0]):
Sfdt[i, ...] += (f[i, ...] + f[i - 1, ...]) * ((t[i] - t[i - 1]) / 2.0)
return Sfdt
def Wigner3j(j_1, j_2, j_3, m_1, m_2, m_3):
"""Calculate the Wigner 3j symbol `Wigner3j(j_1,j_2,j_3,m_1,m_2,m_3)`
This function is copied with minor modification from
sympy.physics.Wigner, as written by Jens Rasch.
The inputs must be integers. (Half integer arguments are
sacrificed so that we can use numba.) Nonzero return quantities
only occur when the `j`s obey the triangle inequality (any two
must add up to be as big as or bigger than the third).
Examples
========
>>> from spherical_functions import Wigner3j
>>> Wigner3j(2, 6, 4, 0, 0, 0)
0.186989398002
>>> Wigner3j(2, 6, 4, 0, 0, 1)
0
"""
if (m_1 + m_2 + m_3 != 0):
return 0
if ( (abs(m_1) > j_1) or (abs(m_2) > j_2) or (abs(m_3) > j_3) ):
return 0
prefid = (1 if (j_1 - j_2 - m_3) % 2 == 0 else -1)
m_3 = -m_3
a1 = j_1 + j_2 - j_3
if (a1 < 0):
return 0
a2 = j_1 - j_2 + j_3;
if (a2 < 0):
return 0
a3 = -j_1 + j_2 + j_3;
if (a3 < 0):
return 0
argsqrt = ( factorials[j_1 + j_2 - j_3] *
factorials[j_1 - j_2 + j_3] *
factorials[-j_1 + j_2 + j_3] *
factorials[j_1 - m_1] *
factorials[j_1 + m_1] *
factorials[j_2 - m_2] *
factorials[j_2 + m_2] *
factorials[j_3 - m_3] *
factorials[j_3 + m_3] ) / factorials[j_1 + j_2 + j_3 + 1]
ressqrt = sqrt(argsqrt)
imin = max(-j_3 + j_1 + m_2, max(-j_3 + j_2 - m_1, 0))
imax = min(j_2 + m_2, min(j_1 - m_1, j_1 + j_2 - j_3))
sumres = 0.0;
for ii in xrange(imin, imax + 1):
den = ( factorials[ii] *
factorials[ii + j_3 - j_1 - m_2] *
factorials[j_2 + m_2 - ii] *
factorials[j_1 - ii - m_1] *
factorials[ii + j_3 - j_2 + m_1] *
factorials[j_1 + j_2 - j_3 - ii] )
if (ii % 2 == 0):
sumres = sumres + 1.0 / den
else:
sumres = sumres - 1.0 / den
return ressqrt * sumres * prefid
def _derivative(f, t, dfdt):
for i in xrange(2):
t_i = t[i]
t1 = t[0]
t2 = t[1]
t3 = t[2]
t4 = t[3]
t5 = t[4]
h1 = t1 - t_i
h2 = t2 - t_i
h3 = t3 - t_i
h4 = t4 - t_i
h5 = t5 - t_i
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
dfdt[i] = (-((h2 * h3 * h4 + h2 * h3 * h5 + h2 * h4 * h5 + h3 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[0]
+ ((h1 * h3 * h4 + h1 * h3 * h5 + h1 * h4 * h5 + h3 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[1]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[2]
+ ((h1 * h2 * h3 + h1 * h2 * h5 + h1 * h3 * h5 + h2 * h3 * h5) / (h14 * h24 * h34 * h45)) * f[3]
- ((h1 * h2 * h3 + h1 * h2 * h4 + h1 * h3 * h4 + h2 * h3 * h4) / (h15 * h25 * h35 * h45)) * f[4])
for i in xrange(2, len(t) - 2):
t1 = t[i - 2]
t2 = t[i - 1]
t3 = t[i]
t4 = t[i + 1]
t5 = t[i + 2]
h1 = t1 - t3
h2 = t2 - t3
h4 = t4 - t3
h5 = t5 - t3
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
dfdt[i] = (-((h2 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[i - 2]
+ ((h1 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[i - 1]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[i]
+ ((h1 * h2 * h5) / (h14 * h24 * h34 * h45)) * f[i + 1]
- ((h1 * h2 * h4) / (h15 * h25 * h35 * h45)) * f[i + 2])
for i in xrange(len(t) - 2, len(t)):
t_i = t[i]
t1 = t[-5]
t2 = t[-4]
t3 = t[-3]
t4 = t[-2]
t5 = t[-1]
h1 = t1 - t_i
h2 = t2 - t_i
h3 = t3 - t_i
h4 = t4 - t_i
h5 = t5 - t_i
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
dfdt[i] = (-((h2 * h3 * h4 + h2 * h3 * h5 + h2 * h4 * h5 + h3 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[-5]
+ ((h1 * h3 * h4 + h1 * h3 * h5 + h1 * h4 * h5 + h3 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[-4]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[-3]
+ ((h1 * h2 * h3 + h1 * h2 * h5 + h1 * h3 * h5 + h2 * h3 * h5) / (h14 * h24 * h34 * h45)) * f[-2]
- ((h1 * h2 * h3 + h1 * h2 * h4 + h1 * h3 * h4 + h2 * h3 * h4) / (h15 * h25 * h35 * h45)) * f[-1])
return
def _LLMatrix(data, lm, LL):
"""Helper function for the LLMatrix function"""
# Big, bad, ugly, obvious way to do the calculation
# =================================================
# L+ = Lx + i Ly Lx = (L+ + L-) / 2 Im(Lx) = ( Im(L+) + Im(L-) ) / 2
# L- = Lx - i Ly Ly = -i (L+ - L-) / 2 Im(Ly) = -( Re(L+) - Re(L-) ) / 2
# Lz = Lz Lz = Lz Im(Lz) = Im(Lz)
# LxLx = (L+ + L-)(L+ + L-) / 4
# LxLy = -i(L+ + L-)(L+ - L-) / 4
# LxLz = (L+ + L-)( Lz ) / 2
# LyLx = -i(L+ - L-)(L+ + L-) / 4
# LyLy = -(L+ - L-)(L+ - L-) / 4
# LyLz = -i(L+ - L-)( Lz ) / 2
# LzLx = ( Lz )(L+ + L-) / 2
# LzLy = -i( Lz )(L+ - L-) / 2
# LzLz = ( Lz )( Lz )
for i_mode in xrange(lm.shape[0]):
L = lm[i_mode, 0]
M = lm[i_mode, 1]
for i_time in xrange(data.shape[0]):
# Compute first in (+,-,z) basis
LpLp = (np.conjugate(data[i_time, i_mode + 2]) *
data[i_time, i_mode] *
(ladder(L, M + 1) * ladder(L, M)) if M + 2 <= L else 0.0 +
0.0j)
LpLm = (np.conjugate(data[i_time, i_mode]) * data[i_time, i_mode] *
(ladder(L, M - 1) * ladder(L, -M))
if M - 1 >= -L else 0.0 + 0.0j)
LmLp = (np.conjugate(data[i_time, i_mode]) * data[i_time, i_mode] *
(ladder(L, -(M + 1)) * ladder(L, M))
if M + 1 <= L else 0.0 + 0.0j)
LmLm = (np.conjugate(data[i_time, i_mode - 2]) *
data[i_time, i_mode] *
(ladder(L, -(M - 1)) * ladder(L, -M))
if M - 2 >= -L else 0.0 + 0.0j)
LpLz = (np.conjugate(data[i_time, i_mode + 1]) *
data[i_time, i_mode] *
(ladder(L, M) * M) if M + 1 <= L else 0.0 + 0.0j)
LzLp = (np.conjugate(data[i_time, i_mode + 1]) *
data[i_time, i_mode] *
((M + 1) * ladder(L, M)) if M + 1 <= L else 0.0 + 0.0j)
LmLz = (np.conjugate(data[i_time, i_mode - 1]) *
data[i_time, i_mode] *
(ladder(L, -M) * M) if M - 1 >= -L else 0.0 + 0.0j)
LzLm = (np.conjugate(data[i_time, i_mode - 1]) *
data[i_time, i_mode] *
((M - 1) * ladder(L, -M)) if M - 1 >= -L else 0.0 + 0.0j)
LzLz = np.conjugate(data[i_time, i_mode]) * data[i_time,
i_mode] * M**2
# Convert into (x,y,z) basis
LxLx = 0.25 * (LpLp + LmLm + LmLp + LpLm)
LxLy = -0.25j * (LpLp - LmLm + LmLp - LpLm)
LxLz = 0.5 * (LpLz + LmLz)
LyLx = -0.25j * (LpLp - LmLp + LpLm - LmLm)
LyLy = -0.25 * (LpLp - LmLp - LpLm + LmLm)
LyLz = -0.5j * (LpLz - LmLz)
LzLx = 0.5 * (LzLp + LzLm)
LzLy = -0.5j * (LzLp - LzLm)
# LzLz = (LzLz)
# Symmetrize
LL[i_time, 0, 0] += LxLx.real
LL[i_time, 0, 1] += (LxLy + LyLx).real / 2.0
LL[i_time, 0, 2] += (LxLz + LzLx).real / 2.0
LL[i_time, 1, 0] += (LyLx + LxLy).real / 2.0
LL[i_time, 1, 1] += LyLy.real
LL[i_time, 1, 2] += (LyLz + LzLy).real / 2.0
LL[i_time, 2, 0] += (LzLx + LxLz).real / 2.0
LL[i_time, 2, 1] += (LzLy + LyLz).real / 2.0
LL[i_time, 2, 2] += LzLz.real
return
def _derivative_3d(f, t, dfdt):
for i in xrange(2):
t_i = t[i]
t1 = t[0]
t2 = t[1]
t3 = t[2]
t4 = t[3]
t5 = t[4]
h1 = t1 - t_i
h2 = t2 - t_i
h3 = t3 - t_i
h4 = t4 - t_i
h5 = t5 - t_i
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
for k in xrange(f.shape[1]):
for m in xrange(f.shape[1]):
dfdt[i, k, m] = (
-((h2 * h3 * h4 + h2 * h3 * h5 + h2 * h4 * h5 + h3 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[0, k, m]
+ ((h1 * h3 * h4 + h1 * h3 * h5 + h1 * h4 * h5 + h3 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[1, k, m]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[2, k, m]
+ ((h1 * h2 * h3 + h1 * h2 * h5 + h1 * h3 * h5 + h2 * h3 * h5) / (h14 * h24 * h34 * h45)) * f[3, k, m]
- ((h1 * h2 * h3 + h1 * h2 * h4 + h1 * h3 * h4 + h2 * h3 * h4) / (h15 * h25 * h35 * h45)) * f[4, k, m])
for i in xrange(2, len(t) - 2):
t1 = t[i - 2]
t2 = t[i - 1]
t3 = t[i]
t4 = t[i + 1]
t5 = t[i + 2]
h1 = t1 - t3
h2 = t2 - t3
h4 = t4 - t3
h5 = t5 - t3
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
for k in xrange(f.shape[1]):
for m in xrange(f.shape[1]):
dfdt[i, k, m] = (-((h2 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[i - 2, k, m]
+ ((h1 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[i - 1, k, m]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35))
* f[i, k, m]
+ ((h1 * h2 * h5) / (h14 * h24 * h34 * h45)) * f[i + 1, k, m]
- ((h1 * h2 * h4) / (h15 * h25 * h35 * h45)) * f[i + 2, k, m])
for i in xrange(len(t) - 2, len(t)):
t_i = t[i]
t1 = t[-5]
t2 = t[-4]
t3 = t[-3]
t4 = t[-2]
t5 = t[-1]
h1 = t1 - t_i
h2 = t2 - t_i
h3 = t3 - t_i
h4 = t4 - t_i
h5 = t5 - t_i
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
for k in xrange(f.shape[1]):
for m in xrange(f.shape[1]):
dfdt[i, k, m] = (
-((h2 * h3 * h4 + h2 * h3 * h5 + h2 * h4 * h5 + h3 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[-5, k, m]
+ ((h1 * h3 * h4 + h1 * h3 * h5 + h1 * h4 * h5 + h3 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[-4, k, m]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[-3, k, m]
+ ((h1 * h2 * h3 + h1 * h2 * h5 + h1 * h3 * h5 + h2 * h3 * h5) / (h14 * h24 * h34 * h45)) * f[-2, k, m]
- ((h1 * h2 * h3 + h1 * h2 * h4 + h1 * h3 * h4 + h2 * h3 * h4) / (h15 * h25 * h35 * h45)) * f[-1, k, m])
return
def definite_integral(f, t):
Sfdt = np.zeros_like(f)
for i in xrange(1, f.shape[0]):
Sfdt[i, ...] += (f[i, ...] + f[i - 1, ...]) * ((t[i] - t[i - 1]) / 2.0)
return Sfdt
def complex_array_norm(c, s):
for i in xrange(len(s)):
s[i] = 0.0
for j in xrange(c.shape[1]):
s[i] += c[i, j].real ** 2 + c[i, j].imag ** 2
return
def _SWSH(Ra, Rb, s, indices, values):
"""Compute spin-weighted spherical harmonics from rotor components
This is the core function that does all the work in the
computation, but it is strict about its inputs, and does not check
them for validity -- though numba provides some degree of safety.
_SWSH(Ra, Rb, s, indices, values)
Parameters
----------
Ra : complex
Component `a` of the rotor
Rb : complex
Component `b` of the rotor
s : int
Spin weight of the field to evaluate
indices : 2-d array of int
Array of (ell,m) values to evaluate
values : 1-d array of complex
Output array to contain values. Length must equal first dimension of `indices`. Needed because numba cannot
create arrays at the moment.
Returns
-------
void
The input/output array `values` is modified in place.
"""
N = indices.shape[0]
# These constants are the recurring quantities in the computation
# of the matrix values, so we calculate them here just once
ra, phia = cmath.polar(Ra)
rb, phib = cmath.polar(Rb)
if ra <= epsilon:
for i in xrange(N):
ell, m = indices[i, 0:2]
if (m != s or abs(m) > ell or abs(s) > ell):
values[i] = 0.0j
else:
if (ell) % 2 == 0:
values[i] = math.sqrt((2 * ell + 1) / (4 * np.pi)) * Rb ** (-2 * s)
else:
values[i] = -math.sqrt((2 * ell + 1) / (4 * np.pi)) * Rb ** (-2 * s)
elif rb <= epsilon:
for i in xrange(N):
ell, m = indices[i, 0:2]
if (m != -s or abs(m) > ell or abs(s) > ell):
values[i] = 0.0j
else:
if (s) % 2 == 0:
values[i] = math.sqrt((2 * ell + 1) / (4 * np.pi)) * Ra ** (-2 * s)
else:
values[i] = -math.sqrt((2 * ell + 1) / (4 * np.pi)) * Ra ** (-2 * s)
elif ra < rb:
# We have to have these two versions (both this ra<rb branch,
# and ra>=rb below) to avoid overflows and underflows
absRRatioSquared = -ra * ra / (rb * rb)
for i in xrange(N):
ell, m = indices[i, 0:2]
if (abs(m) > ell or abs(s) > ell):
values[i] = 0.0j
else:
rhoMin = max(0, -m + s)
# Protect against overflow by decomposing Ra,Rb as
# abs,angle components and pulling out the factor of
# absRRatioSquared**rhoMin. Here, ra might be quite
# small, in which case ra**(-s+m) could be enormous
# when the exponent (-s+m) is very negative; adding
# 2*rhoMin to the exponent ensures that it is always
# positive, which protects from overflow. Meanwhile,
# underflow just goes to zero, which is fine since
# nothing else should be very large.
Prefactor = cmath.rect(
coeff(ell, -m, -s) * rb ** (2 * ell + s - m - 2 * rhoMin) * ra ** (-s + m + 2 * rhoMin),
phib * (-s - m) + phia * (-s + m))
if (Prefactor == 0.0j):
values[i] = 0.0j
else:
if ((ell + rhoMin) % 2 != 0):
Prefactor *= -1
rhoMax = min(ell - m, ell + s)
N1 = ell - m + 1
N2 = ell + s + 1
M = -s + m
Sum = 1.0
for rho in xrange(rhoMax, rhoMin, -1):
Sum *= absRRatioSquared * ((N1 - rho) * (N2 - rho)) / (rho * (M + rho))
Sum += 1
# Sum = 0.0
# for rho in xrange(rhoMax, rhoMin-1, -1):
# Sum = ( binomial_coefficient(ell-m,rho) * binomial_coefficient(ell+m, ell-rho+s)
# + Sum * absRRatioSquared )
values[i] = math.sqrt((2 * ell + 1) / (4 * np.pi)) * Prefactor * Sum
else: # ra >= rb
# We have to have these two versions (both this ra>=rb branch,
# and ra<rb above) to avoid overflows and underflows
absRRatioSquared = -rb * rb / (ra * ra)
for i in xrange(N):
ell, m = indices[i, 0:2]
if (abs(m) > ell or abs(s) > ell):
values[i] = 0.0j
else:
rhoMin = max(0, m + s)
# Protect against overflow by decomposing Ra,Rb as
# abs,angle components and pulling out the factor of
# absRRatioSquared**rhoMin. Here, rb might be quite
# small, in which case rb**(-s-m) could be enormous
# when the exponent (-s-m) is very negative; adding
# 2*rhoMin to the exponent ensures that it is always
# positive, which protects from overflow. Meanwhile,
# underflow just goes to zero, which is fine since
# nothing else should be very large.
Prefactor = cmath.rect(
coeff(ell, m, -s) * ra ** (2 * ell + s + m - 2 * rhoMin) * rb ** (-s - m + 2 * rhoMin),
phia * (-s + m) + phib * (-s - m))
if (Prefactor == 0.0j):
values[i] = 0.0j
else:
if ((rhoMin + s) % 2 != 0):
Prefactor *= -1
rhoMax = min(ell + m, ell + s)
N1 = ell + m + 1
N2 = ell + s + 1
M = -s - m
Sum = 1.0
for rho in xrange(rhoMax, rhoMin, -1):
Sum *= absRRatioSquared * ((N1 - rho) * (N2 - rho)) / (rho * (M + rho))
Sum += 1
# Sum = 0.0
# for rho in xrange(rhoMax, rhoMin-1, -1):
# Sum = ( binomial_coefficient(ell+m,rho) * binomial_coefficient(ell-m, ell-rho+s)
# + Sum * absRRatioSquared )
values[i] = math.sqrt((2 * ell + 1) / (4 * np.pi)) * Prefactor * Sum
def _Wigner_D_elements(Rs, ell, mp, m, values):
"""Main work function for computing Wigner D matrix elements
This is the core function that does all the work in the
computation, but it is strict about its input, and does not check
them for validity.
Input arguments
===============
_Wigner_D_matrices(Rs, ell, m, mp, elements)
The `matrices` variable is needed because numba cannot create
arrays at the moment, but this is modified in place, so after
calling this function, the input array will contain the correct
values.
"""
N = Rs.shape[0]
if (abs(m) > ell or abs(mp) > ell):
for i in xrange(N):
values[i] = 0.0j
else:
rhoMin_a = max(0, mp - m)
rhoMax_a = min(ell + mp, ell - m)
coefficient_a = coeff(ell, mp, m)
if (rhoMin_a % 2 != 0):
coefficient_a *= -1
N1_a = ell + mp + 1
N2_a = ell - m + 1
M_a = m - mp
rhoMin_b = max(0, -m - mp)
rhoMax_b = min(ell - mp, ell - m)
coefficient_b = coeff(ell, -mp, m)
if ((ell + m + rhoMin_b) % 2 != 0):
coefficient_b *= -1
N1_b = ell - mp + 1
N2_b = ell - m + 1
M_b = m + mp
for i in xrange(N):
Ra = complex(Rs[i, 0], Rs[i, 3])
ra, phia = cmath.polar(Ra)
Rb = complex(Rs[i, 2], Rs[i, 1])
rb, phib = cmath.polar(Rb)
if ra <= epsilon:
if mp != -m:
values[i] = 0.0j
elif (ell + mp) % 2 == 0:
values[i] = Rb**(-2 * mp)
else:
values[i] = -(Rb**(-2 * mp))
elif rb <= epsilon:
if mp != m:
values[i] = 0.0j
else:
values[i] = Ra**(2 * mp)
elif ra < rb:
absRRatioSquared = -ra * ra / (rb * rb)
d = coefficient_b * rb**(2 * ell - m - mp - 2 * rhoMin_b
) * ra**(m + mp + 2 * rhoMin_b)
if d == 0.0j:
values[i] = 0.0j
else:
Prefactor1 = cmath.rect(d,
phib * (m - mp) + phia * (m + mp))
Prefactor2 = cmath.rect(d, (phib + np.pi) * (m - mp) -
phia * (m + mp))
Sum = 1.0
for rho in xrange(rhoMax_b, rhoMin_b, -1):
Sum *= absRRatioSquared * (
(N1_b - rho) * (N2_b - rho)) / (rho * (M_b + rho))
Sum += 1
values[i] = Prefactor1 * Sum
else: # ra >= rb
absRRatioSquared = -rb * rb / (ra * ra)
d = coefficient_a * ra**(2 * ell - m + mp - 2 * rhoMin_a
) * rb**(m - mp + 2 * rhoMin_a)
if d == 0.0j:
values[i] = 0.0j
else:
Prefactor1 = cmath.rect(d,
phia * (m + mp) + phib * (m - mp))
Prefactor2 = cmath.rect(
d, -phia * (m + mp) + (phib + np.pi) * (m - mp))
Sum = 1.0
for rho in xrange(rhoMax_a, rhoMin_a, -1):
Sum *= absRRatioSquared * (
(N1_a - rho) * (N2_a - rho)) / (rho * (M_a + rho))
Sum += 1
values[i] = Prefactor1 * Sum
def _Wigner_D_matrices(Ra, Rb, ell_min, ell_max, matrices):
"""Main work function for `Wigner_D_matrices`
This is the core function that does all the work in the
computation, but it is strict about its input, and does not check
them for validity.
Input arguments
===============
_Wigner_D_matrices(Ra, Rb, ell_min, ell_max, elements)
* Ra, Rb are the complex components of the rotor
* ell_min, ell_max are the limits of the matrices
* matrix is a one-dimensional array of complex numbers to be
filled with the elements of the matrices; the correct shape is
assumed
The `matrices` variable is needed because numba cannot create
arrays at the moment, but this is modified in place, so after
calling this function, the input array will contain the correct
values.
"""
# These constants are the recurring quantities in the computation
# of the matrix elements, so we calculate them here just once
ra, phia = cmath.polar(Ra)
rb, phib = cmath.polar(Rb)
if (ra <= epsilon):
for ell in xrange(ell_min, ell_max + 1):
i_ell = _linear_matrix_offset(ell, ell_min)
for i in xrange((2 * ell + 1)**2):
matrices[i_ell + i] = 0j
for mpmm in xrange(-ell, ell + 1):
i_mpmm = _linear_matrix_index(ell, mpmm, -mpmm)
if (ell + mpmm) % 2 == 0:
matrices[i_ell + i_mpmm] = Rb**(-2 * mpmm)
else:
matrices[i_ell + i_mpmm] = -(Rb**(-2 * mpmm))
elif (rb <= epsilon):
for ell in xrange(ell_min, ell_max + 1):
i_ell = _linear_matrix_offset(ell, ell_min)
for i in xrange((2 * ell + 1)**2):
matrices[i_ell + i] = 0j
for mpm in xrange(-ell, ell + 1):
i_mpm = _linear_matrix_diagonal_index(ell, mpm)
matrices[i_ell + i_mpm] = Ra**(2 * mpm)
elif (ra < rb):
# We have to have these two versions (both this ra<rb branch,
# and ra>=rb below) to avoid overflows and underflows
absRRatioSquared = -ra * ra / (rb * rb)
for ell in xrange(ell_min, ell_max + 1):
i_ell = _linear_matrix_offset(ell, ell_min)
for mp in xrange(-ell, 1):
for m in xrange(mp, -mp + 1):
i_mpm = _linear_matrix_index(ell, mp, m)
rhoMin = max(0, -mp - m)
# Protect against overflow by decomposing Ra,Rb as
# abs,angle components and pulling out the factor of
# absRRatioSquared**rhoMin. Here, ra might be quite
# small, in which case ra**(m+mp) could be enormous
# when the exponent (m+mp) is very negative; adding
# 2*rhoMin to the exponent ensures that it is always
# positive, which protects from overflow. Meanwhile,
# underflow just goes to zero, which is fine since
# nothing else should be very large.
d = coeff(ell, -mp, m) * rb**(2 * ell - m - mp - 2 * rhoMin
) * ra**(m + mp + 2 * rhoMin)
if (d == 0.0j):
matrices[i_ell + i_mpm] = 0.0j
if (abs(m) != abs(mp)):
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell +
_linear_matrix_index(ell, -mp, -m)] = 0.0j
# D_{m,mp}(R) = \bar{D}_{mp,m}(\bar{R})
matrices[i_ell +
_linear_matrix_index(ell, m, mp)] = 0.0j
# D_{-m,-mp}(R) = (-1)^{mp+m} D_{mp,m}(\bar{R})
matrices[i_ell +
_linear_matrix_index(ell, -m, -mp)] = 0.0j
elif (m != 0):
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell +
_linear_matrix_index(ell, -mp, -m)] = 0.0j
else:
if ((ell + m + rhoMin) % 2 != 0):
d *= -1
Prefactor1 = cmath.rect(
d,
phib * (m - mp) + phia * (m + mp))
Prefactor2 = cmath.rect(d, (phib + np.pi) * (m - mp) -
phia * (m + mp))
rhoMax = min(ell - mp, ell - m)
N1 = ell - mp + 1
N2 = ell - m + 1
M = m + mp
Sum = 1.0
for rho in xrange(rhoMax, rhoMin, -1):
Sum *= absRRatioSquared * (
(N1 - rho) * (N2 - rho)) / (rho * (M + rho))
Sum += 1
# Sum *= absRRatioSquared**rhoMin
matrices[i_ell + i_mpm] = Prefactor1 * Sum
if (abs(m) != abs(mp)):
if ((m + mp) % 2 == 0):
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell + _linear_matrix_index(
ell, -mp,
-m)] = Prefactor1.conjugate() * Sum
# D_{-m,-mp}(R) = (-1)^{mp+m} D_{mp,m}(\bar{R})
matrices[i_ell + _linear_matrix_index(
ell, -m, -mp)] = Prefactor2 * Sum
else:
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell + _linear_matrix_index(
ell, -mp,
-m)] = -Prefactor1.conjugate() * Sum
# D_{-m,-mp}(R) = (-1)^{mp+m} D_{mp,m}(\bar{R})
matrices[i_ell + _linear_matrix_index(
ell, -m, -mp)] = -Prefactor2 * Sum
# D_{m,mp}(R) = \bar{D}_{mp,m}(\bar{R})
matrices[i_ell + _linear_matrix_index(
ell, m, mp)] = Prefactor2.conjugate() * Sum
elif (m != 0):
if ((m + mp) % 2 == 0):
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell + _linear_matrix_index(
ell, -mp,
-m)] = Prefactor1.conjugate() * Sum
else:
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell + _linear_matrix_index(
ell, -mp,
-m)] = -Prefactor1.conjugate() * Sum
else: # ra >= rb
# We have to have these two versions (both this ra>=rb branch,
# and ra<rb above) to avoid overflows and underflows
absRRatioSquared = -rb * rb / (ra * ra)
for ell in xrange(ell_min, ell_max + 1):
i_ell = _linear_matrix_offset(ell, ell_min)
for mp in xrange(-ell, 1):
for m in xrange(mp, -mp + 1):
i_mpm = _linear_matrix_index(ell, mp, m)
rhoMin = max(0, mp - m)
# Protect against overflow by decomposing Ra,Rb as
# abs,angle components and pulling out the factor of
# absRRatioSquared**rhoMin. Here, rb might be quite
# small, in which case rb**(m-mp) could be enormous
# when the exponent (m-mp) is very negative; adding
# 2*rhoMin to the exponent ensures that it is always
# positive, which protects from overflow. Meanwhile,
# underflow just goes to zero, which is fine since
# nothing else should be very large.
d = coeff(ell, mp, m) * ra**(2 * ell - m + mp - 2 * rhoMin
) * rb**(m - mp + 2 * rhoMin)
if (d == 0.0j):
matrices[i_ell + i_mpm] = 0.0j
if (abs(m) != abs(mp)):
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell +
_linear_matrix_index(ell, -mp, -m)] = 0.0j
# D_{m,mp}(R) = \bar{D}_{mp,m}(\bar{R})
matrices[i_ell +
_linear_matrix_index(ell, m, mp)] = 0.0j
# D_{-m,-mp}(R) = (-1)^{mp+m} D_{mp,m}(\bar{R})
matrices[i_ell +
_linear_matrix_index(ell, -m, -mp)] = 0.0j
elif (m != 0):
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell +
_linear_matrix_index(ell, -mp, -m)] = 0.0j
else:
if (rhoMin % 2 != 0):
d *= -1
Prefactor1 = cmath.rect(
d,
phia * (m + mp) + phib * (m - mp))
Prefactor2 = cmath.rect(
d, -phia * (m + mp) + (phib + np.pi) * (m - mp))
rhoMax = min(ell + mp, ell - m)
N1 = ell + mp + 1
N2 = ell - m + 1
M = m - mp
Sum = 1.0
for rho in xrange(rhoMax, rhoMin, -1):
Sum *= absRRatioSquared * (
(N1 - rho) * (N2 - rho)) / (rho * (M + rho))
Sum += 1
# Sum *= absRRatioSquared**rhoMin
matrices[i_ell + i_mpm] = Prefactor1 * Sum
if (abs(m) != abs(mp)):
if (m + mp) % 2 == 0:
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell + _linear_matrix_index(
ell, -mp,
-m)] = Prefactor1.conjugate() * Sum
# D_{-m,-mp}(R) = (-1)^{mp+m} D_{mp,m}(\bar{R})
matrices[i_ell + _linear_matrix_index(
ell, -m, -mp)] = Prefactor2 * Sum
else:
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell + _linear_matrix_index(
ell, -mp,
-m)] = -Prefactor1.conjugate() * Sum
# D_{-m,-mp}(R) = (-1)^{mp+m} D_{mp,m}(\bar{R})
matrices[i_ell + _linear_matrix_index(
ell, -m, -mp)] = -Prefactor2 * Sum
# D_{m,mp}(R) = \bar{D}_{mp,m}(\bar{R})
matrices[i_ell + _linear_matrix_index(
ell, m, mp)] = Prefactor2.conjugate() * Sum
elif (m != 0):
if (m + mp) % 2 == 0:
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell + _linear_matrix_index(
ell, -mp,
-m)] = Prefactor1.conjugate() * Sum
else:
# D_{-mp,-m}(R) = (-1)^{mp+m} \bar{D}_{mp,m}(R)
matrices[i_ell + _linear_matrix_index(
ell, -mp,
-m)] = -Prefactor1.conjugate() * Sum
def Wigner_D_element(*args):
"""Return elements of the Wigner D matrices
The conventions used for this function are discussed more fully on
<http://moble.github.io/spherical_functions/>.
Input arguments
===============
The input can be in any of the following forms:
Wigner_D_element(R, ell, mp, m)
Wigner_D_element(Rs, ell, mp, m)
Wigner_D_element(R, indices)
Wigner_D_element(Ra, Rb, ell, mp, m)
Wigner_D_element(Ra, Rb, indices)
Wigner_D_element(alpha, beta, gamma, ell, mp, m)
Wigner_D_element(alpha, beta, gamma, indices)
Where
* R is a unit quaternion (no checking of norm is done)
* Rs is an array of unit quaternions (no checking of norm is done)
* Ra and Rb are the complex parts of a unit quaternion
* alpha, beta, gamma are the Euler angles [shudder...]
* ell, mp, m are the integer or half-integer indices of the
D matrix element
* indices is an array of [ell,mp,m] indices as above, or simply
a list of ell modes, in which case all valid [mp,m] values
will be returned
Note that, by default, a ValueError will be raised if the input
(ell, mp, m) values are not valid. (For example, |m|>ell.)
If instead, you would simply like a return value of 0.0, after
importing this module as sf, simply evaluate
>>> sf.error_on_bad_indices = False
Return value
============
One complex number is returned for each component requested. If a
single quaternion and the (ell,mp,m) arguments were given explicitly,
this means that a single complex scalar is returned. If more than
one component was requested, a one-dimensional numpy array of complex
scalars is returned, in the same order as the input.
"""
# Find the rotation from the args
if isinstance(args[0], np.ndarray):
elements = np.empty_like(args[0], dtype=complex)
_Wigner_D_elements(quaternion.as_float_array(args[0]), args[1],
args[2], args[3], elements)
return elements
elif isinstance(args[0], np.quaternion):
# The rotation is input as a single quaternion
Ra = args[0].a
Rb = args[0].b
mode_offset = 1
elif isinstance(args[0], numbers.Number) and isinstance(args[1], numbers.Number)\
and isinstance(args[2], numbers.Number):
# UUUUGGGGLLLLYYYY. The rotation is input as Euler angles
R = quaternion.from_euler_angles(args[0], args[1], args[2])
Ra = R.a
Rb = R.b
mode_offset = 3
elif isinstance(args[0], numbers.Complex) and isinstance(
args[1], numbers.Complex):
# The rotation is input as the two parts of a single quaternion
Ra = args[0]
Rb = args[1]
mode_offset = 2
else:
raise ValueError("Can't understand input rotation")
# Find the indices
return_scalar = False
if (len(args) - mode_offset == 3):
# Assume these are the (ell, mp, m) indices
ell, mp, m = args[mode_offset:]
indices = np.array([
[round(2 * ell), round(2 * mp),
round(2 * m)],
],
dtype=int)
if (error_on_bad_indices and not _check_valid_indices(*(indices[0]))):
raise ValueError(
"(ell,mp,m)=({0},{1},{2})".format(ell, mp, m) +
" is not a valid set of indices for Wigner's D matrix")
return_scalar = True
elif (len(args) - mode_offset == 1):
indices = np.round(2 * np.asarray(args[mode_offset])).astype(int)
if (indices.ndim == 0 and indices.size == 1):
# This was just a single ell value
twoell = indices[0] # already multiplied by 2
indices = np.array([[twoell, twomp, twom]
for twomp in xrange(-twoell, twoell + 1, 2)
for twom in xrange(-twoell, twoell + 1, 2)])
if (twoell == 0):
return_scalar = True
elif (indices.ndim == 1 and indices.size > 0):
# This a list of ell values
indices = np.array([[twoell, twomp, twom] for twoell in indices
for twomp in xrange(-twoell, twoell + 1, 2)
for twom in xrange(-twoell, twoell + 1, 2)])
elif (indices.ndim == 2):
# This is an array of [ell,mp,m] values
if (error_on_bad_indices):
for twoell, twomp, twom in indices:
if not _check_valid_indices(twoell, twomp, twom):
raise ValueError(
"(ell,mp,m)=({0},{1},{2}) is not a valid set".
format(twoell / 2, twomp / 2, twom / 2) +
" of indices for Wigner's D matrix")
else:
raise ValueError("Can't understand input indices")
else:
raise ValueError("Can't understand input indices")
elements = np.empty((len(indices), ), dtype=complex)
_Wigner_D_element(Ra, Rb, indices, elements)
if (return_scalar):
return elements[0]
return elements
def _derivative(f, t, dfdt):
for i in xrange(2):
t_i = t[i]
t1 = t[0]
t2 = t[1]
t3 = t[2]
t4 = t[3]
t5 = t[4]
h1 = t1 - t_i
h2 = t2 - t_i
h3 = t3 - t_i
h4 = t4 - t_i
h5 = t5 - t_i
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
dfdt[i] = (-((h2 * h3 * h4 + h2 * h3 * h5 + h2 * h4 * h5 + h3 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[0]
+ ((h1 * h3 * h4 + h1 * h3 * h5 + h1 * h4 * h5 + h3 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[1]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[2]
+ ((h1 * h2 * h3 + h1 * h2 * h5 + h1 * h3 * h5 + h2 * h3 * h5) / (h14 * h24 * h34 * h45)) * f[3]
- ((h1 * h2 * h3 + h1 * h2 * h4 + h1 * h3 * h4 + h2 * h3 * h4) / (h15 * h25 * h35 * h45)) * f[4])
for i in xrange(2, len(t) - 2):
t1 = t[i - 2]
t2 = t[i - 1]
t3 = t[i]
t4 = t[i + 1]
t5 = t[i + 2]
h1 = t1 - t3
h2 = t2 - t3
h4 = t4 - t3
h5 = t5 - t3
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
dfdt[i] = (-((h2 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[i - 2]
+ ((h1 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[i - 1]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[i]
+ ((h1 * h2 * h5) / (h14 * h24 * h34 * h45)) * f[i + 1]
- ((h1 * h2 * h4) / (h15 * h25 * h35 * h45)) * f[i + 2])
for i in xrange(len(t) - 2, len(t)):
t_i = t[i]
t1 = t[-5]
t2 = t[-4]
t3 = t[-3]
t4 = t[-2]
t5 = t[-1]
h1 = t1 - t_i
h2 = t2 - t_i
h3 = t3 - t_i
h4 = t4 - t_i
h5 = t5 - t_i
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
dfdt[i] = (-((h2 * h3 * h4 + h2 * h3 * h5 + h2 * h4 * h5 + h3 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[-5]
+ ((h1 * h3 * h4 + h1 * h3 * h5 + h1 * h4 * h5 + h3 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[-4]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[-3]
+ ((h1 * h2 * h3 + h1 * h2 * h5 + h1 * h3 * h5 + h2 * h3 * h5) / (h14 * h24 * h34 * h45)) * f[-2]
- ((h1 * h2 * h3 + h1 * h2 * h4 + h1 * h3 * h4 + h2 * h3 * h4) / (h15 * h25 * h35 * h45)) * f[-1])
return
def _SWSHs(Rs, s, ell, m, values):
"""Compute spin-weighted spherical harmonics from rotor components
This is the core function that does all the work in the
computation, but it is strict about its inputs, and does not check
them for validity -- though numba provides some degree of safety.
_SWSHs(Rs, s, ell, m, values)
Parameters
----------
Rs : 2-d array of float
Components of the rotors, with the 0 index iterating over rotor, and the 1 index iterating over component.
s : int
Spin weight of the field to evaluate
ell : int
m : int
Values of (ell,m) to output
values : 1-d array of complex
Output array to contain values. Length must equal 0 dimension of `Rs`. Needed because numba cannot create
arrays at the moment.
Returns
-------
void
The input/output array `values` is modified in place.
"""
N = Rs.shape[0]
if (abs(m) > ell or abs(s) > ell):
for i in xrange(N):
values[i] = 0.0j
else:
constant = math.sqrt((2 * ell + 1) / (4 * np.pi))
ell_even = (ell % 2 == 0)
s_even = (s % 2 == 0)
rhoMin_a = max(0, m + s)
rhoMax_a = min(ell + m, ell + s)
coefficient_a = coeff(ell, m, -s)
if ((rhoMin_a + s) % 2 != 0):
coefficient_a *= -1
N1_a = ell + m + 1
N2_a = ell + s + 1
M_a = -s - m
rhoMin_b = max(0, -m + s)
rhoMax_b = min(ell - m, ell + s)
coefficient_b = coeff(ell, -m, -s)
if ((ell + rhoMin_b) % 2 != 0):
coefficient_b *= -1
N1_b = ell - m + 1
N2_b = ell + s + 1
M_b = -s + m
for i in xrange(N):
Ra = complex(Rs[i, 0], Rs[i, 3])
ra, phia = cmath.polar(Ra)
Rb = complex(Rs[i, 2], Rs[i, 1])
rb, phib = cmath.polar(Rb)
if ra <= epsilon:
if m != s:
values[i] = 0.0j
elif ell_even:
values[i] = constant * Rb ** (-2 * s)
else:
values[i] = -constant * Rb ** (-2 * s)
elif rb <= epsilon:
if m != -s:
values[i] = 0.0j
elif s_even:
values[i] = constant * Ra ** (-2 * s)
else:
values[i] = -constant * Ra ** (-2 * s)
elif ra < rb:
if (coefficient_b == 0.0j):
values[i] = 0.0j
else:
absRRatioSquared = -ra * ra / (rb * rb)
Prefactor = cmath.rect(
coefficient_b * rb ** (2 * ell + s - m - 2 * rhoMin_b) * ra ** (-s + m + 2 * rhoMin_b),
phib * (-s - m) + phia * (-s + m))
Sum = 1.0
for rho in xrange(rhoMax_b, rhoMin_b, -1):
Sum *= absRRatioSquared * ((N1_b - rho) * (N2_b - rho)) / (rho * (M_b + rho))
Sum += 1
values[i] = constant * Prefactor * Sum
else: # ra >= rb
if (coefficient_a == 0.0j):
values[i] = 0.0j
else:
absRRatioSquared = -rb * rb / (ra * ra)
Prefactor = cmath.rect(
coefficient_a * ra ** (2 * ell + s + m - 2 * rhoMin_a) * rb ** (-s - m + 2 * rhoMin_a),
phia * (-s + m) + phib * (-s - m))
Sum = 1.0
for rho in xrange(rhoMax_a, rhoMin_a, -1):
Sum *= absRRatioSquared * ((N1_a - rho) * (N2_a - rho)) / (rho * (M_a + rho))
Sum += 1
values[i] = constant * Prefactor * Sum
def _derivative_3d(f, t, dfdt):
for i in xrange(2):
t_i = t[i]
t1 = t[0]
t2 = t[1]
t3 = t[2]
t4 = t[3]
t5 = t[4]
h1 = t1 - t_i
h2 = t2 - t_i
h3 = t3 - t_i
h4 = t4 - t_i
h5 = t5 - t_i
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
for k in xrange(f.shape[1]):
for m in xrange(f.shape[1]):
dfdt[i, k, m] = (
-((h2 * h3 * h4 + h2 * h3 * h5 + h2 * h4 * h5 + h3 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[0, k, m]
+ ((h1 * h3 * h4 + h1 * h3 * h5 + h1 * h4 * h5 + h3 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[1, k, m]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[2, k, m]
+ ((h1 * h2 * h3 + h1 * h2 * h5 + h1 * h3 * h5 + h2 * h3 * h5) / (h14 * h24 * h34 * h45)) * f[3, k, m]
- ((h1 * h2 * h3 + h1 * h2 * h4 + h1 * h3 * h4 + h2 * h3 * h4) / (h15 * h25 * h35 * h45)) * f[4, k, m])
for i in xrange(2, len(t) - 2):
t1 = t[i - 2]
t2 = t[i - 1]
t3 = t[i]
t4 = t[i + 1]
t5 = t[i + 2]
h1 = t1 - t3
h2 = t2 - t3
h4 = t4 - t3
h5 = t5 - t3
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
for k in xrange(f.shape[1]):
for m in xrange(f.shape[1]):
dfdt[i, k, m] = (-((h2 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[i - 2, k, m]
+ ((h1 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[i - 1, k, m]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35))
* f[i, k, m]
+ ((h1 * h2 * h5) / (h14 * h24 * h34 * h45)) * f[i + 1, k, m]
- ((h1 * h2 * h4) / (h15 * h25 * h35 * h45)) * f[i + 2, k, m])
for i in xrange(len(t) - 2, len(t)):
t_i = t[i]
t1 = t[-5]
t2 = t[-4]
t3 = t[-3]
t4 = t[-2]
t5 = t[-1]
h1 = t1 - t_i
h2 = t2 - t_i
h3 = t3 - t_i
h4 = t4 - t_i
h5 = t5 - t_i
h12 = t1 - t2
h13 = t1 - t3
h14 = t1 - t4
h15 = t1 - t5
h23 = t2 - t3
h24 = t2 - t4
h25 = t2 - t5
h34 = t3 - t4
h35 = t3 - t5
h45 = t4 - t5
for k in xrange(f.shape[1]):
for m in xrange(f.shape[1]):
dfdt[i, k, m] = (
-((h2 * h3 * h4 + h2 * h3 * h5 + h2 * h4 * h5 + h3 * h4 * h5) / (h12 * h13 * h14 * h15)) * f[-5, k, m]
+ ((h1 * h3 * h4 + h1 * h3 * h5 + h1 * h4 * h5 + h3 * h4 * h5) / (h12 * h23 * h24 * h25)) * f[-4, k, m]
- ((h1 * h2 * h4 + h1 * h2 * h5 + h1 * h4 * h5 + h2 * h4 * h5) / (h13 * h23 * h34 * h35)) * f[-3, k, m]
+ ((h1 * h2 * h3 + h1 * h2 * h5 + h1 * h3 * h5 + h2 * h3 * h5) / (h14 * h24 * h34 * h45)) * f[-2, k, m]
- ((h1 * h2 * h3 + h1 * h2 * h4 + h1 * h3 * h4 + h2 * h3 * h4) / (h15 * h25 * h35 * h45)) * f[-1, k, m])
return
def _SWSHs(Rs, s, ell, m, values):
"""Compute spin-weighted spherical harmonics from rotor components
This is the core function that does all the work in the
computation, but it is strict about its inputs, and does not check
them for validity -- though numba provides some degree of safety.
_SWSHs(Rs, s, ell, m, values)
Parameters
----------
Rs : 2-d array of float
Components of the rotors, with the 0 index iterating over rotor, and the 1 index iterating over component.
s : int
Spin weight of the field to evaluate
ell : int
m : int
Values of (ell,m) to output
values : 1-d array of complex
Output array to contain values. Length must equal 0 dimension of `Rs`. Needed because numba cannot create
arrays at the moment.
Returns
-------
void
The input/output array `values` is modified in place.
"""
N = Rs.shape[0]
if (abs(m) > ell or abs(s) > ell):
for i in xrange(N):
values[i] = 0.0j
else:
constant = math.sqrt((2 * ell + 1) / (4 * np.pi))
ell_even = (ell % 2 == 0)
s_even = (s % 2 == 0)
rhoMin_a = max(0, m + s)
rhoMax_a = min(ell + m, ell + s)
coefficient_a = coeff(ell, m, -s)
if ((rhoMin_a + s) % 2 != 0):
coefficient_a *= -1
N1_a = ell + m + 1
N2_a = ell + s + 1
M_a = -s - m
rhoMin_b = max(0, -m + s)
rhoMax_b = min(ell - m, ell + s)
coefficient_b = coeff(ell, -m, -s)
if ((ell + rhoMin_b) % 2 != 0):
coefficient_b *= -1
N1_b = ell - m + 1
N2_b = ell + s + 1
M_b = -s + m
for i in xrange(N):
Ra = complex(Rs[i, 0], Rs[i, 3])
ra, phia = cmath.polar(Ra)
Rb = complex(Rs[i, 2], Rs[i, 1])
rb, phib = cmath.polar(Rb)
if ra <= epsilon:
if m != s:
values[i] = 0.0j
elif ell_even:
values[i] = constant * Rb**(-2 * s)
else:
values[i] = -constant * Rb**(-2 * s)
elif rb <= epsilon:
if m != -s:
values[i] = 0.0j
elif s_even:
values[i] = constant * Ra**(-2 * s)
else:
values[i] = -constant * Ra**(-2 * s)
elif ra < rb:
if (coefficient_b == 0.0j):
values[i] = 0.0j
else:
absRRatioSquared = -ra * ra / (rb * rb)
Prefactor = cmath.rect(
coefficient_b * rb**(2 * ell + s - m - 2 * rhoMin_b) *
ra**(-s + m + 2 * rhoMin_b),
phib * (-s - m) + phia * (-s + m))
Sum = 1.0
for rho in xrange(rhoMax_b, rhoMin_b, -1):
Sum *= absRRatioSquared * (
(N1_b - rho) * (N2_b - rho)) / (rho * (M_b + rho))
Sum += 1
values[i] = constant * Prefactor * Sum
else: # ra >= rb
if (coefficient_a == 0.0j):
values[i] = 0.0j
else:
absRRatioSquared = -rb * rb / (ra * ra)
Prefactor = cmath.rect(
coefficient_a * ra**(2 * ell + s + m - 2 * rhoMin_a) *
rb**(-s - m + 2 * rhoMin_a),
phia * (-s + m) + phib * (-s - m))
Sum = 1.0
for rho in xrange(rhoMax_a, rhoMin_a, -1):
Sum *= absRRatioSquared * (
(N1_a - rho) * (N2_a - rho)) / (rho * (M_a + rho))
Sum += 1
values[i] = constant * Prefactor * Sum
def _SWSH(Ra, Rb, s, indices, values):
"""Compute spin-weighted spherical harmonics from rotor components
This is the core function that does all the work in the
computation, but it is strict about its inputs, and does not check
them for validity -- though numba provides some degree of safety.
_SWSH(Ra, Rb, s, indices, values)
Parameters
----------
Ra : complex
Component `a` of the rotor
Rb : complex
Component `b` of the rotor
s : int
Spin weight of the field to evaluate
indices : 2-d array of int
Array of (ell,m) values to evaluate
values : 1-d array of complex
Output array to contain values. Length must equal first dimension of `indices`. Needed because numba cannot
create arrays at the moment.
Returns
-------
void
The input/output array `values` is modified in place.
"""
N = indices.shape[0]
# These constants are the recurring quantities in the computation
# of the matrix values, so we calculate them here just once
ra, phia = cmath.polar(Ra)
rb, phib = cmath.polar(Rb)
if ra <= epsilon:
for i in xrange(N):
ell, m = indices[i, 0:2]
if (m != s or abs(m) > ell or abs(s) > ell):
values[i] = 0.0j
else:
if (ell) % 2 == 0:
values[i] = math.sqrt(
(2 * ell + 1) / (4 * np.pi)) * Rb**(-2 * s)
else:
values[i] = -math.sqrt(
(2 * ell + 1) / (4 * np.pi)) * Rb**(-2 * s)
elif rb <= epsilon:
for i in xrange(N):
ell, m = indices[i, 0:2]
if (m != -s or abs(m) > ell or abs(s) > ell):
values[i] = 0.0j
else:
if (s) % 2 == 0:
values[i] = math.sqrt(
(2 * ell + 1) / (4 * np.pi)) * Ra**(-2 * s)
else:
values[i] = -math.sqrt(
(2 * ell + 1) / (4 * np.pi)) * Ra**(-2 * s)
elif ra < rb:
# We have to have these two versions (both this ra<rb branch,
# and ra>=rb below) to avoid overflows and underflows
absRRatioSquared = -ra * ra / (rb * rb)
for i in xrange(N):
ell, m = indices[i, 0:2]
if (abs(m) > ell or abs(s) > ell):
values[i] = 0.0j
else:
rhoMin = max(0, -m + s)
# Protect against overflow by decomposing Ra,Rb as
# abs,angle components and pulling out the factor of
# absRRatioSquared**rhoMin. Here, ra might be quite
# small, in which case ra**(-s+m) could be enormous
# when the exponent (-s+m) is very negative; adding
# 2*rhoMin to the exponent ensures that it is always
# positive, which protects from overflow. Meanwhile,
# underflow just goes to zero, which is fine since
# nothing else should be very large.
Prefactor = cmath.rect(
coeff(ell, -m, -s) * rb**(2 * ell + s - m - 2 * rhoMin) *
ra**(-s + m + 2 * rhoMin),
phib * (-s - m) + phia * (-s + m))
if (Prefactor == 0.0j):
values[i] = 0.0j
else:
if ((ell + rhoMin) % 2 != 0):
Prefactor *= -1
rhoMax = min(ell - m, ell + s)
N1 = ell - m + 1
N2 = ell + s + 1
M = -s + m
Sum = 1.0
for rho in xrange(rhoMax, rhoMin, -1):
Sum *= absRRatioSquared * ((N1 - rho) *
(N2 - rho)) / (rho *
(M + rho))
Sum += 1
# Sum = 0.0
# for rho in xrange(rhoMax, rhoMin-1, -1):
# Sum = ( binomial_coefficient(ell-m,rho) * binomial_coefficient(ell+m, ell-rho+s)
# + Sum * absRRatioSquared )
values[i] = math.sqrt(
(2 * ell + 1) / (4 * np.pi)) * Prefactor * Sum
else: # ra >= rb
# We have to have these two versions (both this ra>=rb branch,
# and ra<rb above) to avoid overflows and underflows
absRRatioSquared = -rb * rb / (ra * ra)
for i in xrange(N):
ell, m = indices[i, 0:2]
if (abs(m) > ell or abs(s) > ell):
values[i] = 0.0j
else:
rhoMin = max(0, m + s)
# Protect against overflow by decomposing Ra,Rb as
# abs,angle components and pulling out the factor of
# absRRatioSquared**rhoMin. Here, rb might be quite
# small, in which case rb**(-s-m) could be enormous
# when the exponent (-s-m) is very negative; adding
# 2*rhoMin to the exponent ensures that it is always
# positive, which protects from overflow. Meanwhile,
# underflow just goes to zero, which is fine since
# nothing else should be very large.
Prefactor = cmath.rect(
coeff(ell, m, -s) * ra**(2 * ell + s + m - 2 * rhoMin) *
rb**(-s - m + 2 * rhoMin),
phia * (-s + m) + phib * (-s - m))
if (Prefactor == 0.0j):
values[i] = 0.0j
else:
if ((rhoMin + s) % 2 != 0):
Prefactor *= -1
rhoMax = min(ell + m, ell + s)
N1 = ell + m + 1
N2 = ell + s + 1
M = -s - m
Sum = 1.0
for rho in xrange(rhoMax, rhoMin, -1):
Sum *= absRRatioSquared * ((N1 - rho) *
(N2 - rho)) / (rho *
(M + rho))
Sum += 1
# Sum = 0.0
# for rho in xrange(rhoMax, rhoMin-1, -1):
# Sum = ( binomial_coefficient(ell+m,rho) * binomial_coefficient(ell-m, ell-rho+s)
# + Sum * absRRatioSquared )
values[i] = math.sqrt(
(2 * ell + 1) / (4 * np.pi)) * Prefactor * Sum
def _Wigner_D_element(Ra, Rb, indices, elements):
"""Main work function for computing Wigner D matrix elements
This is the core function that does all the work in the
computation, but it is strict about its input, and does not check
them for validity. Note that the indices should be integers,
representing the (2*ell, 2*mp, 2*m) values, meaning that
(ell, mp, m) can be half-integer.
Input arguments
===============
_Wigner_D_element(Ra, Rb, indices, elements)
* Ra, Rb are the complex components of the rotor
* indices is an array of integer sets [2*ell, 2*mp, 2*m]
* elements is an array of complex with length equal to the first
dimension of indices
The `elements` variable is needed because numba cannot create
arrays at the moment, but this is modified in place.
"""
N = indices.shape[0]
# These constants are the recurring quantities in the computation
# of the matrix elements, so we calculate them here just once
ra, phia = cmath.polar(Ra)
rb, phib = cmath.polar(Rb)
if (ra <= epsilon):
for i in xrange(N):
twoell, twomp, twom = indices[i, 0:3]
if (twomp != -twom or abs(twomp) > twoell or abs(twom) > twoell):
elements[i] = 0.0j
else:
if (twoell - twom) % 4 == 0:
elements[i] = Rb**twom
else:
elements[i] = -Rb**twom
elif (rb <= epsilon):
for i in xrange(N):
twoell, twomp, twom = indices[i, 0:3]
if (twomp != twom or abs(twomp) > twoell or abs(twom) > twoell):
elements[i] = 0.0j
else:
elements[i] = Ra**twom
elif (ra < rb):
# We have to have these two versions (both this ra<rb branch,
# and ra>=rb below) to avoid overflows and underflows
absRRatioSquared = -ra * ra / (rb * rb)
for i in xrange(N):
twoell, twomp, twom = indices[i, 0:3]
if (abs(twomp) > twoell or abs(twom) > twoell):
elements[i] = 0.0j
else:
tworhoMin = max(0, -twomp - twom)
# Protect against overflow by decomposing Ra,Rb as
# abs,angle components and pulling out the factor of
# absRRatioSquared**rhoMin. Here, ra might be quite
# small, in which case ra**(m+mp) could be enormous
# when the exponent (m+mp) is very negative; adding
# 2*rhoMin to the exponent ensures that it is always
# positive, which protects from overflow. Meanwhile,
# underflow just goes to zero, which is fine since
# nothing else should be very large.
Prefactor = cmath.rect(
_coeff(twoell, -twomp, twom) *
rb**(twoell - (twom + twomp) / 2 - tworhoMin) *
ra**((twom + twomp) / 2 + tworhoMin),
phib * (twom - twomp) / 2 + phia * (twom + twomp) / 2)
if (Prefactor == 0.0j):
elements[i] = 0.0j
else:
if ((twoell - twom - tworhoMin) % 4 != 0):
Prefactor *= -1
tworhoMax = min(twoell - twomp, twoell - twom)
twoN1 = twoell - twomp + 2
twoN2 = twoell - twom + 2
twoM = twom + twomp
Sum = 1.0
for tworho in xrange(tworhoMax, tworhoMin, -2):
Sum *= absRRatioSquared * (
(twoN1 - tworho) *
(twoN2 - tworho)) / (tworho * (twoM + tworho))
Sum += 1
elements[i] = Prefactor * Sum
else: # ra >= rb
# We have to have these two versions (both this ra>=rb branch,
# and ra<rb above) to avoid overflows and underflows
absRRatioSquared = -rb * rb / (ra * ra)
for i in xrange(N):
twoell, twomp, twom = indices[i, 0:3]
if (abs(twomp) > twoell or abs(twom) > twoell):
elements[i] = 0.0j
else:
tworhoMin = max(0, twomp - twom)
# Protect against overflow by decomposing Ra,Rb as
# abs,angle components and pulling out the factor of
# absRRatioSquared**rhoMin. Here, rb might be quite
# small, in which case rb**(m-mp) could be enormous
# when the exponent (m-mp) is very negative; adding
# 2*rhoMin to the exponent ensures that it is always
# positive, which protects from overflow. Meanwhile,
# underflow just goes to zero, which is fine since
# nothing else should be very large.
Prefactor = cmath.rect(
_coeff(twoell, twomp, twom) *
ra**(twoell - twom / 2 + twomp / 2 - tworhoMin) *
rb**(twom / 2 - twomp / 2 + tworhoMin),
phia * (twom + twomp) / 2 + phib * (twom - twomp) / 2)
if (Prefactor == 0.0j):
elements[i] = 0.0j
else:
if (tworhoMin % 4 != 0):
Prefactor *= -1
tworhoMax = min(twoell + twomp, twoell - twom)
twoN1 = twoell + twomp + 2
twoN2 = twoell - twom + 2
twoM = twom - twomp
Sum = 1.0
for tworho in xrange(tworhoMax, tworhoMin, -2):
Sum *= absRRatioSquared * (
(twoN1 - tworho) *
(twoN2 - tworho)) / (tworho * (twoM + tworho))
Sum += 1
elements[i] = Prefactor * Sum