Commit 87e409c3 authored by Chris Kerr's avatar Chris Kerr

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
parents c3a88a37 b3624c21
Pipeline #15410 passed with stages
in 2 minutes and 9 seconds
#!/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],
)
! -*- 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
"""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
......@@ -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):
......
# 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)
# 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)
......@@ -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
......
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