Merge branch '38-do-not-require-fortran-compiler-f2py-for-wigner-3j-symbols' into 'master'

Resolve "Do not require Fortran compiler / f2py for Wigner 3j symbols"

Closes #38

See merge request !29

setup.py

 #!/usr/bin/env python
 
 # SPDX-FileCopyrightText: 2016-2017 European Molecular Biology Laboratory (EMBL)
-# SPDX-FileCopyrightText: 2018-2019 Christopher Kerr
+# SPDX-FileCopyrightText: 2018-2020 Christopher Kerr
 #
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-from numpy.distutils.core import setup as np_setup
-from numpy.distutils.extension import Extension as F2PyExtension
+from setuptools import setup
 
 
-wigner_ext = F2PyExtension(
-    name='wignerSymbols',
-    sources=[
-        'src/fortran/wignerSymbols.pyf',
-        'src/fortran/wignerSymbols.f',
-    ],
-)
-
-
-np_setup(
+setup(
     name='galumph',
     version='0.2.0',
     description='Calculate ALM using GPU acceleration',
@@ -43,6 +33,4 @@ np_setup(
     package_data={
         'galumph': ["cl-src/*.cl"],
     },
-    ext_package='galumph',
-    ext_modules=[wigner_ext],
 )

src/fortran/wignerSymbols.pyf

-!    -*- f90 -*-
-
-! SPDX-FileCopyrightText: 2018 Christopher Kerr
-!
-! SPDX-License-Identifier: CC0-1.0
-
-python module wignerSymbols
-    interface
-        subroutine shzmat(LMAX, DLMKP)
-            integer, intent(in) :: LMAX
-            double precision, intent(out), depend(LMAX), dimension((LMAX+1)*(LMAX+2)*(LMAX+2)*(LMAX+3)/12) :: DLMKP
-        end subroutine shzmat
-    end interface
-end python module wignerSymbols

src/python/gaunt.py

+"""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_2 + 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 magnitude and sign 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),
+    ))
+    sq_mag = 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)
+    sign = (-1) ** ((l_1 + l_2 - l_3) // 2)
+    if non_sqrt_part < 0:
+        sign *= -1
+    return sq_mag, sign
+
+
+def modified_gaunt(l_1, l_2, l_3, m, *, factorials=None):
+    """Modified Gaunt coefficient used in the Z translation matrix."""
+    sq_mag, sign = _modified_gaunt_squared(l_1, l_2, l_3, m, factorials=factorials)
+    return sign * math.sqrt(sq_mag)
+
+
+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

src/python/zshift.py

@@ -12,7 +12,7 @@ import pyopencl.array as cl_array
 from .almarray import AlmArray
 from .sphericalbessel import SphericalBessel
 from .util import ClProgram, LMAXMixin, check_array, with_some_queue
-from .wignerSymbols import shzmat
+from .gaunt import shzmat
 
 
 class PackedLMKPMixin(LMAXMixin):

test/test_wigner.py → test/test_formulae.py

-# SPDX-FileCopyrightText: 2018 Christopher Kerr
+"""Test various mathematical formulae used in GALUMPH.
+
+This uses symbolic algebra to verify various simplified formulae.
+"""
+
+# SPDX-FileCopyrightText: 2018-2020 Christopher Kerr
 #
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
-from hypothesis import assume, given, HealthCheck, settings
-import hypothesis.strategies as st
 import pytest
 
-from galumph import wignerSymbols
-
 sympy = pytest.importorskip('sympy')
-from sympy.physics.wigner import wigner_3j  # noqa: E402
-
-
-LMAX_TEST = 63
 
 
 def iLMKP(L, M, K, P, check=True):
@@ -37,25 +34,6 @@ def assert_sympy_equal(a, b):
     assert sympy.simplify(a - b) == 0
 
 
-@pytest.fixture(scope='module')
-def wigner_matrix_large():
-    """Fixture to calculate the Wigner matrix once and reuse it."""
-    mat = wignerSymbols.shzmat(LMAX_TEST)
-    LMAX1 = LMAX_TEST + 1
-    LMAX2 = LMAX_TEST + 2
-    LMAX3 = LMAX_TEST + 3
-    assert mat.shape == (LMAX1 * LMAX2**2 * LMAX3 / 12,)
-    return mat
-
-
-def sympy_matrix_element(L, M, K, P):
-    """Calculate an element of the dlmkp matrix using sympy."""
-    wigner0 = wigner_3j(L, P, K, 0, 0, 0)
-    wignerM = wigner_3j(L, P, K, -M, 0, M)
-    prefactor = (2*P+1) * sympy.sqrt((2*L+1) * (2*K+1))
-    return prefactor * wigner0 * wignerM
-
-
 def test_packed_sizes_offsets():
     """Check formulae for sizes of and offsets into packed dlmkp matrices."""
     def psize(K):
@@ -147,28 +125,3 @@ def test_packed_sizes_offsets():
 
     check_sizes()
     check_formulae()
-
-
-@settings(suppress_health_check=[HealthCheck.filter_too_much])
-@given(
-    L=st.integers(min_value=0, max_value=LMAX_TEST),
-    K=st.integers(min_value=0, max_value=LMAX_TEST),
-    M=st.integers(min_value=0, max_value=LMAX_TEST),
-    P=st.integers(min_value=0, max_value=2*LMAX_TEST),
-)
-def test_wigner_sympy_large(wigner_matrix_large, L, K, M, P):
-    """Test elements of the dlmkp matrix using sympy."""
-    assume(L >= M)
-    assume(M >= 0)
-    assume(K >= M)
-    assume((L+K) >= P)
-    assume(abs(L-K) <= P)
-    expected_symbolic = sympy_matrix_element(L, M, K, P)
-    if (L+K+P) % 2 != 0:
-        assert expected_symbolic == 0
-        return
-    expected = float(expected_symbolic)
-    if L >= K:
-        assert wigner_matrix_large[iLMKP(L, M, K, P)] == pytest.approx(expected)
-    else:
-        assert wigner_matrix_large[iLMKP(K, M, L, P)] == pytest.approx(expected)

test/test_gaunt.py

+# SPDX-FileCopyrightText: 2018-2020 Christopher Kerr
+#
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+from hypothesis import given
+import hypothesis.strategies as st
+import pytest
+
+from galumph import gaunt
+
+sympy = pytest.importorskip('sympy')
+import sympy.physics.wigner  # noqa: E402
+
+LMAX_TEST = 63
+
+
+def assert_sympy_equal(a, b):
+    """Assert that two sympy expressions are equal."""
+    assert sympy.simplify(a - b) == 0
+
+
+@st.composite
+def LMKP(draw):
+    L = draw(st.integers(min_value=0, max_value=LMAX_TEST))
+    M = draw(st.integers(min_value=0, max_value=L))
+    K = draw(st.integers(min_value=M, max_value=L))
+    iP = draw(st.integers(min_value=0, max_value=K))
+    P = L - K + 2 * iP
+    return (L, M, K, P)
+
+
+def sympy_matrix_element_W(L, M, K, P):
+    """Calculate an element of the dlmkp matrix using sympy."""
+    wigner0 = sympy.physics.wigner.wigner_3j(L, P, K, 0, 0, 0)
+    wignerM = sympy.physics.wigner.wigner_3j(L, P, K, -M, 0, M)
+    prefactor = (2*P+1) * sympy.sqrt((2*L+1) * (2*K+1))
+    return prefactor * wigner0 * wignerM
+
+
+def sympy_matrix_element_G(L, M, K, P):
+    """Calculate an element of the dlmkp matrix using sympy."""
+    standard_gaunt = sympy.physics.wigner.gaunt(L, K, P, M, -M, 0)
+    prefactor = sympy.sqrt((2*P+1) * 4 * sympy.pi)
+    return prefactor * standard_gaunt
+
+
+@given(lmkp=LMKP())
+def test_sympy_gaunt_vs_wigner(lmkp):
+    L, M, K, P = lmkp
+    assert_sympy_equal(
+        sympy_matrix_element_W(L, M, K, P),
+        sympy_matrix_element_G(L, M, K, P),
+    )
+
+
+@given(lmkp=LMKP())
+def test_python_vs_sympy_gaunt2(lmkp):
+    L, M, K, P = lmkp
+    sympy_result = sympy_matrix_element_G(L, M, K, P)
+    sympy_sign = sympy.sign(sympy_result)
+    python_result, sign = gaunt._modified_gaunt_squared(L, K, P, M)
+    if sympy_sign != 0:
+        assert sign == sympy_sign
+    assert_sympy_equal(
+        python_result,
+        sympy_result ** 2,
+    )
+
+
+@given(lmkp=LMKP())
+def test_python_vs_sympy_gaunt(lmkp):
+    L, M, K, P = lmkp
+    sympy_result = sympy_matrix_element_G(L, M, K, P)
+    sympy_float = float(sympy_result.n(10))
+    python_result = gaunt.modified_gaunt(L, K, P, M)
+    assert python_result == pytest.approx(sympy_float)

test/test_zshift.py

@@ -11,7 +11,7 @@ import pyopencl.array as cl_array
 import pytest
 
 from galumph.atomicscattering import AtomicScattering
-from galumph.wignerSymbols import shzmat
+from galumph.gaunt import shzmat
 from galumph.zshift import AlmShiftZ
 
 # Import test strategies from test_atomic_scattering