From 595a5fec60079b42816934a26cf4b3143f5ac1e0 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 11 Sep 2024 14:45:46 -0400 Subject: [PATCH 1/6] begin looking at removing pyshtools --- src/aspire/basis/basis_utils.py | 50 +++++++++++++++++++-------------- 1 file changed, 29 insertions(+), 21 deletions(-) diff --git a/src/aspire/basis/basis_utils.py b/src/aspire/basis/basis_utils.py index fe599e9fdc..3813c038ba 100644 --- a/src/aspire/basis/basis_utils.py +++ b/src/aspire/basis/basis_utils.py @@ -9,8 +9,7 @@ import numpy as np from numpy import diff, exp, log, pi from numpy.polynomial.legendre import leggauss -from pyshtools.expand import spharm_lm -from scipy.special import jn, jv, sph_harm +from scipy.special import jn, jv, sph_harm, factorial from aspire.utils import grid_2d, grid_3d @@ -157,6 +156,31 @@ def norm_assoc_legendre(j, m, x): return px +def _sph_harm(m, j, theta, phi): + """ + Compute spherical harmonics. + + :param m: Order |m| <= j + :param j: Harmonic degree, j>=0 + :param theta: longitude coordinate [0, 2*pi] + :param phi: latitude coordinate [0, pi] + """ + + if m < 0: + return (-1)**(m%2)*np.conj(_sph_harm(-m, j, phi, theta)) + + from scipy.special import lpmv + y = np.sqrt( ((2*j+1)/(4*np.pi)) * factorial(j-m)/factorial(j+m)) * lpmv(m, j, np.cos(phi)) * np.exp(1j*m*theta) # OKAY + # _y = norm_assoc_legendre(j, m, np.cos(phi)) * np.exp(1j*m*theta) / np.sqrt( ((2*j+1)/(4*np.pi)) * factorial(j-m)/factorial(j+m)) # not the right factor? + + # # also not the right factor + # k=2 + # if m == 0: + # k =1 + # _y = norm_assoc_legendre(j, m, np.cos(phi)) * np.exp(1j*m*theta) / np.sqrt( k*(2*j+1) * factorial(j-m)/factorial(j+m)) + + + return y def real_sph_harmonic(j, m, theta, phi): """ @@ -174,27 +198,11 @@ def real_sph_harmonic(j, m, theta, phi): # The `scipy` sph_harm implementation is much faster, # but incorrectly returns NaN for high orders. - # For higher order use `pyshtools`. - if j < 86: - y = sph_harm(abs_m, j, phi, theta) - else: - warnings.warn( - "Computing higher order spherical harmonics is slow." - " Consider using `FFBBasis3D` or decreasing volume size.", - stacklevel=1, - ) + y = _sph_harm(abs_m, j, phi, theta) + #y = sph_harm(abs_m, j, phi, theta) # scipy - y = spharm_lm( - j, - abs_m, - theta, - phi, - kind="complex", - degrees=False, - csphase=-1, - normalization="ortho", - ) + if m < 0: y = np.sqrt(2) * np.imag(y) elif m > 0: From 21f789a9f8b2493ae2d5f603919b80ee1efb98a8 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 11 Sep 2024 14:46:59 -0400 Subject: [PATCH 2/6] purge shtools --- .github/workflows/workflow.yml | 1 - environment-accelerate.yml | 1 - pyproject.toml | 1 - 3 files changed, 3 deletions(-) diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml index c41b221ec4..b2263a460b 100644 --- a/.github/workflows/workflow.yml +++ b/.github/workflows/workflow.yml @@ -244,7 +244,6 @@ jobs: run: | conda info conda list - conda install pyshtools # debug depends issues pip install -e ".[dev]" # install aspire pip freeze - name: Test diff --git a/environment-accelerate.yml b/environment-accelerate.yml index 38dd49813d..cfe9631f3c 100644 --- a/environment-accelerate.yml +++ b/environment-accelerate.yml @@ -7,7 +7,6 @@ channels: dependencies: - pip - python=3.8 - - pyshtools - numpy=1.24.1 - scipy=1.10.1 - scikit-learn diff --git a/pyproject.toml b/pyproject.toml index 51fc94dd8f..7dff9dddaf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,6 @@ dependencies = [ "pillow", "psutil", "pymanopt", - "pyshtools<=4.10.4", # 4.11.7 might have a packaging bug "PyWavelets", "ray >= 2.9.2", "scipy >= 1.10.0", From 24a12ba4d949ad4aa5d618a77b907a30de653333 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 11 Sep 2024 20:45:51 -0400 Subject: [PATCH 3/6] tox --- src/aspire/basis/basis_utils.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/aspire/basis/basis_utils.py b/src/aspire/basis/basis_utils.py index 3813c038ba..02297ab814 100644 --- a/src/aspire/basis/basis_utils.py +++ b/src/aspire/basis/basis_utils.py @@ -4,12 +4,11 @@ """ import logging -import warnings import numpy as np from numpy import diff, exp, log, pi from numpy.polynomial.legendre import leggauss -from scipy.special import jn, jv, sph_harm, factorial +from scipy.special import factorial, jn, jv from aspire.utils import grid_2d, grid_3d @@ -156,6 +155,7 @@ def norm_assoc_legendre(j, m, x): return px + def _sph_harm(m, j, theta, phi): """ Compute spherical harmonics. @@ -167,10 +167,15 @@ def _sph_harm(m, j, theta, phi): """ if m < 0: - return (-1)**(m%2)*np.conj(_sph_harm(-m, j, phi, theta)) + return (-1) ** (m % 2) * np.conj(_sph_harm(-m, j, phi, theta)) + + from scipy.special import lpmv - from scipy.special import lpmv - y = np.sqrt( ((2*j+1)/(4*np.pi)) * factorial(j-m)/factorial(j+m)) * lpmv(m, j, np.cos(phi)) * np.exp(1j*m*theta) # OKAY + y = ( + np.sqrt(((2 * j + 1) / (4 * np.pi)) * factorial(j - m) / factorial(j + m)) + * lpmv(m, j, np.cos(phi)) + * np.exp(1j * m * theta) + ) # OKAY # _y = norm_assoc_legendre(j, m, np.cos(phi)) * np.exp(1j*m*theta) / np.sqrt( ((2*j+1)/(4*np.pi)) * factorial(j-m)/factorial(j+m)) # not the right factor? # # also not the right factor @@ -179,9 +184,9 @@ def _sph_harm(m, j, theta, phi): # k =1 # _y = norm_assoc_legendre(j, m, np.cos(phi)) * np.exp(1j*m*theta) / np.sqrt( k*(2*j+1) * factorial(j-m)/factorial(j+m)) - return y + def real_sph_harmonic(j, m, theta, phi): """ Evaluate a real spherical harmonic @@ -199,10 +204,9 @@ def real_sph_harmonic(j, m, theta, phi): # The `scipy` sph_harm implementation is much faster, # but incorrectly returns NaN for high orders. y = _sph_harm(abs_m, j, phi, theta) - #y = sph_harm(abs_m, j, phi, theta) # scipy - + # from scipy.special import sph_harm + # y = sph_harm(abs_m, j, phi, theta) # scipy - if m < 0: y = np.sqrt(2) * np.imag(y) elif m > 0: From 51ece9fdce83df2a0660ea195dd51eb8aff5e834 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 12 Sep 2024 09:01:44 -0400 Subject: [PATCH 4/6] implement sph_harm using recurance form of norm assoc legendre --- src/aspire/basis/basis_utils.py | 43 +++++++++++++----------------- tests/test_basis_utils.py | 46 +++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 25 deletions(-) diff --git a/src/aspire/basis/basis_utils.py b/src/aspire/basis/basis_utils.py index 02297ab814..e629d4d456 100644 --- a/src/aspire/basis/basis_utils.py +++ b/src/aspire/basis/basis_utils.py @@ -8,7 +8,7 @@ import numpy as np from numpy import diff, exp, log, pi from numpy.polynomial.legendre import leggauss -from scipy.special import factorial, jn, jv +from scipy.special import jn, jv from aspire.utils import grid_2d, grid_3d @@ -156,33 +156,29 @@ def norm_assoc_legendre(j, m, x): return px -def _sph_harm(m, j, theta, phi): +def sph_harm(j, m, theta, phi): """ Compute spherical harmonics. + Note call signature convention may be different from other packages. + :param m: Order |m| <= j :param j: Harmonic degree, j>=0 - :param theta: longitude coordinate [0, 2*pi] - :param phi: latitude coordinate [0, pi] + :param theta: latitude coordinate [0, pi] + :param phi: longitude coordinate [0, 2*pi] + :return: Complex array of evaluated spherical harmonics. """ - if m < 0: - return (-1) ** (m % 2) * np.conj(_sph_harm(-m, j, phi, theta)) - - from scipy.special import lpmv - + # Compute sph_harm for positive `abs(m)` y = ( - np.sqrt(((2 * j + 1) / (4 * np.pi)) * factorial(j - m) / factorial(j + m)) - * lpmv(m, j, np.cos(phi)) - * np.exp(1j * m * theta) - ) # OKAY - # _y = norm_assoc_legendre(j, m, np.cos(phi)) * np.exp(1j*m*theta) / np.sqrt( ((2*j+1)/(4*np.pi)) * factorial(j-m)/factorial(j+m)) # not the right factor? - - # # also not the right factor - # k=2 - # if m == 0: - # k =1 - # _y = norm_assoc_legendre(j, m, np.cos(phi)) * np.exp(1j*m*theta) / np.sqrt( k*(2*j+1) * factorial(j-m)/factorial(j+m)) + norm_assoc_legendre(j, abs(m), np.cos(theta)) + * np.exp(1j * abs(m) * phi) + * np.sqrt(0.5 / np.pi) + ) + + # Use identity for negative `m` + if m < 0: + (-1) ** (m % 2) * np.conj(y) return y @@ -201,11 +197,8 @@ def real_sph_harmonic(j, m, theta, phi): """ abs_m = abs(m) - # The `scipy` sph_harm implementation is much faster, - # but incorrectly returns NaN for high orders. - y = _sph_harm(abs_m, j, phi, theta) - # from scipy.special import sph_harm - # y = sph_harm(abs_m, j, phi, theta) # scipy + # Note the calling convention here may not match other `sph_harm` packages + y = sph_harm(j, abs_m, theta, phi) if m < 0: y = np.sqrt(2) * np.imag(y) diff --git a/tests/test_basis_utils.py b/tests/test_basis_utils.py index 2d6a3efdb2..6f475f89fe 100644 --- a/tests/test_basis_utils.py +++ b/tests/test_basis_utils.py @@ -1,6 +1,7 @@ from unittest import TestCase import numpy as np +from scipy.special import sph_harm as sp_sph_harm from aspire.basis.basis_utils import ( all_besselj_zeros, @@ -9,10 +10,55 @@ norm_assoc_legendre, real_sph_harmonic, sph_bessel, + sph_harm, unique_coords_nd, ) +def test_sph_harm_low_order(): + """ + Test the `sph_harm` implementation matches `scipy` at lower orders. + """ + m = 3 + j = 5 + x = np.linspace(0, np.pi, 42) + y = np.linspace(0, 2 * np.pi, 42) + + ref = sp_sph_harm(m, j, y, x) # Note calling convention is different + np.testing.assert_allclose(sph_harm(j, m, x, y), ref) + + +def test_sph_harm_high_order(): + """ + Test we remain finite at higher orders where `scipy.special.sph_harm` overflows. + """ + m = 87 + j = 87 + x = 0.12345 + y = 0.56789 + + # If scipy fixed their implementation for higher orders in the future, + # this check will we may wish to take it. + ref = sp_sph_harm(m, j, y, x) # Note calling convention is different + assert not np.isfinite(ref) + + # Can manually check against pyshtools, + # but we are avoiding that package dependency. + # y = spharm_lm( + # j, + # abs_m, + # theta, + # phi, + # kind="complex", + # degrees=False, + # csphase=-1, + # normalization="ortho", + # ) + + # Check we are finite. + assert np.isfinite(sph_harm(j, m, x, y)) + + class BesselTestCase(TestCase): def setUp(self): pass From 521f5cbbe9e3c61cbc3234ba3f9419c70269d8c6 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 12 Sep 2024 09:41:36 -0400 Subject: [PATCH 5/6] self review cleanup --- tests/test_basis_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_basis_utils.py b/tests/test_basis_utils.py index 6f475f89fe..e0693fce7a 100644 --- a/tests/test_basis_utils.py +++ b/tests/test_basis_utils.py @@ -38,12 +38,13 @@ def test_sph_harm_high_order(): y = 0.56789 # If scipy fixed their implementation for higher orders in the future, - # this check will we may wish to take it. + # this check should fail and we can reconsider that package. ref = sp_sph_harm(m, j, y, x) # Note calling convention is different assert not np.isfinite(ref) # Can manually check against pyshtools, # but we are avoiding that package dependency. + # Leaving this here intentionally for future developers. # y = spharm_lm( # j, # abs_m, From 2199c3beb5216a0d0d62f56157e4d82a140e331b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 12 Sep 2024 10:27:39 -0400 Subject: [PATCH 6/6] Fixup -m sph_harm case and add test. --- src/aspire/basis/basis_utils.py | 2 +- tests/test_basis_utils.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/aspire/basis/basis_utils.py b/src/aspire/basis/basis_utils.py index e629d4d456..c32e73fbf3 100644 --- a/src/aspire/basis/basis_utils.py +++ b/src/aspire/basis/basis_utils.py @@ -178,7 +178,7 @@ def sph_harm(j, m, theta, phi): # Use identity for negative `m` if m < 0: - (-1) ** (m % 2) * np.conj(y) + y = (-1) ** (m % 2) * np.conj(y) return y diff --git a/tests/test_basis_utils.py b/tests/test_basis_utils.py index e0693fce7a..2e416a5dee 100644 --- a/tests/test_basis_utils.py +++ b/tests/test_basis_utils.py @@ -27,6 +27,11 @@ def test_sph_harm_low_order(): ref = sp_sph_harm(m, j, y, x) # Note calling convention is different np.testing.assert_allclose(sph_harm(j, m, x, y), ref) + # negative m + m *= -1 + ref = sp_sph_harm(m, j, y, x) # Note calling convention is different + np.testing.assert_allclose(sph_harm(j, m, x, y), ref) + def test_sph_harm_high_order(): """