Verified Commit cd5df099 authored by Chris Kerr's avatar Chris Kerr

Pure Python implementation of the Z translation coefficients

These numbers are the Gaunt coefficients, multiplied by a factor of
sqrt((2 * l_3 + 1) / (4 * pi))
parent 7ce1d26b
"""Pure Python implementation of Gaunt coefficients.
Gaunt coefficients give the spatial integral of the product of three
spherical harmonic functions. In GALUMPH, they are used in the ALM
translation function.
The code here is based on Formula 4.2 of:
J. Rasch and A. C. H. Yu, 'Efficient Storage Scheme for
Pre-calculated Wigner 3j, 6j and Gaunt Coefficients', SIAM
J. Sci. Comput. Volume 25, Issue 4, pp. 1416-1428 (2003)
https://doi.org/10.1137/S1064827503422932
However, the formula in that article has some errors, which are corrected here.
For a translation along the Z axis, we only need values with:
m_3 == 0, m_1 == -m_2
This allows simplifying the formula and referring to m_1 as plain `m`.
"""
# SPDX-FileCopyrightText: 2020 Christopher Kerr
#
# SPDX-License-Identifier: LGPL-3.0-or-later
__copyright__ = "Christopher Kerr"
__license__ = "LGPLv3+"
from fractions import Fraction
import functools
import math
import operator
def prod(iterable, *, start=1):
"""Product of a sequence of numbers.
Reimplementation of math.prod, which is only available for Python >= 3.8.
"""
return functools.reduce(operator.mul, iterable, start)
def _calculate_factorials(N):
"""Return a tuple containing the factorials of all numbers from 0 to N.
Since Python integers are arbitrary-precision, the values are exact.
"""
assert N >= 0
# Special case for 0
fac_i = 1
fac_list = [fac_i]
for i in range(1, N+1):
fac_i *= i
fac_list.append(fac_i)
return tuple(fac_list)
def delta_squared(l_1, l_2, l_3, *, factorials=None):
"""Square of the Delta function defined in Formula 2.3 of Rasch2003."""
if factorials is None:
factorials = _calculate_factorials(l_1 + l_2 + l_3 + 1)
return Fraction(
prod((
factorials[l_1 + l_2 - l_3],
factorials[l_2 + l_3 - l_1],
factorials[l_3 + l_1 - l_2],
)),
factorials[l_1 + l_2 + l_3 + 1],
)
def _gaunt_ksum(l_1, l_2, l_3, m, *, factorials):
"""Part of the Gaunt coefficient function involving a sum over k."""
k_min = max(l_1 - m - l_3, l_2 - m - l_3, 0)
k_max = min(l_1 - m, l_2 - m, l_1 + l_2 - l_3)
def kpart(k):
return Fraction(
-1 ** k,
prod((
factorials[k],
factorials[k + l_3 + m - l_1],
factorials[k + l_3 + m - l_2],
factorials[l_1 + l_2 - l_3 - k],
factorials[l_1 - m - k],
factorials[l_2 - m - k],
)),
)
return sum(kpart(k) for k in range(k_min, k_max + 1))
def _gaunt_sqrt_arg(l_1, l_2, l_3, m, *, factorials):
"""Calculate the part of the formula that needs to be square-rooted."""
return prod((
(2 * l_1 + 1),
(2 * l_1 + 1),
(2 * l_3 + 1),
factorials[l_1 + m],
factorials[l_1 - m],
factorials[l_2 + m],
factorials[l_2 - m],
factorials[l_3 + 0],
factorials[l_3 - 0],
))
def _gaunt_big_L_part(l_1, l_2, l_3, *, factorials):
"""The part of the Gaunt coefficient formula involving the capital L."""
assert (l_1 + l_2 + l_3) % 2 == 0
L = (l_1 + l_2 + l_3) // 2
return Fraction(
factorials[L],
prod((
factorials[L - l_1],
factorials[L - l_2],
factorials[L - l_3],
)),
)
def _modified_gaunt_squared(l_1, l_2, l_3, m, *, factorials=None):
"""Square of the modified Gaunt coefficient.
Here we are calculating a modified Gaunt coefficient which differs from
the standard coefficient by a factor of sqrt((2 * l_3 + 1) / 4 * pi). That
modification is applied here.
"""
if factorials is None:
factorials = _calculate_factorials(l_1 + l_2 + l_3 + 1)
non_sqrt_part = prod((
delta_squared(l_1, l_2, l_3, factorials=factorials),
_gaunt_ksum(l_1, l_2, l_3, m, factorials=factorials),
_gaunt_big_L_part(l_1, l_2, l_3, factorials=factorials),
))
return float(prod((
_gaunt_sqrt_arg(l_1, l_2, l_3, m, factorials=factorials),
non_sqrt_part * non_sqrt_part,
(2 * l_3 + 1), # For standard Gaunt, omit this
))) # For standard Gaunt, divide by (4 * math.pi)
def modified_gaunt(l_1, l_2, l_3, m, *, factorials=None):
"""Modified Gaunt coefficient used in the Z translation matrix."""
return math.sqrt(_modified_gaunt_squared(l_1, l_2, l_3, m, factorials=factorials))
def shzmat(LMAX):
"""Calculate the array of Gaunt coefficients used for ALM Z translation."""
factorials = _calculate_factorials(4 * LMAX + 1)
dlmkp = list()
for L in range(LMAX + 1):
for M in range(L + 1):
for K in range(M, L + 1):
for P in range(L - K, L + K + 1, 2):
dlmkp.append(modified_gaunt(L, K, P, M, factorials=factorials))
return dlmkp
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