From 99c75adf10ef1b914ec0b6cc1bb62e46d0943c58 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 28 Apr 2026 15:58:13 +0100 Subject: [PATCH 01/66] Add Clenshaw-Curtis quadrature rule implementation --- pyproject.toml | 1 + s2fft/utils/quadrature.py | 58 ++++++++++++++++++++++++++++++++++- s2fft/utils/quadrature_jax.py | 44 +++++++++++++++++++++++++- 3 files changed, 101 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a7d640b5f..62f63c2f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,6 +5,7 @@ requires = [ "scikit-build-core>=0.4.3", "nanobind>=1.3.2", "jax >= 0.5.0", + "array-api-extra >= 0.10.1", ] build-backend = "scikit_build_core.build" diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index 5b2e15783..c2f48d420 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -1,3 +1,6 @@ +from types import ModuleType + +import array_api_extra as xpx import numpy as np import numpy.fft as fft @@ -18,7 +21,7 @@ def quad_weights_transform( L (int): Harmonic band-limit. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mwss", "dh", "gl", "healpix}. Defaults to "mwss". + {"mwss", "dh", "gl", "healpix", "cc"}. Defaults to "mwss". spin (int, optional): Harmonic spin. Defaults to 0. @@ -46,6 +49,9 @@ def quad_weights_transform( elif sampling.lower() == "healpix": return quad_weights_hp(nside) + elif sampling.lower == "cc": + return quad_weights_cc(L) + else: raise ValueError(f"Sampling scheme sampling={sampling} not supported") @@ -318,3 +324,53 @@ def mw_weights(m: int) -> float: else: return 0 + + +def quad_weights_cc(L: int, xp: ModuleType = np) -> np.ndarray: + r""" + Compute Clenshaw-Curtis quadrature weights for :math:`\theta` and :math:`\phi` integration. + + Args: + L (int): Harmonic band-limit. + xp (module): Array namespace to use to operations. Defaults to NumPy. + + Returns: + np.ndarray: Weights computed for each :math:`\theta` (weights are identical + as :math:`\phi` varies for given :math:`\theta`). + + """ + return quad_weights_cc_theta_only(L, xp) * xp.pi / L + + +def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: + r""" + Compute Clenshaw-curtis quadrature weights for :math:`\theta` integration (only). + + Args: + L (int): Harmonic band-limit. + xp (module): Array namespace to use to operations. Defaults to NumPy. + + Returns: + np.ndarray: Weights computed for each :math:`\theta`. + + References: + Waldvogel, J. (2006). Fast construction of the Fejer and Clenshaw-Curtis quadrature rules. + BIT Numerical Mathematics, 46(1), 195–202. https://doi.org/10.1007/s10543-006-0045-4. + + """ + if L < 1: + raise ValueError(f"Bandlimit must be at least 1: L = {L}") + + if L == 1: + return xp.array([1.0, 1.0]) + + t = xp.arange(1, L, 2) + w0 = -2 / (t * (t - 2)) + w1 = xp.array([-2 / (L - 1)]) if L % 2 == 0 else xp.array([-1 / (L - 2)] * 2) + w = xp.concatenate([w0, w1, w0[-1:0:-1]]) + g = -xp.ones(L) + g = xpx.at(g)[L // 2].add(L) + g = xpx.at(g)[L - L // 2].add(L) + g /= L**2 - 1 + (L % 2) + weights = xp.fft.ifft(w + g) + return xp.concatenate([weights.real, weights.real[:1]]) diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py index bae0bc3b9..941af6159 100644 --- a/s2fft/utils/quadrature_jax.py +++ b/s2fft/utils/quadrature_jax.py @@ -5,6 +5,10 @@ from jax import jit as _jit from s2fft.sampling import s2_samples as samples +from s2fft.utils.quadrature import quad_weights_cc as _quad_weights_cc +from s2fft.utils.quadrature import ( + quad_weights_cc_theta_only as _quad_weights_cc_theta_only, +) @_partial(_jit, static_argnums=(0, 1, 2)) @@ -23,7 +27,7 @@ def quad_weights_transform( L (int): Harmonic band-limit. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mwss", "dh", "gl", "healpix}. Defaults to "mwss". + {"mwss", "dh", "gl", "healpix", "cc"}. Defaults to "mwss". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -95,6 +99,9 @@ def quad_weights(L: int = None, sampling: str = "mw", nside: int = None) -> jnp. elif sampling.lower() == "healpix": return quad_weights_hp(nside) + elif sampling.lower() == "cc": + return quad_weights_cc(L) + else: raise ValueError(f"Sampling scheme sampling={sampling} not implemented") @@ -341,3 +348,38 @@ def mw_weights(m: int) -> float: ), m, ) + + +@_partial(_jit, static_argnums=(0)) +def quad_weights_cc(L: int) -> jnp.ndarray: + r""" + Compute Clenshaw-Curtis quadrature weights for :math:`\theta` and :math:`\phi` integration. + + Args: + L (int): Harmonic band-limit. + + Returns: + np.ndarray: Weights computed for each :math:`\theta` (weights are identical + as :math:`\phi` varies for given :math:`\theta`). + + """ + return _quad_weights_cc(L, xp=jnp) + + +@_partial(_jit, static_argnums=(0)) +def quad_weights_cc_theta_only(L: int) -> jnp.ndarray: + r""" + Compute Clenshaw-curtis quadrature weights for :math:`\theta` integration (only). + + Args: + L (int): Harmonic band-limit. + + Returns: + jnp.ndarray: Weights computed for each :math:`\theta`. + + References: + Waldvogel, J. (2006). Fast construction of the Fejer and Clenshaw-Curtis quadrature rules. + BIT Numerical Mathematics, 46(1), 195–202. https://doi.org/10.1007/s10543-006-0045-4. + + """ + return _quad_weights_cc_theta_only(L, xp=jnp) From ae28cc1f473c56a8b335802ec84365a132a0794d Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 30 Apr 2026 09:03:55 +0100 Subject: [PATCH 02/66] Change to inverse RFFT based implementation --- s2fft/utils/quadrature.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index c2f48d420..f0fa77862 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -332,7 +332,7 @@ def quad_weights_cc(L: int, xp: ModuleType = np) -> np.ndarray: Args: L (int): Harmonic band-limit. - xp (module): Array namespace to use to operations. Defaults to NumPy. + xp (module): Array namespace to use for operations. Defaults to NumPy. Returns: np.ndarray: Weights computed for each :math:`\theta` (weights are identical @@ -348,7 +348,7 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: Args: L (int): Harmonic band-limit. - xp (module): Array namespace to use to operations. Defaults to NumPy. + xp (module): Array namespace to use for operations. Defaults to NumPy. Returns: np.ndarray: Weights computed for each :math:`\theta`. @@ -366,11 +366,10 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: t = xp.arange(1, L, 2) w0 = -2 / (t * (t - 2)) - w1 = xp.array([-2 / (L - 1)]) if L % 2 == 0 else xp.array([-1 / (L - 2)] * 2) - w = xp.concatenate([w0, w1, w0[-1:0:-1]]) - g = -xp.ones(L) + w1 = xp.array([-2 / (L - 1)]) if L % 2 == 0 else xp.array([-1 / (L - 2)]) + w = xp.concatenate([w0, w1]) + g = -xp.ones(L // 2 + 1) g = xpx.at(g)[L // 2].add(L) - g = xpx.at(g)[L - L // 2].add(L) g /= L**2 - 1 + (L % 2) - weights = xp.fft.ifft(w + g) - return xp.concatenate([weights.real, weights.real[:1]]) + weights = xp.fft.irfft(w + g, n=L) + return xp.concatenate([weights, weights[:1]]) From f18db3535984198be02ec17ecb858c46e7d735c7 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 30 Apr 2026 14:07:56 +0100 Subject: [PATCH 03/66] Correct CC weights expression for even L --- s2fft/utils/quadrature.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index f0fa77862..e99401a26 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -369,7 +369,7 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: w1 = xp.array([-2 / (L - 1)]) if L % 2 == 0 else xp.array([-1 / (L - 2)]) w = xp.concatenate([w0, w1]) g = -xp.ones(L // 2 + 1) - g = xpx.at(g)[L // 2].add(L) + g = xpx.at(g)[L // 2].add(2 * L if L % 2 == 0 else L) g /= L**2 - 1 + (L % 2) weights = xp.fft.irfft(w + g, n=L) return xp.concatenate([weights, weights[:1]]) From 2e47ee4bd3ec7ccbbc1c3ed08f1d39a4978bd1a4 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 19 May 2026 14:01:37 +0100 Subject: [PATCH 04/66] Add Fejer-2 weights and refactor CC implementation to reuse --- s2fft/utils/quadrature.py | 73 ++++++++++++++++++++++++++++++++--- s2fft/utils/quadrature_jax.py | 56 ++++++++++++++++++++++++++- 2 files changed, 121 insertions(+), 8 deletions(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index e99401a26..a7f25cc5c 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -326,6 +326,29 @@ def mw_weights(m: int) -> float: return 0 +def _fejer_second_rule_rfft_weights(L: int, xp: ModuleType = np): + r""" + Compute real-valued FFT of weights for Fejer second quadrature rule. + + Args: + L (int): Harmonic band-limit. + xp (module): Array namespace to use for operations. Defaults to NumPy. + + Returns: + np.ndarray: Real-valued FFT of weights computed for each :math:`\theta` + (weights are identical as :math:`\phi` varies for given :math:`\theta`). + + References: + Waldvogel, J. (2006). Fast construction of the Fejer and Clenshaw-Curtis quadrature rules. + BIT Numerical Mathematics, 46(1), 195–202. https://doi.org/10.1007/s10543-006-0045-4. + + """ + t = xp.arange(1, L, 2) + w0 = -2 / (t * (t - 2)) + w1 = xp.array([-2 / (L - 1)]) if L % 2 == 0 else xp.array([-1 / (L - 2)]) + return xp.concatenate([w0, w1]) + + def quad_weights_cc(L: int, xp: ModuleType = np) -> np.ndarray: r""" Compute Clenshaw-Curtis quadrature weights for :math:`\theta` and :math:`\phi` integration. @@ -339,12 +362,12 @@ def quad_weights_cc(L: int, xp: ModuleType = np) -> np.ndarray: as :math:`\phi` varies for given :math:`\theta`). """ - return quad_weights_cc_theta_only(L, xp) * xp.pi / L + return quad_weights_cc_theta_only(L, xp) * 2 * xp.pi / samples.nphi_equiang(L, "cc") def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: r""" - Compute Clenshaw-curtis quadrature weights for :math:`\theta` integration (only). + Compute Clenshaw-Curtis quadrature weights for :math:`\theta` integration (only). Args: L (int): Harmonic band-limit. @@ -364,12 +387,50 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: if L == 1: return xp.array([1.0, 1.0]) - t = xp.arange(1, L, 2) - w0 = -2 / (t * (t - 2)) - w1 = xp.array([-2 / (L - 1)]) if L % 2 == 0 else xp.array([-1 / (L - 2)]) - w = xp.concatenate([w0, w1]) + w = _fejer_second_rule_rfft_weights(L, xp) g = -xp.ones(L // 2 + 1) g = xpx.at(g)[L // 2].add(2 * L if L % 2 == 0 else L) g /= L**2 - 1 + (L % 2) weights = xp.fft.irfft(w + g, n=L) return xp.concatenate([weights, weights[:1]]) + + +def quad_weights_f2(L: int, xp: ModuleType = np) -> np.ndarray: + r""" + Compute Fejér's second rule quadrature weights for :math:`\theta` and :math:`\phi` integration. + + Args: + L (int): Harmonic band-limit. + xp (module): Array namespace to use for operations. Defaults to NumPy. + + Returns: + np.ndarray: Weights computed for each :math:`\theta` (weights are identical + as :math:`\phi` varies for given :math:`\theta`). + + """ + return quad_weights_f2_theta_only(L, xp) * 2 * xp.pi / samples.nphi_equiang(L, "f2") + + +def quad_weights_f2_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: + r""" + Compute Fejér's second rule quadrature weights for :math:`\theta` integration (only). + + Args: + L (int): Harmonic band-limit. + xp (module): Array namespace to use for operations. Defaults to NumPy. + + Returns: + np.ndarray: Weights computed for each :math:`\theta`. + + References: + Waldvogel, J. (2006). Fast construction of the Fejer and Clenshaw-Curtis quadrature rules. + BIT Numerical Mathematics, 46(1), 195–202. https://doi.org/10.1007/s10543-006-0045-4. + + """ + if L < 1: + raise ValueError(f"Bandlimit must be at least 1: L = {L}") + + w = _fejer_second_rule_rfft_weights(L, xp) + weights = xp.fft.irfft(w, n=L) + # Weights are zero at poles which we assume are excluded from nodes + return weights[1:] diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py index 941af6159..1c4d7761f 100644 --- a/s2fft/utils/quadrature_jax.py +++ b/s2fft/utils/quadrature_jax.py @@ -5,10 +5,18 @@ from jax import jit as _jit from s2fft.sampling import s2_samples as samples -from s2fft.utils.quadrature import quad_weights_cc as _quad_weights_cc +from s2fft.utils.quadrature import ( + quad_weights_cc as _quad_weights_cc, +) from s2fft.utils.quadrature import ( quad_weights_cc_theta_only as _quad_weights_cc_theta_only, ) +from s2fft.utils.quadrature import ( + quad_weights_f2 as _quad_weights_f2, +) +from s2fft.utils.quadrature import ( + quad_weights_f2_theta_only as _quad_weights_f2_theta_only, +) @_partial(_jit, static_argnums=(0, 1, 2)) @@ -53,6 +61,12 @@ def quad_weights_transform( elif sampling.lower() == "healpix": return quad_weights_hp(nside) + elif sampling.lower() == "cc": + return quad_weights_cc(L) + + elif sampling.lower() == "f2": + return quad_weights_f2(L) + else: raise ValueError(f"Sampling scheme sampling={sampling} not supported") @@ -102,6 +116,9 @@ def quad_weights(L: int = None, sampling: str = "mw", nside: int = None) -> jnp. elif sampling.lower() == "cc": return quad_weights_cc(L) + elif sampling.lower() == "f2": + return quad_weights_f2(L) + else: raise ValueError(f"Sampling scheme sampling={sampling} not implemented") @@ -369,7 +386,7 @@ def quad_weights_cc(L: int) -> jnp.ndarray: @_partial(_jit, static_argnums=(0)) def quad_weights_cc_theta_only(L: int) -> jnp.ndarray: r""" - Compute Clenshaw-curtis quadrature weights for :math:`\theta` integration (only). + Compute Clenshaw-Curtis quadrature weights for :math:`\theta` integration (only). Args: L (int): Harmonic band-limit. @@ -383,3 +400,38 @@ def quad_weights_cc_theta_only(L: int) -> jnp.ndarray: """ return _quad_weights_cc_theta_only(L, xp=jnp) + + +@_partial(_jit, static_argnums=(0)) +def quad_weights_f2(L: int) -> jnp.ndarray: + r""" + Compute Fejér's second rule quadrature weights for :math:`\theta` and :math:`\phi` integration. + + Args: + L (int): Harmonic band-limit. + + Returns: + np.ndarray: Weights computed for each :math:`\theta` (weights are identical + as :math:`\phi` varies for given :math:`\theta`). + + """ + return _quad_weights_f2(L, xp=jnp) + + +@_partial(_jit, static_argnums=(0)) +def quad_weights_f2_theta_only(L: int) -> jnp.ndarray: + r""" + Compute Fejér's second rule quadrature weights for :math:`\theta` integration (only). + + Args: + L (int): Harmonic band-limit. + + Returns: + jnp.ndarray: Weights computed for each :math:`\theta`. + + References: + Waldvogel, J. (2006). Fast construction of the Fejer and Clenshaw-Curtis quadrature rules. + BIT Numerical Mathematics, 46(1), 195–202. https://doi.org/10.1007/s10543-006-0045-4. + + """ + return _quad_weights_f2_theta_only(L, xp=jnp) From efa2fb44d2688fa58fda07cf0e2f3ff134a4143b Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 19 May 2026 14:02:22 +0100 Subject: [PATCH 05/66] Add dimensions for CC and F2 sampling grids --- s2fft/sampling/s2_samples.py | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 6e8c8dc3e..ababd7b23 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -31,12 +31,18 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: f"Sampling scheme sampling={sampling} with L={L} not supported" ) - if sampling.lower() in ["mw", "gl"]: + if sampling.lower() in ("mw", "gl"): return L + elif sampling.lower() == "f2": + return L - 1 + elif sampling.lower() == "mwss": return L + 1 + elif sampling.lower() == "cc": + return L + 1 + elif sampling.lower() == "dh": return 2 * L @@ -94,7 +100,7 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int: L (int): Harmonic band-limit. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "cc", "f2"}. Defaults to "mw". Raises: ValueError: HEALPix sampling scheme. @@ -105,15 +111,12 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int: int: Number of :math:`\phi` samples. """ - if sampling.lower() in ["mw", "gl"]: + if sampling.lower() in ("mw", "dh", "gl"): return 2 * L - 1 - elif sampling.lower() == "mwss": + elif sampling.lower() in ("mwss", "cc", "f2"): return 2 * L - elif sampling.lower() == "dh": - return 2 * L - 1 - elif sampling.lower() == "healpix": raise ValueError(f"Sampling scheme sampling={sampling} not supported") @@ -131,7 +134,7 @@ def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> tuple[int, int L (int): Harmonic band-limit. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. @@ -143,11 +146,11 @@ def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> tuple[int, int that here "healpix" defaults to :math:`2L = 4nside` phi samples for ftm. """ - if sampling.lower() in ["mwss", "healpix"]: + if sampling.lower() == "healpix": return ntheta(L, sampling, nside), 2 * L - elif sampling.lower() in ["mw", "dh", "gl"]: - return ntheta(L, sampling, nside), 2 * L - 1 + elif sampling.lower() in ["mw", "dh", "gl", "mwss", "cc", "f2"]: + return ntheta(L, sampling, nside), nphi_equiang(L, sampling) else: raise ValueError(f"Sampling scheme sampling={sampling} not supported") @@ -209,7 +212,7 @@ def thetas(L: int = None, sampling: str = "mw", nside: int = None) -> np.ndarray Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -239,7 +242,7 @@ def t2theta( Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -260,10 +263,10 @@ def t2theta( f"Sampling scheme sampling={sampling} with L={L} not supported" ) - if sampling.lower() == "mw": + if sampling.lower() in ("mw", "f2"): return (2 * t + 1) * np.pi / (2 * L - 1) - elif sampling.lower() == "mwss": + elif sampling.lower() in ("mwss", "cc"): return 2 * t * np.pi / (2 * L) elif sampling.lower() == "dh": From 8f2c2d97cee9cae1b1fe8567facfc40e7a6cbe32 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 19 May 2026 15:47:44 +0100 Subject: [PATCH 06/66] Add CC and F2 sampling methods to numpy quadrature wrapper --- s2fft/utils/quadrature.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index a7f25cc5c..c13a8f74b 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -49,9 +49,12 @@ def quad_weights_transform( elif sampling.lower() == "healpix": return quad_weights_hp(nside) - elif sampling.lower == "cc": + elif sampling.lower() == "cc": return quad_weights_cc(L) + elif sampling.lower() == "f2": + return quad_weights_f2(L) + else: raise ValueError(f"Sampling scheme sampling={sampling} not supported") @@ -95,6 +98,12 @@ def quad_weights( elif sampling.lower() == "gl": return quad_weights_gl(L) + elif sampling.lower() == "cc": + return quad_weights_cc(L) + + elif sampling.lower() == "f2": + return quad_weights_f2(L) + elif sampling.lower() == "healpix": return quad_weights_hp(nside) From 2311a4f5936cb717c9b7ad7c26dbc4a409e56d26 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 19 May 2026 15:48:41 +0100 Subject: [PATCH 07/66] Fix F2 theta positions --- s2fft/sampling/s2_samples.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index ababd7b23..1411735d6 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -263,9 +263,12 @@ def t2theta( f"Sampling scheme sampling={sampling} with L={L} not supported" ) - if sampling.lower() in ("mw", "f2"): + if sampling.lower() == "mw": return (2 * t + 1) * np.pi / (2 * L - 1) + elif sampling.lower() == "f2": + return (t + 1) * np.pi / L + elif sampling.lower() in ("mwss", "cc"): return 2 * t * np.pi / (2 * L) From 8e7259a15af8f06e181300371cb2a943f8622c82 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 19 May 2026 15:49:06 +0100 Subject: [PATCH 08/66] Separate out CC and F2 phi grid sizes --- s2fft/sampling/s2_samples.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 1411735d6..e9a2c2dfd 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -114,7 +114,10 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int: if sampling.lower() in ("mw", "dh", "gl"): return 2 * L - 1 - elif sampling.lower() in ("mwss", "cc", "f2"): + elif sampling.lower() == "mwss": + return 2 * L + + elif sampling.lower() in ("cc", "f2"): return 2 * L elif sampling.lower() == "healpix": From 3462e9551dc131ef4a0cbf2b994dff028926dab6 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 20 May 2026 13:33:15 +0100 Subject: [PATCH 09/66] Add tests for quadrature rules on known 1D integrals --- tests/test_quadrature.py | 53 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index a26f25346..264cd72f3 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -10,7 +10,7 @@ @pytest.mark.parametrize("L", [5, 6]) -@pytest.mark.parametrize("sampling", ["mw", "mwss", "dh", "gl"]) +@pytest.mark.parametrize("sampling", ["mw", "mwss", "dh", "gl", "cc"]) @pytest.mark.parametrize("method", ["numpy", "jax", "torch"]) def test_quadrature_mw_weights(flm_generator, L: int, sampling: str, method: str): spin = 0 @@ -48,3 +48,54 @@ def test_quadrature_exceptions(): with pytest.raises(ValueError): quadrature.quad_weights(L, sampling="foo") + + +def check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module, tol=1e-12): + f, true_integral = f_and_integral + thetas = samples.thetas(n_points, sampling=rule) + xs = np.cos(thetas) + weights = { + "cc": quadrature_module.quad_weights_cc_theta_only, + "f2": quadrature_module.quad_weights_f2_theta_only, + "mw": quadrature_module.quad_weights_mw_theta_only, + "mwss": quadrature_module.quad_weights_mwss_theta_only, + "gl": lambda L: quadrature_module.quad_weights_gl(L) + * (2 * L - 1) + / (2 * np.pi), + "dh": lambda L: quadrature_module.quad_weight_dh_theta_only( + samples.thetas(L, sampling="dh"), L + ), + }[rule](n_points) + quad_integral = (f(xs) * weights).sum() + assert abs(quad_integral - true_integral) < tol + + +@pytest.mark.parametrize( + "f_and_integral", + [ + (lambda x: x, 0.0), + (lambda x: x**2, 2 / 3), + (lambda x: x**3 - 2 * x**2 + x - 1, -10 / 3), + (lambda x: x**4 - x**2, -4 / 15), + ], +) +@pytest.mark.parametrize("rule", ["cc", "f2", "mw", "mwss", "gl", "dh"]) +@pytest.mark.parametrize("n_points", [6, 8, 16]) +@pytest.mark.parametrize("quadrature_module", [quadrature, quadrature_jax]) +def test_quadrature_polynomial(f_and_integral, rule, n_points, quadrature_module): + check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module) + + +@pytest.mark.parametrize( + "f_and_integral", + [ + (lambda x: np.cos(x), np.sin(1) * 2), + (lambda x: np.exp(x), np.exp(1) - np.exp(-1)), + (lambda x: np.log(1 + x**2), 2 * np.log(2) - 4 + np.pi), + ], +) +@pytest.mark.parametrize("rule", ["cc", "f2", "mw", "mwss", "gl", "dh"]) +@pytest.mark.parametrize("n_points", [32, 64, 128]) +@pytest.mark.parametrize("quadrature_module", [quadrature, quadrature_jax]) +def test_quadrature_non_polynomial(f_and_integral, rule, n_points, quadrature_module): + check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module) From 00baa27b52cd26c30bfc19ce452059f645c02e02 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 08:51:27 +0100 Subject: [PATCH 10/66] Link quadrature weight and ntheta defs for CC and F2 rules --- s2fft/utils/quadrature.py | 39 ++++++++++++++++++++++++--------------- 1 file changed, 24 insertions(+), 15 deletions(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index c13a8f74b..d203398c3 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -335,26 +335,31 @@ def mw_weights(m: int) -> float: return 0 -def _fejer_second_rule_rfft_weights(L: int, xp: ModuleType = np): +def _fejer_second_rule_rfft_weights(n_points: int, xp: ModuleType = np): r""" Compute real-valued FFT of weights for Fejer second quadrature rule. Args: - L (int): Harmonic band-limit. + n_points (int): Number of quadrature points / nodes (including left boundary point). xp (module): Array namespace to use for operations. Defaults to NumPy. Returns: - np.ndarray: Real-valued FFT of weights computed for each :math:`\theta` - (weights are identical as :math:`\phi` varies for given :math:`\theta`). + np.ndarray: Real-valued FFT of array of weights for quadrature nodes. Corresponding weigh + array computed using `irfft` includes zero weight in first element for point at left (-1) + boundary but excludes zero weight for point at right (1) boundary. References: Waldvogel, J. (2006). Fast construction of the Fejer and Clenshaw-Curtis quadrature rules. BIT Numerical Mathematics, 46(1), 195–202. https://doi.org/10.1007/s10543-006-0045-4. """ - t = xp.arange(1, L, 2) + t = xp.arange(1, n_points, 2) w0 = -2 / (t * (t - 2)) - w1 = xp.array([-2 / (L - 1)]) if L % 2 == 0 else xp.array([-1 / (L - 2)]) + w1 = ( + xp.array([-2 / (n_points - 1)]) + if n_points % 2 == 0 + else xp.array([-1 / (n_points - 2)]) + ) return xp.concatenate([w0, w1]) @@ -393,14 +398,17 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: if L < 1: raise ValueError(f"Bandlimit must be at least 1: L = {L}") - if L == 1: + n_theta = samples.ntheta(L, "cc") + + if n_theta == 2: return xp.array([1.0, 1.0]) - w = _fejer_second_rule_rfft_weights(L, xp) - g = -xp.ones(L // 2 + 1) - g = xpx.at(g)[L // 2].add(2 * L if L % 2 == 0 else L) - g /= L**2 - 1 + (L % 2) - weights = xp.fft.irfft(w + g, n=L) + n = n_theta - 1 + w = _fejer_second_rule_rfft_weights(n, xp) + g = -xp.ones(n // 2 + 1) + g = xpx.at(g)[n // 2].add(2 * n if n % 2 == 0 else n) + g /= n**2 - 1 + (n % 2) + weights = xp.fft.irfft(w + g, n=n) return xp.concatenate([weights, weights[:1]]) @@ -439,7 +447,8 @@ def quad_weights_f2_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: if L < 1: raise ValueError(f"Bandlimit must be at least 1: L = {L}") - w = _fejer_second_rule_rfft_weights(L, xp) - weights = xp.fft.irfft(w, n=L) - # Weights are zero at poles which we assume are excluded from nodes + n_theta = samples.ntheta(L, "f2") + w = _fejer_second_rule_rfft_weights(n_theta + 1, xp) + weights = xp.fft.irfft(w, n=n_theta + 1) + # Weight is zero at left boundary / north pole which we assume is excluded from nodes return weights[1:] From 21ed46158f3dd6c68ca8ef6e3686266821f7a082 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 08:52:31 +0100 Subject: [PATCH 11/66] Simplify index to phi computation using nphi def --- s2fft/sampling/s2_samples.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index e9a2c2dfd..4eae21c5b 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -402,18 +402,8 @@ def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray: np.ndarray: :math:`\phi` sample(s) for given sampling scheme. """ - if sampling.lower() in ["mw", "gl"]: - return 2 * p * np.pi / (2 * L - 1) - - elif sampling.lower() == "mwss": - return 2 * p * np.pi / (2 * L) - - elif sampling.lower() == "dh": - return 2 * p * np.pi / (2 * L - 1) - - elif sampling.lower() == "healpix": - raise ValueError(f"Sampling scheme sampling={sampling} not supported") - + if sampling.lower() in ["mw", "gl", "dh", "mw", "mwss", "cc", "f2"]: + return 2 * p * np.pi / nphi_equiang(L, sampling) else: raise ValueError(f"Sampling scheme sampling={sampling} not supported") From 9e1f84ce341346b8a739504cf9a13351ef1fe8e2 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 08:53:09 +0100 Subject: [PATCH 12/66] Increase number of theta samples for F2 rule by one --- s2fft/sampling/s2_samples.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 4eae21c5b..ca9ab59b0 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -35,7 +35,7 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: return L elif sampling.lower() == "f2": - return L - 1 + return L elif sampling.lower() == "mwss": return L + 1 @@ -270,7 +270,7 @@ def t2theta( return (2 * t + 1) * np.pi / (2 * L - 1) elif sampling.lower() == "f2": - return (t + 1) * np.pi / L + return (t + 1) * np.pi / (L + 1) elif sampling.lower() in ("mwss", "cc"): return 2 * t * np.pi / (2 * L) From f5e72787c873e59c12d3d837d5679ecb8af7c6f7 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 08:57:32 +0100 Subject: [PATCH 13/66] Factor out m offset 1 and pole singularity sets --- s2fft/base_transforms/spherical.py | 8 +++--- s2fft/precompute_transforms/spherical.py | 8 +++--- s2fft/precompute_transforms/wigner.py | 9 ++++--- s2fft/sampling/s2_samples.py | 4 +++ s2fft/sampling/so3_samples.py | 2 +- s2fft/transforms/otf_recursions.py | 31 +++++++++++++----------- s2fft/transforms/spherical.py | 12 ++++----- 7 files changed, 41 insertions(+), 33 deletions(-) diff --git a/s2fft/base_transforms/spherical.py b/s2fft/base_transforms/spherical.py index 98024c1e7..cb674546c 100644 --- a/s2fft/base_transforms/spherical.py +++ b/s2fft/base_transforms/spherical.py @@ -486,7 +486,7 @@ def _compute_inverse_sov_fft( assert L >= 2 * nside ftm = np.zeros(samples.ftm_shape(L, sampling, nside), dtype=np.complex128) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 for t, theta in enumerate(thetas): phi_ring_offset = ( @@ -576,7 +576,7 @@ def _compute_inverse_sov_fft_vectorized( """ ftm = np.zeros(samples.ftm_shape(L, sampling, nside), dtype=np.complex128) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 for t, theta in enumerate(thetas): phase_shift = ( @@ -844,7 +844,7 @@ def _compute_forward_sov_fft( flm = np.zeros(samples.flm_shape(L), dtype=np.complex128) ftm = np.zeros_like(f).astype(np.complex128) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 if sampling.lower() == "healpix": ftm = hp.healpix_fft(f, L, nside, "numpy", reality) @@ -960,7 +960,7 @@ def _compute_forward_sov_fft_vectorized( flm = np.zeros(samples.flm_shape(L), dtype=np.complex128) ftm = np.zeros_like(f).astype(np.complex128) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 if reality: m_conj = (-1) ** (np.arange(1, L) % 2) diff --git a/s2fft/precompute_transforms/spherical.py b/s2fft/precompute_transforms/spherical.py index 24952d805..cb193c8e0 100644 --- a/s2fft/precompute_transforms/spherical.py +++ b/s2fft/precompute_transforms/spherical.py @@ -120,7 +120,7 @@ def inverse_transform( np.ndarray: Pixel-space coefficients. """ - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 m_start_ind = L - 1 if reality else 0 ftm = np.zeros(samples.ftm_shape(L, sampling, nside), dtype=np.complex128) @@ -186,7 +186,7 @@ def inverse_transform_jax( jnp.ndarray: Pixel-space coefficients with shape. """ - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 m_start_ind = L - 1 if reality else 0 ftm = jnp.zeros( @@ -356,7 +356,7 @@ def forward_transform( sampling = "mwss" f = resampling.upsample_by_two_mwss(f, L, spin) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 m_start_ind = L - 1 if reality else 0 if sampling.lower() == "healpix": @@ -426,7 +426,7 @@ def forward_transform_jax( sampling = "mwss" f = resampling_jax.upsample_by_two_mwss(f, L, spin) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 m_start_ind = L - 1 if reality else 0 if sampling.lower() == "healpix": diff --git a/s2fft/precompute_transforms/wigner.py b/s2fft/precompute_transforms/wigner.py index 4222392ea..91c452472 100644 --- a/s2fft/precompute_transforms/wigner.py +++ b/s2fft/precompute_transforms/wigner.py @@ -5,6 +5,7 @@ from jax import jit from s2fft.precompute_transforms import construct +from s2fft.sampling import s2_samples from s2fft.sampling import so3_samples as samples from s2fft.utils import healpix_ffts as hp from s2fft.utils import resampling, resampling_jax, torch_wrapper @@ -112,7 +113,7 @@ def inverse_transform( np.ndarray: Pixel-space coefficients. """ - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in s2_samples.M_OFFSET_1_SCHEMES else 0 n_start_ind = N - 1 if reality else 0 fnab = np.zeros(samples.fnab_shape(L, N, sampling, nside), dtype=np.complex128) @@ -174,7 +175,7 @@ def inverse_transform_jax( jnp.ndarray: Pixel-space coefficients. """ - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in s2_samples.M_OFFSET_1_SCHEMES else 0 n_start_ind = N - 1 if reality else 0 fnab = jnp.zeros(samples.fnab_shape(L, N, sampling, nside), dtype=jnp.complex128) @@ -339,7 +340,7 @@ def forward_transform( sampling = "mwss" fban = resampling.upsample_by_two_mwss(fban, L, spins) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in s2_samples.M_OFFSET_1_SCHEMES else 0 if sampling.lower() in "healpix": temp = np.zeros(samples.fnab_shape(L, N, sampling, nside), dtype=np.complex128) @@ -422,7 +423,7 @@ def forward_transform_jax( sampling = "mwss" fban = resampling_jax.upsample_by_two_mwss(fban, L, spins) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in s2_samples.M_OFFSET_1_SCHEMES else 0 if sampling.lower() in "healpix": temp = jnp.zeros( diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index ca9ab59b0..25bcdc47d 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -1,5 +1,9 @@ import numpy as np +M_OFFSET_1_SCHEMES = frozenset(("mwss", "healpix")) +INCLUDES_SOUTH_POLE_SCHEMES = frozenset(("mw", "mwss", "cc")) +INCLUDES_NORTH_POLE_SCHEMES = frozenset(("mwss", "cc")) + def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: r""" diff --git a/s2fft/sampling/so3_samples.py b/s2fft/sampling/so3_samples.py index cd849125a..3796c853c 100644 --- a/s2fft/sampling/so3_samples.py +++ b/s2fft/sampling/so3_samples.py @@ -90,7 +90,7 @@ def fnab_shape( :math:`SO(3)`. """ - if sampling.lower() in ["mwss", "healpix"]: + if sampling.lower() in samples.M_OFFSET_1_SCHEMES: return _ngamma(N), samples.ntheta(L, sampling, nside), 2 * L elif sampling.lower() in ["mw", "dh", "gl"]: diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py index 8eae9cfae..6b9c99072 100644 --- a/s2fft/transforms/otf_recursions.py +++ b/s2fft/transforms/otf_recursions.py @@ -94,7 +94,7 @@ def inverse_latitudinal_step( for i in range(2): if not (reality and i == 0): - m_offset = 1 if sampling in ["mwss", "healpix"] and i == 0 else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES and i == 0 else 0 lind = L - 1 - L_lower sind = lims[i] @@ -237,8 +237,9 @@ def inverse_latitudinal_step_jax( """ mm = -spin # switch to match convention ntheta = len(beta) # Number of theta samples - m_count = 2 * L if sampling.lower() in ["mwss", "healpix"] else 2 * L - 1 - ftm = jnp.zeros((ntheta, m_count), dtype=jnp.complex128) + ftm = jnp.zeros( + (ntheta, samples.ftm_shape(L, sampling, nside)[1]), dtype=jnp.complex128 + ) el = jnp.arange(L_lower, L) # Trigonometric constant adopted throughout @@ -258,7 +259,7 @@ def inverse_latitudinal_step_jax( for i in range(2): if not (reality and i == 0): - m_offset = 1 if sampling in ["mwss", "healpix"] and i == 0 else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES and i == 0 else 0 lind = L - 1 - L_lower sind = lims[i] @@ -410,8 +411,8 @@ def eval_recursion_step( ) # Remove south pole singularity - m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0 - if sampling.lower() in ["mw", "mwss"]: + m_offset = 1 if sampling.lower() in samples.M_OFFSET_1_SCHEMES else 0 + if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: ftm = ftm.at[-1].set(0) ftm = ftm.at[-1, L - 1 + spin + m_offset].set( jnp.nansum( @@ -420,7 +421,7 @@ def eval_recursion_step( ) # Remove north pole singularity - if sampling.lower() == "mwss": + if sampling.lower() in samples.INCLUDES_NORTH_POLE_SCHEMES: ftm = ftm.at[0].set(0) ftm = ftm.at[0, L - 1 - spin + m_offset].set( jnp.nansum(flm[L_lower:, L - 1 - spin]) @@ -504,7 +505,7 @@ def forward_latitudinal_step( for i in range(2): if not (reality and i == 0): - m_offset = 1 if sampling in ["mwss", "healpix"] and i == 0 else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES and i == 0 else 0 lind = L - 1 - L_lower sind = lims[i] @@ -653,10 +654,12 @@ def forward_latitudinal_step_jax( """ # Avoid pole-singularities for MWSS sampling - if sampling.lower() == "mwss": + if sampling.lower() in ( + samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES + ): ftm = ftm_in[1:-1] beta = beta_in[1:-1] - elif sampling.lower() == "mw": + elif sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: ftm = ftm_in[:-1] beta = beta_in[:-1] else: @@ -685,7 +688,7 @@ def forward_latitudinal_step_jax( for i in range(2): if not (reality and i == 0): - m_offset = 1 if sampling in ["mwss", "healpix"] and i == 0 else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES and i == 0 else 0 lind = L - 1 - L_lower sind = lims[i] @@ -874,13 +877,13 @@ def eval_recursion_step( ) # Include both pole singularities explicitly - m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0 - if sampling.lower() in ["mw", "mwss"]: + m_offset = 1 if sampling.lower() in samples.M_OFFSET_1_SCHEMES else 0 + if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: flm = flm.at[L_lower:, L - 1 + spin].add( (-1) ** abs(jnp.arange(L_lower, L) - spin) * ftm_in[-1, L - 1 + spin + m_offset] ) - if sampling.lower() == "mwss": + if sampling.lower() in samples.INCLUDES_NORTH_POLE_SCHEMES: flm = flm.at[L_lower:, L - 1 - spin].add(ftm_in[0, L - 1 - spin + m_offset]) return flm diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py index 92126ce65..940f59cf5 100644 --- a/s2fft/transforms/spherical.py +++ b/s2fft/transforms/spherical.py @@ -155,7 +155,7 @@ def inverse_numpy( """ # Define latitudinal sample positions and Fourier offsets thetas = samples.thetas(L, sampling, nside) - m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling.lower() in samples.M_OFFSET_1_SCHEMES else 0 m_start_ind = L - 1 if reality else 0 L0 = L_lower @@ -173,13 +173,13 @@ def inverse_numpy( ) # Remove south pole singularity - if sampling.lower() in ["mw", "mwss"]: + if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: ftm[-1] = 0 ftm[-1, L - 1 + spin + m_offset] = np.nansum( (-1) ** abs(np.arange(L0, L) - spin) * flm[L0:, L - 1 + spin] ) # Remove north pole singularity - if sampling.lower() == "mwss": + if sampling.lower() in samples.INCLUDES_NORTH_POLE_SCHEMES: ftm[0] = 0 ftm[0, L - 1 - spin + m_offset] = jnp.nansum(flm[L0:, L - 1 - spin]) @@ -274,7 +274,7 @@ def inverse_jax( """ # Define latitudinal sample positions and Fourier offsets thetas = samples.thetas(L, sampling, nside) - m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling.lower() in samples.M_OFFSET_1_SCHEMES else 0 m_start_ind = L - 1 if reality else 0 # Apply harmonic normalisation @@ -520,7 +520,7 @@ def forward_numpy( # Define latitudinal sample positions and Fourier offsets weights = quadrature.quad_weights_transform(L, sampling, 0, nside) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 m_start_ind = L - 1 if reality else 0 L0 = L_lower @@ -664,7 +664,7 @@ def forward_jax( # Define latitudinal sample positions and Fourier offsets weights = quadrature_jax.quad_weights_transform(L, sampling, nside) - m_offset = 1 if sampling in ["mwss", "healpix"] else 0 + m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 m_start_ind = L - 1 if reality else 0 # Perform longitundal Fast Fourier Transforms From 63790b851c4fd15f0bfbd55248d2c59fcaf442ce Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 10:45:54 +0100 Subject: [PATCH 14/66] Factor out Gl quadrature theta only to match other rules --- s2fft/utils/quadrature.py | 24 +++++++++++++++++++----- s2fft/utils/quadrature_jax.py | 30 +++++++++++++++++++++++++----- tests/test_quadrature.py | 4 +--- 3 files changed, 45 insertions(+), 13 deletions(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index d203398c3..100964b6a 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -137,16 +137,15 @@ def quad_weights_hp(nside: int) -> np.ndarray: return hp_weights -def quad_weights_gl(L: int) -> np.ndarray: +def quad_weights_gl_theta_only(L: int) -> np.ndarray: r""" - Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration. + Compute GL quadrature weights for :math:`\theta` integration. Args: L (int): Harmonic band-limit. Returns: - np.ndarray: Weights computed for each :math:`\theta` (weights are identical - as :math:`\phi` varies for given :math:`\theta`). + np.ndarray: Weights computed for each :math:`\theta`. """ x1, x2 = -1.0, 1.0 @@ -173,7 +172,22 @@ def quad_weights_gl(L: int) -> np.ndarray: weights[i - 1] = 2.0 * x1 / ((1.0 - z**2) * pp * pp) weights[L + 1 - i - 1] = weights[i - 1] - return weights * 2 * np.pi / (2 * L - 1) + return weights + + +def quad_weights_gl(L: int) -> np.ndarray: + r""" + Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration. + + Args: + L (int): Harmonic band-limit. + + Returns: + np.ndarray: Weights computed for each :math:`\theta` (weights are identical + as :math:`\phi` varies for given :math:`\theta`). + + """ + return quad_weights_gl_theta_only(L) * 2 * np.pi / samples.nphi_equiang(L, "gl") def quad_weights_dh(L: int) -> np.ndarray: diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py index 1c4d7761f..fae2c8903 100644 --- a/s2fft/utils/quadrature_jax.py +++ b/s2fft/utils/quadrature_jax.py @@ -148,16 +148,15 @@ def quad_weights_hp(nside: int) -> jnp.ndarray: @_partial(_jit, static_argnums=(0)) -def quad_weights_gl(L: int) -> jnp.ndarray: +def quad_weights_gl_theta_only(L: int) -> jnp.ndarray: r""" - Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration. + Compute GL quadrature weights for :math:`\theta` integration. Args: L (int): Harmonic band-limit. Returns: - jnp.ndarray: Weights computed for each :math:`\theta` (weights are identical - as :math:`\phi` varies for given :math:`\theta`). + jnp.ndarray: Weights computed for each :math:`\theta`. """ x1, x2 = -1.0, 1.0 @@ -198,7 +197,28 @@ def inner_loop_step(j, p1_p2): weights = weights.at[i - 1].set(2.0 * x1 / ((1.0 - z**2) * pp * pp)) weights = weights.at[L + 1 - i - 1].set(weights[i - 1]) - return weights * 2 * jnp.pi / (2 * L - 1) + return weights + + +@_partial(_jit, static_argnums=(0)) +def quad_weights_gl(L: int) -> jnp.ndarray: + r""" + Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration. + + Args: + L (int): Harmonic band-limit. + + Returns: + jnp.ndarray: Weights computed for each :math:`\theta` (weights are identical + as :math:`\phi` varies for given :math:`\theta`). + + """ + return ( + quad_weights_gl_theta_only(L) + * 2 + * jnp.pi + / samples.nphi_equiang(L, sampling="gl") + ) @_partial(_jit, static_argnums=(0)) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index 264cd72f3..20cd33be1 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -59,9 +59,7 @@ def check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module, tol "f2": quadrature_module.quad_weights_f2_theta_only, "mw": quadrature_module.quad_weights_mw_theta_only, "mwss": quadrature_module.quad_weights_mwss_theta_only, - "gl": lambda L: quadrature_module.quad_weights_gl(L) - * (2 * L - 1) - / (2 * np.pi), + "gl": quadrature_module.quad_weights_gl_theta_only, "dh": lambda L: quadrature_module.quad_weight_dh_theta_only( samples.thetas(L, sampling="dh"), L ), From a0e0d723f8604fdfa229b9196aceb18eb9631237 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 11:07:33 +0100 Subject: [PATCH 15/66] Match signature of DH theta only quadrature weight function to other rules --- docs/api/utility/index.rst | 4 ++-- s2fft/utils/quadrature.py | 20 +++++++------------- s2fft/utils/quadrature_jax.py | 23 ++++++++++------------- tests/test_quadrature.py | 4 +--- 4 files changed, 20 insertions(+), 31 deletions(-) diff --git a/docs/api/utility/index.rst b/docs/api/utility/index.rst index ae76052cb..16082c362 100644 --- a/docs/api/utility/index.rst +++ b/docs/api/utility/index.rst @@ -58,8 +58,8 @@ Utility Functions - Compute MW quadrature weights for :math:`\theta` and :math:`\phi` integration. * - :func:`~s2fft.utils.quadrature.quad_weights_mwss` - Compute MWSS quadrature weights for :math:`\theta` and :math:`\phi` integration. - * - :func:`~s2fft.utils.quadrature.quad_weight_dh_theta_only` - - Compute DH quadrature weight for :math:`\theta` integration (only), for given :math:`\theta`. + * - :func:`~s2fft.utils.quadrature.quad_weights_dh_theta_only` + - Compute DH quadrature weight for :math:`\theta` integration (only). * - :func:`~s2fft.utils.quadrature.quad_weights_mw_theta_only` - Compute MW quadrature weights for :math:`\theta` integration (only). * - :func:`~s2fft.utils.quadrature.quad_weights_mwss_theta_only` diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index 100964b6a..2635d4fd5 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -202,31 +202,25 @@ def quad_weights_dh(L: int) -> np.ndarray: as :math:`\phi` varies for given :math:`\theta`). """ - q = quad_weight_dh_theta_only(samples.thetas(L, sampling="dh"), L) + return quad_weights_dh_theta_only(L) * 2 * np.pi / samples.nphi_equiang(L, "dh") - return q * 2 * np.pi / (2 * L - 1) - -def quad_weight_dh_theta_only(theta: float, L: int) -> float: +def quad_weights_dh_theta_only(L: int) -> np.ndarray: r""" - Compute DH quadrature weight for :math:`\theta` integration (only), for given - :math:`\theta`. + Compute DH quadrature weights for :math:`\theta` integration (only). Args: - theta (float): :math:`\theta` angle for which to compute weight. - L (int): Harmonic band-limit. Returns: - float: Weight computed for each :math:`\theta`. + float: Weights computed for each :math:`\theta`. """ + thetas = samples.thetas(L, sampling="dh") w = 0.0 for k in range(0, L): - w += np.sin((2 * k + 1) * theta) / (2 * k + 1) - - w *= 2 / L * np.sin(theta) - + w += np.sin((2 * k + 1) * thetas) / (2 * k + 1) + w *= 2 / L * np.sin(thetas) return w diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py index fae2c8903..0a9cea9f1 100644 --- a/s2fft/utils/quadrature_jax.py +++ b/s2fft/utils/quadrature_jax.py @@ -235,33 +235,30 @@ def quad_weights_dh(L: int) -> jnp.ndarray: as :math:`\phi` varies for given :math:`\theta`). """ - q = quad_weight_dh_theta_only(samples.thetas(L, sampling="dh"), L) + return quad_weights_dh_theta_only(L) * 2 * jnp.pi / samples.nphi_equiang(L, "dh") - return q * 2 * jnp.pi / (2 * L - 1) - -@_partial(_jit, static_argnums=(1)) -def quad_weight_dh_theta_only(theta: float, L: int) -> float: +@_partial(_jit, static_argnums=(0)) +def quad_weights_dh_theta_only(L: int) -> jnp.ndarray: r""" - Compute DH quadrature weight for :math:`\theta` integration (only), for given - :math:`\theta`. JAX implementation of :func:`s2fft.quadrature.quad_weights_dh_theta_only`. + Compute DH quadrature weights for :math:`\theta` integration (only). + JAX implementation of :func:`s2fft.quadrature.quad_weights_dh_theta_only`. Args: - theta (float): :math:`\theta` angle for which to compute weight. - L (int): Harmonic band-limit. Returns: - float: Weight computed for each :math:`\theta`. + jnp.ndarray: Weight computed for each :math:`\theta`. """ + thetas = samples.thetas(L, sampling="dh") def increment_weights(k, w): - w += jnp.sin((2 * k + 1) * theta) / (2 * k + 1) + w += jnp.sin((2 * k + 1) * thetas) / (2 * k + 1) return w - w = jax.lax.fori_loop(0, L, increment_weights, jnp.zeros_like(theta)) - w *= 2 / L * jnp.sin(theta) + w = jax.lax.fori_loop(0, L, increment_weights, jnp.zeros_like(thetas)) + w *= 2 / L * jnp.sin(thetas) return w diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index 20cd33be1..36e2fdf87 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -60,9 +60,7 @@ def check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module, tol "mw": quadrature_module.quad_weights_mw_theta_only, "mwss": quadrature_module.quad_weights_mwss_theta_only, "gl": quadrature_module.quad_weights_gl_theta_only, - "dh": lambda L: quadrature_module.quad_weight_dh_theta_only( - samples.thetas(L, sampling="dh"), L - ), + "dh": quadrature_module.quad_weights_dh_theta_only, }[rule](n_points) quad_integral = (f(xs) * weights).sum() assert abs(quad_integral - true_integral) < tol From 37bac505f50855d407f8add7b31dc176fc7b99f8 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 11:08:06 +0100 Subject: [PATCH 16/66] Use nphi_equiang function uniformly in quadrature weight implementations --- s2fft/utils/quadrature.py | 11 +++++++++-- s2fft/utils/quadrature_jax.py | 20 +++++++++++--------- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index 2635d4fd5..d380ca949 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -238,7 +238,9 @@ def quad_weights_mw(L: int, spin: int = 0) -> np.ndarray: as :math:`\phi` varies for given :math:`\theta`). """ - return quad_weights_mw_theta_only(L, spin) * 2 * np.pi / (2 * L - 1) + return ( + quad_weights_mw_theta_only(L, spin) * 2 * np.pi / samples.nphi_equiang(L, "mw") + ) def quad_weights_mwss(L: int, spin: int = 0) -> np.ndarray: @@ -255,7 +257,12 @@ def quad_weights_mwss(L: int, spin: int = 0) -> np.ndarray: as :math:`\phi` varies for given :math:`\theta`). """ - return quad_weights_mwss_theta_only(L, spin) * 2 * np.pi / (2 * L) + return ( + quad_weights_mwss_theta_only(L, spin) + * 2 + * np.pi + / samples.nphi_equiang(L, "mwss") + ) def quad_weights_mwss_theta_only(L: int, spin: int = 0) -> np.ndarray: diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py index 0a9cea9f1..b8bf7ff73 100644 --- a/s2fft/utils/quadrature_jax.py +++ b/s2fft/utils/quadrature_jax.py @@ -50,7 +50,12 @@ def quad_weights_transform( """ if sampling.lower() == "mwss": - return quad_weights_mwss_theta_only(2 * L) * 2 * jnp.pi / (2 * L) + return ( + quad_weights_mwss_theta_only(2 * L) + * 2 + * jnp.pi + / samples.nphi_equiang(L, "mw") + ) elif sampling.lower() == "dh": return quad_weights_dh(L) @@ -213,12 +218,7 @@ def quad_weights_gl(L: int) -> jnp.ndarray: as :math:`\phi` varies for given :math:`\theta`). """ - return ( - quad_weights_gl_theta_only(L) - * 2 - * jnp.pi - / samples.nphi_equiang(L, sampling="gl") - ) + return quad_weights_gl_theta_only(L) * 2 * jnp.pi / samples.nphi_equiang(L, "gl") @_partial(_jit, static_argnums=(0)) @@ -279,7 +279,7 @@ def quad_weights_mw(L: int) -> jnp.ndarray: as :math:`\phi` varies for given :math:`\theta`). """ - return quad_weights_mw_theta_only(L) * 2 * jnp.pi / (2 * L - 1) + return quad_weights_mw_theta_only(L) * 2 * jnp.pi / samples.nphi_equiang(L, "mw") @_partial(_jit, static_argnums=(0)) @@ -298,7 +298,9 @@ def quad_weights_mwss(L: int) -> jnp.ndarray: as :math:`\phi` varies for given :math:`\theta`). """ - return quad_weights_mwss_theta_only(L) * 2 * jnp.pi / (2 * L) + return ( + quad_weights_mwss_theta_only(L) * 2 * jnp.pi / samples.nphi_equiang(L, "mwss") + ) @_partial(_jit, static_argnums=(0)) From a0c2891db940c7183262f9fc77e9ad6deb120838 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 13:08:18 +0100 Subject: [PATCH 17/66] Refactor index to theta to use ntheta --- s2fft/sampling/s2_samples.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 25bcdc47d..d21809ed5 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -271,16 +271,16 @@ def t2theta( ) if sampling.lower() == "mw": - return (2 * t + 1) * np.pi / (2 * L - 1) + return (2 * t + 1) * np.pi / (2 * ntheta(L, sampling) - 1) elif sampling.lower() == "f2": - return (t + 1) * np.pi / (L + 1) + return (t + 1) * np.pi / (ntheta(L, sampling) + 1) elif sampling.lower() in ("mwss", "cc"): - return 2 * t * np.pi / (2 * L) + return t * np.pi / (ntheta(L, sampling) - 1) elif sampling.lower() == "dh": - return (2 * t + 1) * np.pi / (4 * L) + return (2 * t + 1) * np.pi / (2 * ntheta(L, sampling)) elif sampling.lower() == "healpix": if nside is None: From 530babe8ba3a329fa178706bd2d20f35963a4172 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 13:18:48 +0100 Subject: [PATCH 18/66] Double CC and F2 scheme number of theta samples --- s2fft/sampling/s2_samples.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index d21809ed5..0ae03b942 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -39,13 +39,13 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: return L elif sampling.lower() == "f2": - return L + return 2 * L - 1 elif sampling.lower() == "mwss": return L + 1 elif sampling.lower() == "cc": - return L + 1 + return 2 * L - 1 elif sampling.lower() == "dh": return 2 * L From 1ac5fc3c6063853a79faeb4be1a2134a3e839369 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 13:19:18 +0100 Subject: [PATCH 19/66] Include CC and F2 in m offset 1 schemes --- s2fft/sampling/s2_samples.py | 2 +- tests/test_quadrature.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 0ae03b942..492033d89 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -1,6 +1,6 @@ import numpy as np -M_OFFSET_1_SCHEMES = frozenset(("mwss", "healpix")) +M_OFFSET_1_SCHEMES = frozenset(("mwss", "healpix", "cc", "f2")) INCLUDES_SOUTH_POLE_SCHEMES = frozenset(("mw", "mwss", "cc")) INCLUDES_NORTH_POLE_SCHEMES = frozenset(("mwss", "cc")) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index 36e2fdf87..ddfde414b 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -10,7 +10,7 @@ @pytest.mark.parametrize("L", [5, 6]) -@pytest.mark.parametrize("sampling", ["mw", "mwss", "dh", "gl", "cc"]) +@pytest.mark.parametrize("sampling", ["mw", "mwss", "dh", "gl", "cc", "f2"]) @pytest.mark.parametrize("method", ["numpy", "jax", "torch"]) def test_quadrature_mw_weights(flm_generator, L: int, sampling: str, method: str): spin = 0 From aa09f7cc0f872e2c067b4a77b542c902c902f0ca Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 13:58:25 +0100 Subject: [PATCH 20/66] Correct bug in determining scheme for dealing with pole singularity --- s2fft/transforms/otf_recursions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py index 6b9c99072..79510a308 100644 --- a/s2fft/transforms/otf_recursions.py +++ b/s2fft/transforms/otf_recursions.py @@ -659,7 +659,9 @@ def forward_latitudinal_step_jax( ): ftm = ftm_in[1:-1] beta = beta_in[1:-1] - elif sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: + elif sampling.lower() in ( + samples.INCLUDES_SOUTH_POLE_SCHEMES - samples.INCLUDES_NORTH_POLE_SCHEMES + ): ftm = ftm_in[:-1] beta = beta_in[:-1] else: From 9366a6c59399516917ddd2a8a96b553f423c6bb7 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 13:58:52 +0100 Subject: [PATCH 21/66] Correct bug in using MW rather than MWSS to compute nphi --- s2fft/utils/quadrature.py | 7 ++++++- s2fft/utils/quadrature_jax.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index d380ca949..dc5eff61d 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -38,7 +38,12 @@ def quad_weights_transform( """ if sampling.lower() == "mwss": - return quad_weights_mwss_theta_only(2 * L, spin=0) * 2 * np.pi / (2 * L) + return ( + quad_weights_mwss_theta_only(2 * L, spin=0) + * 2 + * np.pi + / samples.nphi_equiang(L, "mwss") + ) elif sampling.lower() == "dh": return quad_weights_dh(L) diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py index b8bf7ff73..bbd805193 100644 --- a/s2fft/utils/quadrature_jax.py +++ b/s2fft/utils/quadrature_jax.py @@ -54,7 +54,7 @@ def quad_weights_transform( quad_weights_mwss_theta_only(2 * L) * 2 * jnp.pi - / samples.nphi_equiang(L, "mw") + / samples.nphi_equiang(L, "mwss") ) elif sampling.lower() == "dh": From 417b374d6bb74d06cefa42a09c8094e024d70cc1 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 14:13:22 +0100 Subject: [PATCH 22/66] Generalize additional implicit special casing for handling pole singularities in numpy forward transform --- s2fft/recursions/price_mcewen.py | 7 ++++++- s2fft/transforms/spherical.py | 8 ++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/s2fft/recursions/price_mcewen.py b/s2fft/recursions/price_mcewen.py index 1c98253da..3535a0fd5 100644 --- a/s2fft/recursions/price_mcewen.py +++ b/s2fft/recursions/price_mcewen.py @@ -55,10 +55,15 @@ def generate_precomputes( # Correct for mw to mwss conversion if forward and sampling.lower() in ["mw", "mwss"]: sampling = "mwss" - beta = samples.thetas(2 * L, "mwss")[1:-1] + beta = samples.thetas(2 * L, "mwss") else: beta = samples.thetas(L, sampling, nside) + if forward and sampling.lower() in ( + samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES + ): + beta = beta[1:-1] + ntheta = len(beta) # Number of theta samples el = np.arange(L0, L) diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py index 940f59cf5..f163839e0 100644 --- a/s2fft/transforms/spherical.py +++ b/s2fft/transforms/spherical.py @@ -545,8 +545,12 @@ def forward_numpy( phase_shifts = hp.ring_phase_shifts_hp(L, nside, True, reality) ftm[:, m_start_ind + m_offset :] *= phase_shifts + INCLUDES_BOTH_POLES_SCHEMES = ( + samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES + ) + # Perform latitudinal wigner-d recursions - if sampling.lower() == "mwss": + if sampling.lower() in INCLUDES_BOTH_POLES_SCHEMES: flm = otf.forward_latitudinal_step( ftm[1:-1], thetas[1:-1], @@ -564,7 +568,7 @@ def forward_numpy( ) # Include both pole singularities explicitly - if sampling.lower() == "mwss": + if sampling.lower() in INCLUDES_BOTH_POLES_SCHEMES: flm[L0:, L - 1 + spin] += (-1) ** abs(np.arange(L0, L) - spin) * ftm[ -1, L - 1 + spin + m_offset ] From 0daaeaa00fc8130c46e196544a6ef8e85eafb891 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 14:51:27 +0100 Subject: [PATCH 23/66] Add CC and F2 to sample scheme lists in docstrings --- s2fft/base_transforms/spherical.py | 24 ++++++++++++------------ s2fft/precompute_transforms/construct.py | 4 ++-- s2fft/precompute_transforms/spherical.py | 12 ++++++------ s2fft/recursions/price_mcewen.py | 8 ++++---- s2fft/sampling/s2_samples.py | 8 ++++---- s2fft/transforms/otf_recursions.py | 8 ++++---- s2fft/transforms/spherical.py | 12 ++++++------ s2fft/utils/quadrature.py | 2 +- s2fft/utils/quadrature_jax.py | 2 +- 9 files changed, 40 insertions(+), 40 deletions(-) diff --git a/s2fft/base_transforms/spherical.py b/s2fft/base_transforms/spherical.py index cb674546c..adbf5f40b 100644 --- a/s2fft/base_transforms/spherical.py +++ b/s2fft/base_transforms/spherical.py @@ -31,7 +31,7 @@ def inverse( spin (int, optional): Harmonic spin. Defaults to 0. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -80,7 +80,7 @@ def _inverse( spin (int, optional): Harmonic spin. Defaults to 0. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". method (str, optional): Harmonic transform algorithm. Supported algorithms include {"direct", "sov", "sov_fft", "sov_fft_vectorized"}. Defaults to @@ -154,7 +154,7 @@ def forward( spin (int, optional): Harmonic spin. Defaults to 0. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -217,7 +217,7 @@ def _forward( spin (int, optional): Harmonic spin. Defaults to 0. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". method (str, optional): Harmonic transform algorithm. Supported algorithms include {"direct", "sov", "sov_fft", "sov_fft_vectorized"}. Defaults to @@ -304,7 +304,7 @@ def _compute_inverse_direct( spin (int): Harmonic spin. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. thetas (np.ndarray): Vector of sample positions in :math:`\theta` on the sphere. @@ -391,7 +391,7 @@ def _compute_inverse_sov( spin (int): Harmonic spin. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. thetas (np.ndarray): Vector of sample positions in :math:`\theta` on the sphere. @@ -465,7 +465,7 @@ def _compute_inverse_sov_fft( spin (int): Harmonic spin. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. thetas (np.ndarray): Vector of sample positions in :math:`\theta` on the sphere. @@ -558,7 +558,7 @@ def _compute_inverse_sov_fft_vectorized( spin (int): Harmonic spin. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. thetas (np.ndarray): Vector of sample positions in :math:`\theta` on the sphere. @@ -634,7 +634,7 @@ def _compute_forward_direct( spin (int): Harmonic spin. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. thetas (np.ndarray): Vector of sample positions in :math:`\theta` on the sphere. @@ -726,7 +726,7 @@ def _compute_forward_sov( spin (int): Harmonic spin. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. thetas (np.ndarray): Vector of sample positions in :math:`\theta` on the sphere. @@ -822,7 +822,7 @@ def _compute_forward_sov_fft( spin (int): Harmonic spin. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. thetas (np.ndarray): Vector of sample positions in :math:`\theta` on the sphere. @@ -938,7 +938,7 @@ def _compute_forward_sov_fft_vectorized( spin (int): Harmonic spin. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. thetas (np.ndarray): Vector of sample positions in :math:`\theta` on the sphere. diff --git a/s2fft/precompute_transforms/construct.py b/s2fft/precompute_transforms/construct.py index 90c0179e4..52b965ff3 100644 --- a/s2fft/precompute_transforms/construct.py +++ b/s2fft/precompute_transforms/construct.py @@ -355,7 +355,7 @@ def wigner_kernel( Defaults to False. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int): HEALPix Nside resolution parameter. Only required if sampling="healpix". @@ -494,7 +494,7 @@ def wigner_kernel_jax( Defaults to False. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int): HEALPix Nside resolution parameter. Only required if sampling="healpix". diff --git a/s2fft/precompute_transforms/spherical.py b/s2fft/precompute_transforms/spherical.py index cb193c8e0..4668b5b2e 100644 --- a/s2fft/precompute_transforms/spherical.py +++ b/s2fft/precompute_transforms/spherical.py @@ -40,7 +40,7 @@ def inverse( kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. @@ -106,7 +106,7 @@ def inverse_transform( L (int): Harmonic band-limit. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. @@ -172,7 +172,7 @@ def inverse_transform_jax( L (int): Harmonic band-limit. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. @@ -250,7 +250,7 @@ def forward( kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. @@ -335,7 +335,7 @@ def forward_transform( L (int): Harmonic band-limit. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. @@ -405,7 +405,7 @@ def forward_transform_jax( L (int): Harmonic band-limit. sampling (str): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. diff --git a/s2fft/recursions/price_mcewen.py b/s2fft/recursions/price_mcewen.py index 3535a0fd5..9ee94d037 100644 --- a/s2fft/recursions/price_mcewen.py +++ b/s2fft/recursions/price_mcewen.py @@ -32,7 +32,7 @@ def generate_precomputes( spin (int, optional): Harmonic spin. Defaults to 0. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -142,7 +142,7 @@ def generate_precomputes_jax( spin (int, optional): Harmonic spin. Defaults to 0. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -282,7 +282,7 @@ def generate_precomputes_wigner( N (int): Azimuthal bandlimit sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -334,7 +334,7 @@ def generate_precomputes_wigner_jax( N (int): Azimuthal bandlimit sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 492033d89..0ff60b1ef 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -14,7 +14,7 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -374,7 +374,7 @@ def phis_equiang(L: int, sampling: str = "mw") -> np.ndarray: L (int, optional): Harmonic band-limit. sampling (str, optional): Sampling scheme. Supported equiangular sampling - schemes include {"mw", "mwss", "dh", "gl"}. Defaults to "mw". + schemes include {"mw", "mwss", "dh", "gl", "cc", "f2"}. Defaults to "mw". Returns: np.ndarray: Array of :math:`\phi` samples for given sampling scheme. @@ -395,7 +395,7 @@ def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray: p (int): :math:`\phi` index. sampling (str, optional): Sampling scheme. Supported equiangular sampling - schemes include {"mw", "mwss", "dh", "gl"}. Defaults to "mw". + schemes include {"mw", "mwss", "dh", "gl", "cc", "f2"}. Defaults to "mw". Raises: ValueError: HEALPix sampling not support (only equiangular schemes supported). @@ -454,7 +454,7 @@ def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> tuple[int L (int, optional): Harmonic band-limit. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py index 79510a308..fb9151523 100644 --- a/s2fft/transforms/otf_recursions.py +++ b/s2fft/transforms/otf_recursions.py @@ -48,7 +48,7 @@ def inverse_latitudinal_step( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to @@ -209,7 +209,7 @@ def inverse_latitudinal_step_jax( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to @@ -466,7 +466,7 @@ def forward_latitudinal_step( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to @@ -627,7 +627,7 @@ def forward_latitudinal_step_jax( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py index f163839e0..9bd6e2e12 100644 --- a/s2fft/transforms/spherical.py +++ b/s2fft/transforms/spherical.py @@ -45,7 +45,7 @@ def inverse( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". method (str, optional): Execution mode in {"numpy", "jax", "jax_cuda", "jax_ssht", "jax_healpy"}. Defaults to "numpy". @@ -137,7 +137,7 @@ def inverse_numpy( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to @@ -240,7 +240,7 @@ def inverse_jax( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to @@ -377,7 +377,7 @@ def forward( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". method (str, optional): Execution mode in {"numpy", "jax", "jax_cuda", jax_ssht", "jax_healpy"}. Defaults to "numpy". @@ -492,7 +492,7 @@ def forward_numpy( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to @@ -624,7 +624,7 @@ def forward_jax( if sampling="healpix". Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". reality (bool, optional): Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index dc5eff61d..ce002cc5a 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -76,7 +76,7 @@ def quad_weights( Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". spin (int, optional): Harmonic spin. Defaults to 0. diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py index bbd805193..03f5cf355 100644 --- a/s2fft/utils/quadrature_jax.py +++ b/s2fft/utils/quadrature_jax.py @@ -88,7 +88,7 @@ def quad_weights(L: int = None, sampling: str = "mw", nside: int = None) -> jnp. Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". spin (int, optional): Harmonic spin. Defaults to 0. From 7f64fd72400a7720c675f27c1be3421d37b51773 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 15:01:18 +0100 Subject: [PATCH 24/66] Include CC and F2 schemes in spherical transform tests --- tests/test_spherical_base.py | 2 +- tests/test_spherical_custom_grads.py | 2 +- tests/test_spherical_precompute.py | 2 +- tests/test_spherical_transform.py | 2 +- tests/test_transforms_single_precision.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_spherical_base.py b/tests/test_spherical_base.py index 399e63a1b..07aafacf5 100644 --- a/tests/test_spherical_base.py +++ b/tests/test_spherical_base.py @@ -11,7 +11,7 @@ spin_to_test = [-2, 0, 1] nside_to_test = [4, 5] L_to_nside_ratio = [2, 3] -sampling_to_test = ["mw", "mwss", "dh", "gl"] +sampling_to_test = ["mw", "mwss", "dh", "gl", "cc", "f2"] method_to_test = ["direct", "sov", "sov_fft", "sov_fft_vectorized"] reality_to_test = [False, True] diff --git a/tests/test_spherical_custom_grads.py b/tests/test_spherical_custom_grads.py index 8e0e76b94..d1437d6e2 100644 --- a/tests/test_spherical_custom_grads.py +++ b/tests/test_spherical_custom_grads.py @@ -12,7 +12,7 @@ L_lower_to_test = [2] spin_to_test = [-2, 0, 1] nside_to_test = [8] -sampling_to_test = ["mw", "mwss", "dh", "gl"] +sampling_to_test = ["mw", "mwss", "dh", "gl", "cc", "f2"] reality_to_test = [False, True] diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py index 0d428be63..ee1c38379 100644 --- a/tests/test_spherical_precompute.py +++ b/tests/test_spherical_precompute.py @@ -20,7 +20,7 @@ spin_to_test = [-2, 0, 6] nside_to_test = [4, 5] L_to_nside_ratio = [2, 3] -sampling_to_test = ["mw", "mwss", "dh", "gl"] +sampling_to_test = ["mw", "mwss", "dh", "gl", "cc", "f2"] reality_to_test = [True, False] methods_to_test = ["numpy", "jax", "torch"] recursions_to_test = ["price-mcewen", "risbo", "auto"] diff --git a/tests/test_spherical_transform.py b/tests/test_spherical_transform.py index 73fd3c9d6..d606c287f 100644 --- a/tests/test_spherical_transform.py +++ b/tests/test_spherical_transform.py @@ -15,7 +15,7 @@ L_lower_to_test = [0, 2] spin_to_test = [-2, 0, 1] nside_to_test = [4, 5] -sampling_to_test = ["mw", "mwss", "dh", "gl"] +sampling_to_test = ["mw", "mwss", "dh", "gl", "cc", "f2"] method_to_test = ["numpy", "jax", "torch"] reality_to_test = [False, True] multiple_gpus = [False, True] diff --git a/tests/test_transforms_single_precision.py b/tests/test_transforms_single_precision.py index 29e76d9f2..44ba8e1ee 100644 --- a/tests/test_transforms_single_precision.py +++ b/tests/test_transforms_single_precision.py @@ -10,7 +10,7 @@ jax.config.update("jax_enable_x64", True) -@pytest.mark.parametrize("sampling", ["mw", "mwss", "gl", "dh", "healpix"]) +@pytest.mark.parametrize("sampling", ["mw", "mwss", "gl", "dh", "healpix", "cc", "f2"]) @pytest.mark.parametrize("fwd", [True, False], ids=("forward", "inverse")) @pytest.mark.parametrize( "downsample_kernel", [True, False], ids=("downsample kernel", "full kernel") From f801d6c7a093a6f990a5a8bbb340915ae9263bcc Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 15:38:13 +0100 Subject: [PATCH 25/66] Move array-api-extra to package rather than build dependencies --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 62f63c2f7..967d01b08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,6 @@ requires = [ "scikit-build-core>=0.4.3", "nanobind>=1.3.2", "jax >= 0.5.0", - "array-api-extra >= 0.10.1", ] build-backend = "scikit_build_core.build" @@ -29,6 +28,7 @@ description = "Differentiable and accelerated spherical transforms with JAX" dependencies = [ "numpy>=1.20", "jax>=0.5.0", + "array-api-extra >= 0.10.1", ] dynamic = [ "version", From ed9dcc5630e0c30622d99adf18abddbe6cc8571a Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 17:09:15 +0100 Subject: [PATCH 26/66] Remove CC + F2 from tests using SSHT --- tests/test_spherical_base.py | 2 +- tests/test_spherical_custom_grads.py | 7 ++++--- tests/test_spherical_transform.py | 2 +- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/test_spherical_base.py b/tests/test_spherical_base.py index 07aafacf5..399e63a1b 100644 --- a/tests/test_spherical_base.py +++ b/tests/test_spherical_base.py @@ -11,7 +11,7 @@ spin_to_test = [-2, 0, 1] nside_to_test = [4, 5] L_to_nside_ratio = [2, 3] -sampling_to_test = ["mw", "mwss", "dh", "gl", "cc", "f2"] +sampling_to_test = ["mw", "mwss", "dh", "gl"] method_to_test = ["direct", "sov", "sov_fft", "sov_fft_vectorized"] reality_to_test = [False, True] diff --git a/tests/test_spherical_custom_grads.py b/tests/test_spherical_custom_grads.py index d1437d6e2..131a50e03 100644 --- a/tests/test_spherical_custom_grads.py +++ b/tests/test_spherical_custom_grads.py @@ -12,7 +12,8 @@ L_lower_to_test = [2] spin_to_test = [-2, 0, 1] nside_to_test = [8] -sampling_to_test = ["mw", "mwss", "dh", "gl", "cc", "f2"] +sampling_to_test_ssht = ["mw", "mwss", "dh", "gl"] +sampling_to_test = sampling_to_test_ssht + ["cc", "f2"] reality_to_test = [False, True] @@ -214,7 +215,7 @@ def func(f): @pytest.mark.parametrize("L", L_to_test) @pytest.mark.parametrize("L_lower", L_lower_to_test) @pytest.mark.parametrize("spin", spin_to_test) -@pytest.mark.parametrize("sampling", sampling_to_test) +@pytest.mark.parametrize("sampling", sampling_to_test_ssht) @pytest.mark.parametrize("reality", reality_to_test) @pytest.mark.parametrize("_ssht_backend", [0, 1]) @pytest.mark.filterwarnings("ignore::RuntimeWarning") @@ -264,7 +265,7 @@ def func(flm): @pytest.mark.parametrize("L", L_to_test) @pytest.mark.parametrize("L_lower", L_lower_to_test) @pytest.mark.parametrize("spin", spin_to_test) -@pytest.mark.parametrize("sampling", sampling_to_test) +@pytest.mark.parametrize("sampling", sampling_to_test_ssht) @pytest.mark.parametrize("reality", reality_to_test) @pytest.mark.parametrize("_ssht_backend", [0, 1]) @pytest.mark.filterwarnings("ignore::RuntimeWarning") diff --git a/tests/test_spherical_transform.py b/tests/test_spherical_transform.py index d606c287f..73fd3c9d6 100644 --- a/tests/test_spherical_transform.py +++ b/tests/test_spherical_transform.py @@ -15,7 +15,7 @@ L_lower_to_test = [0, 2] spin_to_test = [-2, 0, 1] nside_to_test = [4, 5] -sampling_to_test = ["mw", "mwss", "dh", "gl", "cc", "f2"] +sampling_to_test = ["mw", "mwss", "dh", "gl"] method_to_test = ["numpy", "jax", "torch"] reality_to_test = [False, True] multiple_gpus = [False, True] From 176b82fb66789e73d027afc85a5b4fe0a889fc5c Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Thu, 21 May 2026 17:25:22 +0100 Subject: [PATCH 27/66] Further removing of CC + F2 from tests using SSHT --- tests/test_spherical_precompute.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py index ee1c38379..7ebe799c8 100644 --- a/tests/test_spherical_precompute.py +++ b/tests/test_spherical_precompute.py @@ -20,7 +20,8 @@ spin_to_test = [-2, 0, 6] nside_to_test = [4, 5] L_to_nside_ratio = [2, 3] -sampling_to_test = ["mw", "mwss", "dh", "gl", "cc", "f2"] +sampling_to_test_ssht = ["mw", "mwss", "dh", "gl"] +sampling_to_test = sampling_to_test_ssht + ["cc", "f2"] reality_to_test = [True, False] methods_to_test = ["numpy", "jax", "torch"] recursions_to_test = ["price-mcewen", "risbo", "auto"] @@ -292,7 +293,7 @@ def test_transform_forward_healpix_torch_gradcheck( @pytest.mark.parametrize("spin", [0, 20, 30, -20, -30]) -@pytest.mark.parametrize("sampling", sampling_to_test) +@pytest.mark.parametrize("sampling", sampling_to_test_ssht) @pytest.mark.parametrize("reality", reality_to_test) def test_transform_inverse_high_spin( cached_ssht_test_case: Callable, spin: int, sampling: str, reality: bool @@ -311,7 +312,7 @@ def test_transform_inverse_high_spin( @pytest.mark.parametrize("spin", [0, 20, 30, -20, -30]) -@pytest.mark.parametrize("sampling", sampling_to_test) +@pytest.mark.parametrize("sampling", sampling_to_test_ssht) @pytest.mark.parametrize("reality", reality_to_test) def test_transform_forward_high_spin( cached_ssht_test_case: Callable, spin: int, sampling: str, reality: bool From 832476b2186d0ebb6bfdd6ca44dc60bfe0921cf8 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 09:38:51 +0100 Subject: [PATCH 28/66] Add equiangular sampling scheme set --- s2fft/sampling/s2_samples.py | 1 + 1 file changed, 1 insertion(+) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 0ff60b1ef..f3d209f3d 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -3,6 +3,7 @@ M_OFFSET_1_SCHEMES = frozenset(("mwss", "healpix", "cc", "f2")) INCLUDES_SOUTH_POLE_SCHEMES = frozenset(("mw", "mwss", "cc")) INCLUDES_NORTH_POLE_SCHEMES = frozenset(("mwss", "cc")) +EQUIANGULAR_SCHEMES = frozenset(("mw", "mwss", "dh", "cc", "f2")) def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: From 304bee57f63fcf42de5b43d128604dff1d2902ab Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 09:39:09 +0100 Subject: [PATCH 29/66] Adjust tolerances for spherical precompute tests --- tests/test_spherical_precompute.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py index 7ebe799c8..f54866c39 100644 --- a/tests/test_spherical_precompute.py +++ b/tests/test_spherical_precompute.py @@ -29,7 +29,12 @@ def get_tol(sampling): - return 1e-8 if sampling.lower() in ["dh", "gl"] else 1e-12 + if sampling.lower() in ("dh", "gl"): + return 1e-8 + elif sampling.lower() in ("cc", "f2"): + return 1e-10 + else: + return 1e-12 def check_spin(recursion, spin): From 3756e228352b8980178a85f9a15e4011ecaed3e4 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 09:40:30 +0100 Subject: [PATCH 30/66] Handle CC and F2 schemes in precompute kernel construction --- s2fft/precompute_transforms/construct.py | 84 ++++++++++++------------ 1 file changed, 43 insertions(+), 41 deletions(-) diff --git a/s2fft/precompute_transforms/construct.py b/s2fft/precompute_transforms/construct.py index 52b965ff3..b441a1313 100644 --- a/s2fft/precompute_transforms/construct.py +++ b/s2fft/precompute_transforms/construct.py @@ -13,6 +13,32 @@ PM_MAX_STABLE_SPIN = 6 +def _n_sample_wigner_fourier_inverse_fft(sampling: str, n_theta: int) -> int: + """ + Number of samples for inverse FFT over Wigner Fourier coefficients. + + Args: + sampling: String specifier of sampling scheme. + n_theta: Number of (co)latitude samples in sampling scheme. + + Returns: + Number of samples. + + """ + if sampling == "mw": + return 2 * n_theta - 1 + elif sampling == "mwss": + return 2 * n_theta - 2 + elif sampling == "dh": + return 2 * n_theta + elif sampling == "cc": + return 2 * n_theta - 2 + elif sampling == "f2": + return 2 * n_theta + 2 + else: + raise ValueError(f"Equiangular sampling scheme {sampling} not recognised") + + def spin_spherical_kernel( L: int, spin: int = 0, @@ -39,7 +65,7 @@ def spin_spherical_kernel( Defaults to False. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int): HEALPix Nside resolution parameter. Only required if sampling="healpix". @@ -110,18 +136,12 @@ def spin_spherical_kernel( delta = recursions.risbo.compute_full_vectorised(delta, thetas, L, el) dl[:, el] = delta[:, m_start_ind:, L - 1 - spin] - # MW, MWSS, and DH sampling ARE uniform in theta therefore CAN be calculated + # MW, MWSS, DH, CC & F2 sampling ARE uniform in theta therefore CAN be calculated # using the Fourier decomposition of Wigner d-functions. # - The complexity of this approach is O(L^3LogL). # - This approach is stable for arbitrary abs(spins) <= L. - if sampling.lower() in ["mw", "mwss", "dh"]: - # Number of samples for inverse FFT over Wigner Fourier coefficients. - if sampling.lower() == "mw": - nsamps = 2 * len(thetas) - 1 - elif sampling.lower() == "mwss": - nsamps = 2 * len(thetas) - 2 - elif sampling.lower() == "dh": - nsamps = 2 * len(thetas) + if sampling.lower() in samples.EQUIANGULAR_SCHEMES: + nsamps = _n_sample_wigner_fourier_inverse_fft(sampling.lower(), len(thetas)) delta = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64) # Calculate the Fourier coefficients of the Wigner d-functions, delta(pi/2). @@ -187,7 +207,7 @@ def spin_spherical_kernel_jax( Defaults to False. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int): HEALPix Nside resolution parameter. Only required if sampling="healpix". @@ -243,12 +263,12 @@ def spin_spherical_kernel_jax( dl = jnp.swapaxes(dl, 0, 1) # North pole singularity - if sampling.lower() == "mwss": + if sampling.lower() in samples.INCLUDES_NORTH_POLE_SCHEMES: dl = dl.at[0].set(0) dl = dl.at[0, :, L - 1 - spin].set(1) # South pole singularity - if sampling.lower() in ["mw", "mwss"]: + if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: dl = dl.at[-1].set(0) dl = dl.at[-1, :, L - 1 + spin].set((-1) ** (jnp.arange(L) - spin)) dl = dl.at[:, : jnp.abs(spin)].multiply(0) @@ -277,14 +297,8 @@ def spin_spherical_kernel_jax( # using the Fourier decomposition of Wigner d-functions. # - The complexity of this approach is O(L^3LogL). # - This approach is stable for arbitrary abs(spins) <= L. - elif sampling.lower() in ["mw", "mwss", "dh"]: - # Number of samples for inverse FFT over Wigner Fourier coefficients. - if sampling.lower() == "mw": - nsamps = 2 * len(thetas) - 1 - elif sampling.lower() == "mwss": - nsamps = 2 * len(thetas) - 2 - elif sampling.lower() == "dh": - nsamps = 2 * len(thetas) + elif sampling.lower() in samples.EQUIANGULAR_SCHEMES: + nsamps = _n_sample_wigner_fourier_inverse_fft(sampling.lower(), len(thetas)) delta = jnp.zeros((2 * L - 1, 2 * L - 1), dtype=jnp.float64) # Calculate the Fourier coefficients of the Wigner d-functions, delta(pi/2). @@ -371,7 +385,7 @@ def wigner_kernel( np.ndarray: Transform kernel for Wigner transform. """ - if mode.lower() == "fft" and sampling.lower() not in ["mw", "mwss", "dh"]: + if mode.lower() == "fft" and sampling.lower() not in samples.EQUIANGULAR_SCHEMES: raise ValueError( f"Fourier based recursion is not valid for {sampling} sampling." ) @@ -379,7 +393,7 @@ def wigner_kernel( # - Can only use the FFT approach when uniformly sampling in theta. # - FFT approach is only more efficient when N <= L/Log(L) roughly. if mode.lower() == "auto": - if sampling.lower() in ["mw", "mwss", "dh"]: + if sampling.lower() in samples.EQUIANGULAR_SCHEMES: mode = "fft" if N <= int(L / np.log(L)) else "direct" else: mode = "direct" @@ -410,20 +424,14 @@ def wigner_kernel( delta = recursions.risbo.compute_full_vectorised(delta, thetas, L, el) dl[:, :, el] = np.moveaxis(delta, -1, 0)[L - 1 + n] - # MW, MWSS, and DH sampling ARE uniform in theta therefore CAN be calculated + # MW, MWSS, DH, CC & F2 sampling ARE uniform in theta therefore CAN be calculated # using the Fourier decomposition of Wigner d-functions. # - The complexity of this approach is O(NL^3LogL). # - This approach is stable for arbitrary abs(spins) <= L. # Therefore when NL^3LogL <= L^4 i.e. when N <= L/LogL, the Fourier based approach # is more efficient. This can be a large difference for large L >> N. elif mode.lower() == "fft": - # Number of samples for inverse FFT over Wigner Fourier coefficients. - if sampling.lower() == "mw": - nsamps = 2 * len(thetas) - 1 - elif sampling.lower() == "mwss": - nsamps = 2 * len(thetas) - 2 - elif sampling.lower() == "dh": - nsamps = 2 * len(thetas) + nsamps = _n_sample_wigner_fourier_inverse_fft(sampling.lower(), len(thetas)) delta = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64) # Calculate the Fourier coefficients of the Wigner d-functions, delta(pi/2). @@ -510,7 +518,7 @@ def wigner_kernel_jax( jnp.ndarray: Transform kernel for Wigner transform. """ - if mode.lower() == "fft" and sampling.lower() not in ["mw", "mwss", "dh"]: + if mode.lower() == "fft" and sampling.lower() not in samples.EQUIANGULAR_SCHEMES: raise ValueError( f"Fourier based recursion is not valid for {sampling} sampling." ) @@ -518,7 +526,7 @@ def wigner_kernel_jax( # - Can only use the FFT approach when uniformly sampling in theta. # - FFT approach is only more efficient when N <= L/Log(L) roughly. if mode.lower() == "auto": - if sampling.lower() in ["mw", "mwss", "dh"]: + if sampling.lower() in samples.EQUIANGULAR_SCHEMES: mode = "fft" if N <= int(L / np.log(L)) else "direct" else: mode = "direct" @@ -550,20 +558,14 @@ def wigner_kernel_jax( delta = vfunc(delta, thetas, L, el) dl = dl.at[:, :, el].set(jnp.moveaxis(delta, -1, 0)[L - 1 + n]) - # MW, MWSS, and DH sampling ARE uniform in theta therefore CAN be calculated + # MW, MWSS, DH, CC & F2 sampling ARE uniform in theta therefore CAN be calculated # using the Fourier decomposition of Wigner d-functions. # - The complexity of this approach is O(NL^3LogL). # - This approach is stable for arbitrary abs(spins) <= L. # Therefore when NL^3LogL <= L^4 i.e. when N <= L/LogL, the Fourier based approach # is more efficient. This can be a large difference for large L >> N. elif mode.lower() == "fft": - # Number of samples for inverse FFT over Wigner Fourier coefficients. - if sampling.lower() == "mw": - nsamps = 2 * len(thetas) - 1 - elif sampling.lower() == "mwss": - nsamps = 2 * len(thetas) - 2 - elif sampling.lower() == "dh": - nsamps = 2 * len(thetas) + nsamps = _n_sample_wigner_fourier_inverse_fft(sampling.lower(), len(thetas)) delta = jnp.zeros((2 * L - 1, 2 * L - 1), dtype=jnp.float64) # Calculate the Fourier coefficients of the Wigner d-functions, delta(pi/2). From 0be15f18349aa49ee85e8381946afdaec72c56ee Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 11:32:38 +0100 Subject: [PATCH 31/66] Generalize handling of poles in JAX precompute function --- s2fft/recursions/price_mcewen.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/s2fft/recursions/price_mcewen.py b/s2fft/recursions/price_mcewen.py index 9ee94d037..8defcbe47 100644 --- a/s2fft/recursions/price_mcewen.py +++ b/s2fft/recursions/price_mcewen.py @@ -165,9 +165,14 @@ def generate_precomputes_jax( if betas is None: if forward and sampling.lower() in ["mw", "mwss"]: sampling = "mwss" - beta = samples.thetas(2 * L, "mwss")[1:-1] + beta = samples.thetas(2 * L, "mwss") else: beta = samples.thetas(L, sampling, nside) + + if forward and sampling.lower() in ( + samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES + ): + beta = beta[1:-1] else: beta = betas From bff50d3c83beeecc9c2513173be35489a4a46cfd Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 11:52:52 +0100 Subject: [PATCH 32/66] Account for L=1 case for CC quadrature weights --- s2fft/utils/quadrature.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index ce002cc5a..111904d86 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -420,7 +420,9 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: n_theta = samples.ntheta(L, "cc") - if n_theta == 2: + if n_theta == 1: + return xp.array([2.0]) + elif n_theta == 2: return xp.array([1.0, 1.0]) n = n_theta - 1 From e0c447d2f399136bc3f6ae4e31778b4e1d718271 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 11:53:09 +0100 Subject: [PATCH 33/66] Add additional quadrature weights tests --- tests/test_quadrature.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index ddfde414b..aba7a2448 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -95,3 +95,19 @@ def test_quadrature_polynomial(f_and_integral, rule, n_points, quadrature_module @pytest.mark.parametrize("quadrature_module", [quadrature, quadrature_jax]) def test_quadrature_non_polynomial(f_and_integral, rule, n_points, quadrature_module): check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module) + + +@pytest.mark.parametrize("sampling", ["cc", "f2", "mw", "mwss", "gl", "dh"]) +@pytest.mark.parametrize("L", [1, 2, 3, 8, 10]) +@pytest.mark.parametrize("quadrature_module", [quadrature, quadrature_jax]) +def test_quadrature_weights(sampling, L, quadrature_module): + weights = quadrature_module.quad_weights(L, sampling) + assert weights.shape[0] == samples.ntheta(L, sampling) + # In general quadrature rules may use negative weights but + # for all currently implemented schemes weights should be non-negative + assert np.all(weights > 0) + # Weights are equal for each longitude phi and so only computed for each theta + # and rescaled to account for integration in phi - we expect the total sum of + # weights (accounting for repeating by number of phi points) to be equal to + # surface area of unit sphere - that is 4 * pi + assert np.isclose((weights * samples.nphi_equiang(L, sampling)).sum(), 4 * np.pi) From 4fbdded26bd4c64f07cce6d10fb38ebfeb6b5705 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 13:40:22 +0100 Subject: [PATCH 34/66] Remove unneeded special case for CC rule with n_theta = 2 --- s2fft/utils/quadrature.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index 111904d86..1e6aa3aa2 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -420,10 +420,10 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: n_theta = samples.ntheta(L, "cc") + # Computing RFFT below with n=0 will fail so explicitly return + # weights for single point case if n_theta == 1: return xp.array([2.0]) - elif n_theta == 2: - return xp.array([1.0, 1.0]) n = n_theta - 1 w = _fejer_second_rule_rfft_weights(n, xp) From f6b4145dae39393ff24bdcaafa820432081e8d79 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 13:41:06 +0100 Subject: [PATCH 35/66] Give quadrature test using transform more descriptive name --- tests/test_quadrature.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index aba7a2448..b5caa0d6c 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -12,7 +12,9 @@ @pytest.mark.parametrize("L", [5, 6]) @pytest.mark.parametrize("sampling", ["mw", "mwss", "dh", "gl", "cc", "f2"]) @pytest.mark.parametrize("method", ["numpy", "jax", "torch"]) -def test_quadrature_mw_weights(flm_generator, L: int, sampling: str, method: str): +def test_quadrature_weights_in_transform( + flm_generator, L: int, sampling: str, method: str +): spin = 0 if method.lower() == "numpy": From 71de50f756610a3e6d2f1beedbab7bc770ca1d9e Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 13:41:32 +0100 Subject: [PATCH 36/66] Add further tests for quadrature rule exceptions --- tests/test_quadrature.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index b5caa0d6c..96da11093 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -51,6 +51,10 @@ def test_quadrature_exceptions(): with pytest.raises(ValueError): quadrature.quad_weights(L, sampling="foo") + for sampling in ("cc", "f2"): + with pytest.raises(ValueError, match="must be at least 1"): + quadrature.quad_weights(0, sampling) + def check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module, tol=1e-12): f, true_integral = f_and_integral From e2aea8888a41776d5b4383bf9d7dde10a40c5c9b Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 13:42:18 +0100 Subject: [PATCH 37/66] Add tests for HEALPix quadrature weights --- tests/test_quadrature.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index 96da11093..5a18e7285 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -103,12 +103,13 @@ def test_quadrature_non_polynomial(f_and_integral, rule, n_points, quadrature_mo check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module) -@pytest.mark.parametrize("sampling", ["cc", "f2", "mw", "mwss", "gl", "dh"]) -@pytest.mark.parametrize("L", [1, 2, 3, 8, 10]) +@pytest.mark.parametrize("sampling", ["cc", "f2", "mw", "mwss", "gl", "dh", "healpix"]) +@pytest.mark.parametrize("L_or_nside", [1, 2, 3, 8, 10]) @pytest.mark.parametrize("quadrature_module", [quadrature, quadrature_jax]) -def test_quadrature_weights(sampling, L, quadrature_module): - weights = quadrature_module.quad_weights(L, sampling) - assert weights.shape[0] == samples.ntheta(L, sampling) +def test_quadrature_weights(sampling, L_or_nside, quadrature_module): + L, nside = (None, L_or_nside) if sampling == "healpix" else (L_or_nside, None) + weights = quadrature_module.quad_weights(L=L, sampling=sampling, nside=nside) + assert weights.shape[0] == samples.ntheta(L=L, sampling=sampling, nside=nside) # In general quadrature rules may use negative weights but # for all currently implemented schemes weights should be non-negative assert np.all(weights > 0) @@ -116,4 +117,13 @@ def test_quadrature_weights(sampling, L, quadrature_module): # and rescaled to account for integration in phi - we expect the total sum of # weights (accounting for repeating by number of phi points) to be equal to # surface area of unit sphere - that is 4 * pi - assert np.isclose((weights * samples.nphi_equiang(L, sampling)).sum(), 4 * np.pi) + if sampling == "healpix": + # HEALPix sampling scheme uses equiarea pixels with uniform weights + assert np.all(weights == weights[0]) + assert np.isclose( + weights[0] * samples.f_shape(L, sampling, nside)[0], 4 * np.pi + ) + else: + assert np.isclose( + (weights * samples.nphi_equiang(L, sampling)).sum(), 4 * np.pi + ) From b3efcf73f173e2447f922014add9eb8ed0aba17c Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 13:43:03 +0100 Subject: [PATCH 38/66] Add tests for Wigner kernel exceptions and n sample function --- tests/test_spherical_precompute.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py index f54866c39..139f4d327 100644 --- a/tests/test_spherical_precompute.py +++ b/tests/test_spherical_precompute.py @@ -351,3 +351,30 @@ def test_inverse_transform_unrecognised_method_raises(): flm = np.zeros(samples.flm_shape(L)) with pytest.raises(ValueError, match=f"{method} not recognised"): inverse(flm, L, method=method) + + +@pytest.mark.parametrize("sampling", ("gl", "healpix")) +def test_wigner_kernel_fft_mode_non_equiangular_scheme_raises(sampling): + mode = "fft" + L = 1 + N = 1 + for wigner_kernel_function in (c.wigner_kernel_jax, c.wigner_kernel): + with pytest.raises(ValueError, match=f"not valid for {sampling} sampling"): + wigner_kernel_function(L, N, sampling=sampling, mode=mode) + + +@pytest.mark.parametrize("sampling", samples.EQUIANGULAR_SCHEMES) +@pytest.mark.parametrize("L", [1, 2, 5]) +def test_n_sample_wigner_fourier_inverse_fft(sampling, L): + n_sample = c._n_sample_wigner_fourier_inverse_fft( + sampling, samples.ntheta(L, sampling) + ) + assert isinstance(n_sample, int) + assert n_sample > 0 + + +def test_n_sample_wigner_fourier_inverse_fft_raises(): + sampling = "invalid_sampling" + n_theta = 1 + with pytest.raises(ValueError, match=f"{sampling} not recognised"): + c._n_sample_wigner_fourier_inverse_fft(sampling, n_theta) From 9049377c49d1debd748ad8aa28de67a213fdc659 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 14:36:47 +0100 Subject: [PATCH 39/66] Relax test condition on number of samples to non-negativity --- tests/test_spherical_precompute.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py index 139f4d327..fe6f1c624 100644 --- a/tests/test_spherical_precompute.py +++ b/tests/test_spherical_precompute.py @@ -370,7 +370,7 @@ def test_n_sample_wigner_fourier_inverse_fft(sampling, L): sampling, samples.ntheta(L, sampling) ) assert isinstance(n_sample, int) - assert n_sample > 0 + assert n_sample >= 0 def test_n_sample_wigner_fourier_inverse_fft_raises(): From 53573edf5667ce3e9c73cf964fb8c7d57d1691a8 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 16:50:29 +0100 Subject: [PATCH 40/66] Add CC and F2 sampling scheme details to docs --- docs/sampling.rst | 62 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/docs/sampling.rst b/docs/sampling.rst index f74be94ff..2d7819858 100644 --- a/docs/sampling.rst +++ b/docs/sampling.rst @@ -58,6 +58,18 @@ We denote by $L$ the band-limit of the signals we are considering. - No - Yes - No + * - :ref:`clenshaw-curtis-cc` + - ``"cc"`` + - :math:`(2L - 1) \times 2L` + - Yes + - No + - Yes + * - :ref:`fejer-rule-2-f2` + - ``"f2"`` + - :math:`(2L - 1) \times 2L` + - Yes + - No + - Yes Specifying sampling schemes in ``S2FFT`` ---------------------------------------- @@ -214,6 +226,56 @@ For the equations defining the exact positioning of the regions, their centres, Further information; `Gorski et al. (2005) `_. +.. _clenshaw-curtis-cc: + +Clenshaw-Curtis (CC) +-------------------- + +Clenshaw-Curtis quadrature `(Cleshaw and Curtis, 1960) `_, +is based on expansion of the integrand in terms of Chebyshev polynomials. +It can be used in a spherical harmonic transform setting to numerically compute integrals over the co-latitude, +:math:`\theta`, dimension. +Clenshaw-Curtis quadrature requires roughly double the number of nodes to exactly integrate a polynomial of a given degree +compared to Gauss-Legendre quadrature, +however the co-latitude nodes are equispaced and grids of different resolutions can be nested. + +Sample positions are defined by + +.. math:: + + \theta_t &= \frac{\pi t}{2L-2}, &\quad t\in\lbrace 0,1,...,2L-2 \rbrace, \\ + \varphi_p &= \frac{2\pi p}{2L}, &\quad p\in\lbrace 0,1,...,2L-1\rbrace. + +Note that here we follow the original definition of the Clenshaw-Curtis quadrature rule +and include samples at both poles. +An analogous scheme which excludes the pole is available as :ref:`F2 `. +Due to the redundancy in samples at the poles, +the total number of distinct sites on the sphere used by this sampling scheme is :math:`(2L-3)(2L)+2 = 4L^2 -6L + 2`. + +For further information see `Hotte and Ujiie (2018) `_. + +.. _fejer-rule-2-f2: + +Fejér's second rule (F2) +------------------------ + +Fejér quadrature `(Fejer, 1933) `_, +is closely related to Clenshaw-Curtis quadrature. +Fejér's second rule can be used to define a sampling scheme analagous to :ref:`CC ` but with the poles excluded. + +Sample positions are defined by + +.. math:: + + \theta_t &= \frac{\pi(t + 1)}{2L}, &\quad t\in\lbrace 0,1,...,2L-2 \rbrace, \\ + \varphi_p &= \frac{2\pi p}{2L}, &\quad p\in\lbrace 0,1,...,2L-1\rbrace. + +This requires :math:`(2L-1) \times 2L = 4L^2 - 2L` elements to be held in memory, +and since the poles are not sample points, +the same number of independent degrees of freedom to represent the signal. + +For further information see `Hotte and Ujiie (2018) `_. + .. rubric:: Footnotes .. [#n-samples-vs-memory-storage] The "number of sample points" listed here refers to the size of the array recording sample values in memory. From 3447e4875434af2709442b80765b4cd259a20c5f Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Fri, 22 May 2026 16:50:49 +0100 Subject: [PATCH 41/66] Minor fixes to sampling scheme doc page formatting --- docs/sampling.rst | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/docs/sampling.rst b/docs/sampling.rst index 2d7819858..cae3b8b67 100644 --- a/docs/sampling.rst +++ b/docs/sampling.rst @@ -9,10 +9,10 @@ Sampling schemes The structure of the algorithms implemented in ``S2FFT`` can support a number of sampling schemes, which we give a brief overview of here. An at-a-glance summary of the differences between the supported sampling schemes is also provided :ref:`in the table below `, with further information available in the dedicated section for each scheme. -A more thorough overview of the schemes can be found in section 4.2 of `Price & McEwen (2025) `_. +A more thorough overview of most of the schemes can be found in section 4.2 of `Price & McEwen (2025) `_. -We adopt the usual ``S2FFT`` conventions for spherical coordinates; :math:`\theta\in[0, \pi]` (colatitude) and :math:`\varphi\in[0,2\pi)` (longitude), with :math:`\theta_t` and :math:`\varphi_p` being the discretised samples (indexed by $t$ and $p$) drawn by the sampling scheme. -We denote by $L$ the band-limit of the signals we are considering. +We adopt the usual ``S2FFT`` conventions for spherical coordinates; :math:`\theta\in[0, \pi]` (colatitude) and :math:`\varphi\in[0,2\pi)` (longitude), with :math:`\theta_t` and :math:`\varphi_p` being the discretised samples (indexed by :math:`t` and :math:`p`) drawn by the sampling scheme. +We denote by :math:`L` the band-limit of the signals we are considering. .. _sampling-comparison-table: @@ -20,12 +20,12 @@ We denote by $L$ the band-limit of the signals we are considering. :header-rows: 1 :align: center :width: 95 - :widths: 20 10 20 20 15 15 + :widths: 20 10 30 15 15 15 * - Scheme - API string - Number of sample points [#n-samples-vs-memory-storage]_ - - Equi- angular + - Equi-angular - Equal region area - Sampling theorem * - :ref:`mcewen-wiaux-mw` @@ -54,7 +54,7 @@ We denote by $L$ the band-limit of the signals we are considering. - Yes * - :ref:`healpix` - ``"healpix"`` - - $12 N_{side}^2$ + - :math:`12 N_{side}^2` - No - Yes - No @@ -159,11 +159,11 @@ Further information; `McEwen & Wiaux (2012) `_. .. _mcewen-wiaux-mwss: -McEwen & Wiaux with Symmetric Sampling (MWSS) ---------------------------------------------- +McEwen & Wiaux Symmetric Sampling (MWSS) +---------------------------------------- -This sampling scheme uses slightly more samples than MW, requiring an array holding :math:`(L+1)\times 2L` elements (with $2(L^2 - L + 1)$ independent degrees of freedom). -Asymptotically, we still only require :math:`\sim 2L^2` elements in memory as $L$ increases. +This sampling scheme uses slightly more samples than MW, requiring an array holding :math:`(L+1)\times 2L` elements (with :math:`2(L^2 - L + 1)` independent degrees of freedom). +Asymptotically, we still only require :math:`\sim 2L^2` elements in memory as :math:`L` increases. In exchange for slightly higher memory usage, the sample locations possess antipodal symmetry. Sample positions are defined by @@ -199,7 +199,7 @@ Gauss-Legendre (GL) The GL sampling theorem also requires an array of :math:`L\times (2L-1) \sim 2L^2` elements to represent the signal. Like :ref:`DH `, there is no redundancy in samples at the poles, so the same number of independent degrees of freedom are needed. -The :math:`\theta_t` are determined by the roots of the Legendre polynomials of order $L$, whilst the :math:`\varphi_p` are defined by +The :math:`\theta_t` are determined by the roots of the Legendre polynomials of order :math:`L`, whilst the :math:`\varphi_p` are defined by .. math:: @@ -219,9 +219,9 @@ However, HEALPix sampling **does not** exhibit a sampling theorem and so round-t An `iterative refinement `_ scheme can be applied to the forward transform to reduce this round-trip error at the cost of additional computation. This can be applied in ``S2FFT``'s forward transforms by setting the `iter` argument to the number of iterations to perform, with more iterations giving a smaller round-trip error. -A HEALPix grid is defined by a resolution parameter $N_{side}$, requiring $12 N_{side}^2$ elements (and independent degrees of freedom) stored in memory. -Given a resolution parameter, the grid will contain $N_{hp} = 12 N_{side}^2$ regions of the same area :math:`\frac{\pi}{3N_{side}^2}`. -The regions will be laid out on $4N_{side}-1$ iso-latitude rings, and the distribution of regions will be symmetric about the equator. +A HEALPix grid is defined by a resolution parameter :math:`N_{side}`, requiring :math:`12 N_{side}^2` elements (and independent degrees of freedom) stored in memory. +Given a resolution parameter, the grid will contain :math:`N_{hp} = 12 N_{side}^2` regions of the same area :math:`\frac{\pi}{3N_{side}^2}`. +The regions will be laid out on :math:`4N_{side}-1` iso-latitude rings, and the distribution of regions will be symmetric about the equator. For the equations defining the exact positioning of the regions, their centres, their boundaries, and how they are organised into an array, see section 5 of `Gorski et al. (2005) `_. Further information; `Gorski et al. (2005) `_. From 6adbe189cd96c5a79742f26e8b9b1e313afe74d5 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 12 Aug 2026 08:54:58 +0100 Subject: [PATCH 42/66] Add F2 sampling scheme to set in docstring --- s2fft/utils/quadrature.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index 1e6aa3aa2..f31b43bc0 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -21,7 +21,7 @@ def quad_weights_transform( L (int): Harmonic band-limit. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mwss", "dh", "gl", "healpix", "cc"}. Defaults to "mwss". + {"mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mwss". spin (int, optional): Harmonic spin. Defaults to 0. From c56a174ed096de3ecff3b208155dbab4d0aa16b5 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 12 Aug 2026 09:48:44 +0100 Subject: [PATCH 43/66] Update m offset logic and docstring to be more explicit --- s2fft/base_transforms/spherical.py | 8 +++---- s2fft/precompute_transforms/spherical.py | 8 +++---- s2fft/precompute_transforms/wigner.py | 8 +++---- s2fft/sampling/s2_samples.py | 28 +++++++++++++++++++----- s2fft/sampling/so3_samples.py | 11 +--------- s2fft/transforms/otf_recursions.py | 12 +++++----- s2fft/transforms/spherical.py | 8 +++---- 7 files changed, 46 insertions(+), 37 deletions(-) diff --git a/s2fft/base_transforms/spherical.py b/s2fft/base_transforms/spherical.py index adbf5f40b..cba443709 100644 --- a/s2fft/base_transforms/spherical.py +++ b/s2fft/base_transforms/spherical.py @@ -486,7 +486,7 @@ def _compute_inverse_sov_fft( assert L >= 2 * nside ftm = np.zeros(samples.ftm_shape(L, sampling, nside), dtype=np.complex128) - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) for t, theta in enumerate(thetas): phi_ring_offset = ( @@ -576,7 +576,7 @@ def _compute_inverse_sov_fft_vectorized( """ ftm = np.zeros(samples.ftm_shape(L, sampling, nside), dtype=np.complex128) - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) for t, theta in enumerate(thetas): phase_shift = ( @@ -844,7 +844,7 @@ def _compute_forward_sov_fft( flm = np.zeros(samples.flm_shape(L), dtype=np.complex128) ftm = np.zeros_like(f).astype(np.complex128) - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) if sampling.lower() == "healpix": ftm = hp.healpix_fft(f, L, nside, "numpy", reality) @@ -960,7 +960,7 @@ def _compute_forward_sov_fft_vectorized( flm = np.zeros(samples.flm_shape(L), dtype=np.complex128) ftm = np.zeros_like(f).astype(np.complex128) - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) if reality: m_conj = (-1) ** (np.arange(1, L) % 2) diff --git a/s2fft/precompute_transforms/spherical.py b/s2fft/precompute_transforms/spherical.py index 4668b5b2e..792269a76 100644 --- a/s2fft/precompute_transforms/spherical.py +++ b/s2fft/precompute_transforms/spherical.py @@ -120,7 +120,7 @@ def inverse_transform( np.ndarray: Pixel-space coefficients. """ - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) m_start_ind = L - 1 if reality else 0 ftm = np.zeros(samples.ftm_shape(L, sampling, nside), dtype=np.complex128) @@ -186,7 +186,7 @@ def inverse_transform_jax( jnp.ndarray: Pixel-space coefficients with shape. """ - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) m_start_ind = L - 1 if reality else 0 ftm = jnp.zeros( @@ -356,7 +356,7 @@ def forward_transform( sampling = "mwss" f = resampling.upsample_by_two_mwss(f, L, spin) - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) m_start_ind = L - 1 if reality else 0 if sampling.lower() == "healpix": @@ -426,7 +426,7 @@ def forward_transform_jax( sampling = "mwss" f = resampling_jax.upsample_by_two_mwss(f, L, spin) - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) m_start_ind = L - 1 if reality else 0 if sampling.lower() == "healpix": diff --git a/s2fft/precompute_transforms/wigner.py b/s2fft/precompute_transforms/wigner.py index 91c452472..3f6ec4518 100644 --- a/s2fft/precompute_transforms/wigner.py +++ b/s2fft/precompute_transforms/wigner.py @@ -113,7 +113,7 @@ def inverse_transform( np.ndarray: Pixel-space coefficients. """ - m_offset = 1 if sampling in s2_samples.M_OFFSET_1_SCHEMES else 0 + m_offset = s2_samples.m_offset(L, sampling) n_start_ind = N - 1 if reality else 0 fnab = np.zeros(samples.fnab_shape(L, N, sampling, nside), dtype=np.complex128) @@ -175,7 +175,7 @@ def inverse_transform_jax( jnp.ndarray: Pixel-space coefficients. """ - m_offset = 1 if sampling in s2_samples.M_OFFSET_1_SCHEMES else 0 + m_offset = s2_samples.m_offset(L, sampling) n_start_ind = N - 1 if reality else 0 fnab = jnp.zeros(samples.fnab_shape(L, N, sampling, nside), dtype=jnp.complex128) @@ -340,7 +340,7 @@ def forward_transform( sampling = "mwss" fban = resampling.upsample_by_two_mwss(fban, L, spins) - m_offset = 1 if sampling in s2_samples.M_OFFSET_1_SCHEMES else 0 + m_offset = s2_samples.m_offset(L, sampling) if sampling.lower() in "healpix": temp = np.zeros(samples.fnab_shape(L, N, sampling, nside), dtype=np.complex128) @@ -423,7 +423,7 @@ def forward_transform_jax( sampling = "mwss" fban = resampling_jax.upsample_by_two_mwss(fban, L, spins) - m_offset = 1 if sampling in s2_samples.M_OFFSET_1_SCHEMES else 0 + m_offset = s2_samples.m_offset(L, sampling) if sampling.lower() in "healpix": temp = jnp.zeros( diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index f3d209f3d..8c9b34ff2 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -1,6 +1,5 @@ import numpy as np -M_OFFSET_1_SCHEMES = frozenset(("mwss", "healpix", "cc", "f2")) INCLUDES_SOUTH_POLE_SCHEMES = frozenset(("mw", "mwss", "cc")) INCLUDES_NORTH_POLE_SCHEMES = frozenset(("mwss", "cc")) EQUIANGULAR_SCHEMES = frozenset(("mw", "mwss", "dh", "cc", "f2")) @@ -119,10 +118,7 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int: if sampling.lower() in ("mw", "dh", "gl"): return 2 * L - 1 - elif sampling.lower() == "mwss": - return 2 * L - - elif sampling.lower() in ("cc", "f2"): + elif sampling.lower() in ("mwss", "cc", "f2"): return 2 * L elif sampling.lower() == "healpix": @@ -134,6 +130,28 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int: return 1 +def m_offset(L: int, sampling: str) -> int: + """ + Offset for `m` index into intermediate `ftm` array. + + Accounts for parity of number of longitudinal samples `n_phi` and whether this + corresponds to an oversampling in `phi` axis. + + Args: + L (int): Harmonic band-limit. + + sampling (str, optional): Sampling scheme. Supported sampling schemes include + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. + + Returns: + Integer binary (0/1) offset. + + """ + return ( + 1 if sampling.lower() == "healpix" or nphi_equiang(L, sampling) % 2 == 0 else 0 + ) + + def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> tuple[int, int]: r""" Shape of intermediate array, before/after latitudinal step. diff --git a/s2fft/sampling/so3_samples.py b/s2fft/sampling/so3_samples.py index 3796c853c..6d583bdcd 100644 --- a/s2fft/sampling/so3_samples.py +++ b/s2fft/sampling/so3_samples.py @@ -90,16 +90,7 @@ def fnab_shape( :math:`SO(3)`. """ - if sampling.lower() in samples.M_OFFSET_1_SCHEMES: - return _ngamma(N), samples.ntheta(L, sampling, nside), 2 * L - - elif sampling.lower() in ["mw", "dh", "gl"]: - return _ngamma(N), samples.ntheta(L, sampling, nside), 2 * L - 1 - - else: - raise ValueError(f"Sampling scheme sampling={sampling} not supported") - - return 1 + return _ngamma(N), *samples.ftm_shape(L, sampling, nside) def flmn_shape_1d(L: int, N: int) -> int: diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py index fb9151523..eea3fb6da 100644 --- a/s2fft/transforms/otf_recursions.py +++ b/s2fft/transforms/otf_recursions.py @@ -94,7 +94,7 @@ def inverse_latitudinal_step( for i in range(2): if not (reality and i == 0): - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES and i == 0 else 0 + m_offset = samples.m_offset(L, sampling) if i == 0 else 0 lind = L - 1 - L_lower sind = lims[i] @@ -259,7 +259,7 @@ def inverse_latitudinal_step_jax( for i in range(2): if not (reality and i == 0): - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES and i == 0 else 0 + m_offset = samples.m_offset(L, sampling) if i == 0 else 0 lind = L - 1 - L_lower sind = lims[i] @@ -411,7 +411,7 @@ def eval_recursion_step( ) # Remove south pole singularity - m_offset = 1 if sampling.lower() in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: ftm = ftm.at[-1].set(0) ftm = ftm.at[-1, L - 1 + spin + m_offset].set( @@ -505,7 +505,7 @@ def forward_latitudinal_step( for i in range(2): if not (reality and i == 0): - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES and i == 0 else 0 + m_offset = samples.m_offset(L, sampling) if i == 0 else 0 lind = L - 1 - L_lower sind = lims[i] @@ -690,7 +690,7 @@ def forward_latitudinal_step_jax( for i in range(2): if not (reality and i == 0): - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES and i == 0 else 0 + m_offset = samples.m_offset(L, sampling) if i == 0 else 0 lind = L - 1 - L_lower sind = lims[i] @@ -879,7 +879,7 @@ def eval_recursion_step( ) # Include both pole singularities explicitly - m_offset = 1 if sampling.lower() in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: flm = flm.at[L_lower:, L - 1 + spin].add( (-1) ** abs(jnp.arange(L_lower, L) - spin) diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py index fb7553ef5..01119ece8 100644 --- a/s2fft/transforms/spherical.py +++ b/s2fft/transforms/spherical.py @@ -156,7 +156,7 @@ def inverse_numpy( """ # Define latitudinal sample positions and Fourier offsets thetas = samples.thetas(L, sampling, nside) - m_offset = 1 if sampling.lower() in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) m_start_ind = L - 1 if reality else 0 L0 = L_lower @@ -275,7 +275,7 @@ def inverse_jax( """ # Define latitudinal sample positions and Fourier offsets thetas = samples.thetas(L, sampling, nside) - m_offset = 1 if sampling.lower() in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) m_start_ind = L - 1 if reality else 0 # Apply harmonic normalisation @@ -501,7 +501,7 @@ def forward_numpy( # Define latitudinal sample positions and Fourier offsets weights = quadrature.quad_weights_transform(L, sampling, 0, nside) - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) m_start_ind = L - 1 if reality else 0 L0 = L_lower @@ -649,7 +649,7 @@ def forward_jax( # Define latitudinal sample positions and Fourier offsets weights = quadrature_jax.quad_weights_transform(L, sampling, nside) - m_offset = 1 if sampling in samples.M_OFFSET_1_SCHEMES else 0 + m_offset = samples.m_offset(L, sampling) m_start_ind = L - 1 if reality else 0 # Perform longitundal Fast Fourier Transforms From 52bcc304bdfd3536b24e272b86993076d203ac67 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 12 Aug 2026 15:21:27 +0100 Subject: [PATCH 44/66] Ensure polynomial quadrature tests check exactness at critical number of points --- tests/test_quadrature.py | 71 +++++++++++++++++++++++++++++++++++----- 1 file changed, 62 insertions(+), 9 deletions(-) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index 5a18e7285..1c219fa56 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -56,9 +56,50 @@ def test_quadrature_exceptions(): quadrature.quad_weights(0, sampling) +def _bandlimit_from_n_theta(n_theta, rule): + # Compute bandlimit L such that ntheta(L, rule) in {n_theta, n_theta + 1} + match rule: + case "mw" | "gl": + return n_theta + case "f2" | "cc": + return n_theta // 2 + 1 + case "mwss": + return n_theta - 1 + case "dh": + return (n_theta + 1) // 2 + + +def _min_n_points_exact(rule, polynomial_degree): + # Minimum number of quadrature points to exactly integrate polynomial of + # specified degree for a given rule + match rule: + case "gl": + return polynomial_degree // 2 + 1 + case "f2" | "cc": + return polynomial_degree + case "dh" | "mw": + return polynomial_degree + 1 + case "mwss": + return polynomial_degree + 2 + + +def _min_n_points(rule): + # Rules including both poles need a minimum of 2 points + if rule in ( + samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES + ): + return 2 + else: + return 1 + + def check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module, tol=1e-12): f, true_integral = f_and_integral - thetas = samples.thetas(n_points, sampling=rule) + L = _bandlimit_from_n_theta(n_points, rule) + # As ntheta may be constrained to be an odd or even number given L, + # we allow that ntheta may be equal to n_points or one more + assert samples.ntheta(L, rule) in (n_points, n_points + 1) + thetas = samples.thetas(L, sampling=rule) xs = np.cos(thetas) weights = { "cc": quadrature_module.quad_weights_cc_theta_only, @@ -67,24 +108,36 @@ def check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module, tol "mwss": quadrature_module.quad_weights_mwss_theta_only, "gl": quadrature_module.quad_weights_gl_theta_only, "dh": quadrature_module.quad_weights_dh_theta_only, - }[rule](n_points) + }[rule](L) quad_integral = (f(xs) * weights).sum() assert abs(quad_integral - true_integral) < tol @pytest.mark.parametrize( - "f_and_integral", + "f_and_integral_and_degree", [ - (lambda x: x, 0.0), - (lambda x: x**2, 2 / 3), - (lambda x: x**3 - 2 * x**2 + x - 1, -10 / 3), - (lambda x: x**4 - x**2, -4 / 15), + (lambda x: 0.0, 0.0, 0), + (lambda x: x + 1, 2.0, 1), + (lambda x: x**2, 2 / 3, 2), + (lambda x: x**3 - 2 * x**2 + x - 1, -10 / 3, 3), + (lambda x: x**4 - x**2, -4 / 15, 4), + (lambda x: x**8 - x**4 + x**3 - 2 * x**2 + x - 1, -158 / 45, 8), ], ) @pytest.mark.parametrize("rule", ["cc", "f2", "mw", "mwss", "gl", "dh"]) -@pytest.mark.parametrize("n_points", [6, 8, 16]) +@pytest.mark.parametrize("n_points_offset", [0, 1, 2, 4]) @pytest.mark.parametrize("quadrature_module", [quadrature, quadrature_jax]) -def test_quadrature_polynomial(f_and_integral, rule, n_points, quadrature_module): +def test_quadrature_polynomial( + f_and_integral_and_degree, rule, n_points_offset, quadrature_module +): + *f_and_integral, polynomial_degree = f_and_integral_and_degree + # Number of quadrature points defined relative to minimum number + # required to exactly integrate polynomial of specified degree + # to ensure we test this minimum (when n_points_offset == 0) + n_points = max( + _min_n_points_exact(rule, polynomial_degree) + n_points_offset, + _min_n_points(rule), + ) check_quadrature_rule(f_and_integral, rule, n_points, quadrature_module) From b4d8bb760708fddbb550b41d9456cfbbafa68f11 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 09:17:06 +0100 Subject: [PATCH 45/66] Factor out constant for schemes with both poles --- s2fft/recursions/price_mcewen.py | 8 ++------ s2fft/sampling/s2_samples.py | 1 + s2fft/transforms/otf_recursions.py | 4 +--- s2fft/transforms/spherical.py | 8 ++------ tests/test_quadrature.py | 4 +--- 5 files changed, 7 insertions(+), 18 deletions(-) diff --git a/s2fft/recursions/price_mcewen.py b/s2fft/recursions/price_mcewen.py index 8defcbe47..e5ea6dbd6 100644 --- a/s2fft/recursions/price_mcewen.py +++ b/s2fft/recursions/price_mcewen.py @@ -59,9 +59,7 @@ def generate_precomputes( else: beta = samples.thetas(L, sampling, nside) - if forward and sampling.lower() in ( - samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES - ): + if forward and sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: beta = beta[1:-1] ntheta = len(beta) # Number of theta samples @@ -169,9 +167,7 @@ def generate_precomputes_jax( else: beta = samples.thetas(L, sampling, nside) - if forward and sampling.lower() in ( - samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES - ): + if forward and sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: beta = beta[1:-1] else: beta = betas diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 8c9b34ff2..08ea4be66 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -2,6 +2,7 @@ INCLUDES_SOUTH_POLE_SCHEMES = frozenset(("mw", "mwss", "cc")) INCLUDES_NORTH_POLE_SCHEMES = frozenset(("mwss", "cc")) +INCLUDES_BOTH_POLES_SCHEMES = INCLUDES_SOUTH_POLE_SCHEMES & INCLUDES_NORTH_POLE_SCHEMES EQUIANGULAR_SCHEMES = frozenset(("mw", "mwss", "dh", "cc", "f2")) diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py index eea3fb6da..2012198d5 100644 --- a/s2fft/transforms/otf_recursions.py +++ b/s2fft/transforms/otf_recursions.py @@ -654,9 +654,7 @@ def forward_latitudinal_step_jax( """ # Avoid pole-singularities for MWSS sampling - if sampling.lower() in ( - samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES - ): + if sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: ftm = ftm_in[1:-1] beta = beta_in[1:-1] elif sampling.lower() in ( diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py index 01119ece8..e51fe297e 100644 --- a/s2fft/transforms/spherical.py +++ b/s2fft/transforms/spherical.py @@ -526,12 +526,8 @@ def forward_numpy( phase_shifts = hp.ring_phase_shifts_hp(L, nside, True, reality) ftm[:, m_start_ind + m_offset :] *= phase_shifts - INCLUDES_BOTH_POLES_SCHEMES = ( - samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES - ) - # Perform latitudinal wigner-d recursions - if sampling.lower() in INCLUDES_BOTH_POLES_SCHEMES: + if sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: flm = otf.forward_latitudinal_step( ftm[1:-1], thetas[1:-1], @@ -549,7 +545,7 @@ def forward_numpy( ) # Include both pole singularities explicitly - if sampling.lower() in INCLUDES_BOTH_POLES_SCHEMES: + if sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: flm[L0:, L - 1 + spin] += (-1) ** abs(np.arange(L0, L) - spin) * ftm[ -1, L - 1 + spin + m_offset ] diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index 1c219fa56..352d887d8 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -85,9 +85,7 @@ def _min_n_points_exact(rule, polynomial_degree): def _min_n_points(rule): # Rules including both poles need a minimum of 2 points - if rule in ( - samples.INCLUDES_NORTH_POLE_SCHEMES & samples.INCLUDES_SOUTH_POLE_SCHEMES - ): + if rule in samples.INCLUDES_BOTH_POLES_SCHEMES: return 2 else: return 1 From 3e684ec9ccfa648b83c724891c81ed52a38e7c40 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 09:22:13 +0100 Subject: [PATCH 46/66] Fix inconsistencies in sampling scheme lists in docstrings --- s2fft/sampling/s2_samples.py | 4 ++-- s2fft/utils/quadrature_jax.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 08ea4be66..0f0293f50 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -239,7 +239,7 @@ def thetas(L: int = None, sampling: str = "mw", nside: int = None) -> np.ndarray Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix", "cc"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. @@ -269,7 +269,7 @@ def t2theta( Defaults to None. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mw", "mwss", "dh", "gl", "healpix", "cc"}. Defaults to "mw". + {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mw". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py index 03f5cf355..b60f7d9f9 100644 --- a/s2fft/utils/quadrature_jax.py +++ b/s2fft/utils/quadrature_jax.py @@ -35,7 +35,7 @@ def quad_weights_transform( L (int): Harmonic band-limit. sampling (str, optional): Sampling scheme. Supported sampling schemes include - {"mwss", "dh", "gl", "healpix", "cc"}. Defaults to "mwss". + {"mwss", "dh", "gl", "healpix", "cc", "f2"}. Defaults to "mwss". nside (int, optional): HEALPix Nside resolution parameter. Only required if sampling="healpix". Defaults to None. From c97fdb481bc52a7dddf5b6dcae371f7c37370361 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 09:22:26 +0100 Subject: [PATCH 47/66] Avoid duplication in sampling scheme list --- s2fft/sampling/s2_samples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 0f0293f50..b3b0432f5 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -426,7 +426,7 @@ def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray: np.ndarray: :math:`\phi` sample(s) for given sampling scheme. """ - if sampling.lower() in ["mw", "gl", "dh", "mw", "mwss", "cc", "f2"]: + if sampling.lower() in ["mw", "gl", "dh", "mwss", "cc", "f2"]: return 2 * p * np.pi / nphi_equiang(L, sampling) else: raise ValueError(f"Sampling scheme sampling={sampling} not supported") From 27cb44087b9e80bd4973a3ef308253cfac9ddf56 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 09:27:21 +0100 Subject: [PATCH 48/66] Add explanatory comments to sampling scheme set constants --- s2fft/sampling/s2_samples.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index b3b0432f5..3b4abe43c 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -1,8 +1,11 @@ import numpy as np +# Sampling schemes which include south pole (θ = π) - handling of pole singularity needs to be special cased INCLUDES_SOUTH_POLE_SCHEMES = frozenset(("mw", "mwss", "cc")) +# Sampling schemes which include north pole (θ = 0) - handling of pole singularity needs to be special cased INCLUDES_NORTH_POLE_SCHEMES = frozenset(("mwss", "cc")) INCLUDES_BOTH_POLES_SCHEMES = INCLUDES_SOUTH_POLE_SCHEMES & INCLUDES_NORTH_POLE_SCHEMES +# Sampling schemes for which the spacing between pairs of co-latitudes (θ) is constant EQUIANGULAR_SCHEMES = frozenset(("mw", "mwss", "dh", "cc", "f2")) From a17772f6498ba0d734fae9474724008d469c11f4 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 09:33:41 +0100 Subject: [PATCH 49/66] Raise early error for L=1 with sampling=cc --- s2fft/sampling/s2_samples.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 3b4abe43c..16be6758f 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -49,6 +49,8 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: return L + 1 elif sampling.lower() == "cc": + if L == 1: + raise ValueError("Sampling scheme sampling='cc' requires at least L=2") return 2 * L - 1 elif sampling.lower() == "dh": From 5d6d078d500fdb5df6788719dfae435be185a7cc Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 09:39:50 +0100 Subject: [PATCH 50/66] Add test for CC sampling L=1 raising error --- tests/test_spherical_transform.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_spherical_transform.py b/tests/test_spherical_transform.py index 73fd3c9d6..414ff9378 100644 --- a/tests/test_spherical_transform.py +++ b/tests/test_spherical_transform.py @@ -233,3 +233,12 @@ def test_sampling_exceptions(flm_generator): with pytest.raises(ValueError): spherical.forward(None, 0, 0, None, method="incorrect") + + +@pytest.mark.parametrize("reality", [True, False]) +@pytest.mark.parametrize("method", ["jax", "numpy"]) +def test_cc_with_L_1_raises_error(flm_generator, reality, method): + L = 1 + flm = flm_generator(L=L, reality=reality) + with pytest.raises(ValueError, match="at least L=2"): + spherical.inverse(flm, L=L, sampling="cc", reality=reality, method=method) From 61c80a2b391278d1e691d851a78f32cd16c6966d Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 10:50:23 +0100 Subject: [PATCH 51/66] Skip L=1 and sampling=cc test configurations --- tests/test_quadrature.py | 2 ++ tests/test_spherical_precompute.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index 352d887d8..d08812cd9 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -159,6 +159,8 @@ def test_quadrature_non_polynomial(f_and_integral, rule, n_points, quadrature_mo @pytest.mark.parametrize("quadrature_module", [quadrature, quadrature_jax]) def test_quadrature_weights(sampling, L_or_nside, quadrature_module): L, nside = (None, L_or_nside) if sampling == "healpix" else (L_or_nside, None) + if sampling == "cc" and L == 1: + pytest.skip("sampling='cc' requires L > 1") weights = quadrature_module.quad_weights(L=L, sampling=sampling, nside=nside) assert weights.shape[0] == samples.ntheta(L=L, sampling=sampling, nside=nside) # In general quadrature rules may use negative weights but diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py index fe6f1c624..bb21ae6cb 100644 --- a/tests/test_spherical_precompute.py +++ b/tests/test_spherical_precompute.py @@ -366,6 +366,8 @@ def test_wigner_kernel_fft_mode_non_equiangular_scheme_raises(sampling): @pytest.mark.parametrize("sampling", samples.EQUIANGULAR_SCHEMES) @pytest.mark.parametrize("L", [1, 2, 5]) def test_n_sample_wigner_fourier_inverse_fft(sampling, L): + if sampling == "cc" and L == 1: + pytest.skip("sampling='cc' requires L > 1") n_sample = c._n_sample_wigner_fourier_inverse_fft( sampling, samples.ntheta(L, sampling) ) From 74f37f3116f6be071a7d7ec543ab8c2eba0c29bf Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 11:05:04 +0100 Subject: [PATCH 52/66] Change guard conditions around L<=1 case for sampling='cc' and remove dead code --- s2fft/sampling/s2_samples.py | 2 +- s2fft/utils/quadrature.py | 14 ++++---------- tests/test_quadrature.py | 5 ++--- 3 files changed, 7 insertions(+), 14 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 16be6758f..a0c146247 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -49,7 +49,7 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: return L + 1 elif sampling.lower() == "cc": - if L == 1: + if L < 2: raise ValueError("Sampling scheme sampling='cc' requires at least L=2") return 2 * L - 1 diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index f31b43bc0..26c668cbe 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -415,15 +415,12 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: BIT Numerical Mathematics, 46(1), 195–202. https://doi.org/10.1007/s10543-006-0045-4. """ - if L < 1: - raise ValueError(f"Bandlimit must be at least 1: L = {L}") - n_theta = samples.ntheta(L, "cc") - # Computing RFFT below with n=0 will fail so explicitly return - # weights for single point case - if n_theta == 1: - return xp.array([2.0]) + # Computing RFFT below with n = n_theta - 1 < 1 will fail so raise + # early in this case + if n_theta < 2: + raise ValueError("Clenshaw-Curtis sampling scheme requires n_theta >= 2") n = n_theta - 1 w = _fejer_second_rule_rfft_weights(n, xp) @@ -466,9 +463,6 @@ def quad_weights_f2_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: BIT Numerical Mathematics, 46(1), 195–202. https://doi.org/10.1007/s10543-006-0045-4. """ - if L < 1: - raise ValueError(f"Bandlimit must be at least 1: L = {L}") - n_theta = samples.ntheta(L, "f2") w = _fejer_second_rule_rfft_weights(n_theta + 1, xp) weights = xp.fft.irfft(w, n=n_theta + 1) diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index d08812cd9..c8d69c01a 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -51,9 +51,8 @@ def test_quadrature_exceptions(): with pytest.raises(ValueError): quadrature.quad_weights(L, sampling="foo") - for sampling in ("cc", "f2"): - with pytest.raises(ValueError, match="must be at least 1"): - quadrature.quad_weights(0, sampling) + with pytest.raises(ValueError, match="at least L=2"): + quadrature.quad_weights(1, "cc") def _bandlimit_from_n_theta(n_theta, rule): From 57a94fb0f063d5a75bbbb099fcc4be695c7cf21d Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 11:06:02 +0100 Subject: [PATCH 53/66] Test ntheta and thetas directly for cc / L=1 case rather than transforms --- tests/test_samples.py | 8 ++++++++ tests/test_spherical_transform.py | 9 --------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/test_samples.py b/tests/test_samples.py index 68d0a1c34..43254fd71 100644 --- a/tests/test_samples.py +++ b/tests/test_samples.py @@ -176,3 +176,11 @@ def test_samples_exceptions(): with pytest.raises(ValueError): samples.t2theta(t=0, sampling="mw") + + +@pytest.mark.parametrize("L", [0, 1]) +def test_cc_with_invalid_L_raises_error(L): + with pytest.raises(ValueError, match="at least L=2"): + samples.ntheta(L, sampling="cc") + with pytest.raises(ValueError, match="at least L=2"): + samples.thetas(L, sampling="cc") diff --git a/tests/test_spherical_transform.py b/tests/test_spherical_transform.py index 414ff9378..73fd3c9d6 100644 --- a/tests/test_spherical_transform.py +++ b/tests/test_spherical_transform.py @@ -233,12 +233,3 @@ def test_sampling_exceptions(flm_generator): with pytest.raises(ValueError): spherical.forward(None, 0, 0, None, method="incorrect") - - -@pytest.mark.parametrize("reality", [True, False]) -@pytest.mark.parametrize("method", ["jax", "numpy"]) -def test_cc_with_L_1_raises_error(flm_generator, reality, method): - L = 1 - flm = flm_generator(L=L, reality=reality) - with pytest.raises(ValueError, match="at least L=2"): - spherical.inverse(flm, L=L, sampling="cc", reality=reality, method=method) From 86465a432817cc46c12e07cd0f3d9264f40cbb84 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:27:41 +0100 Subject: [PATCH 54/66] Use functions instead of sets to define sampling scheme properties --- s2fft/precompute_transforms/construct.py | 16 +-- s2fft/recursions/price_mcewen.py | 4 +- s2fft/sampling/s2_samples.py | 127 ++++++++++++++++++----- s2fft/transforms/otf_recursions.py | 14 +-- s2fft/transforms/spherical.py | 8 +- tests/test_quadrature.py | 2 +- tests/test_spherical_precompute.py | 2 +- 7 files changed, 122 insertions(+), 51 deletions(-) diff --git a/s2fft/precompute_transforms/construct.py b/s2fft/precompute_transforms/construct.py index b441a1313..950e5f1de 100644 --- a/s2fft/precompute_transforms/construct.py +++ b/s2fft/precompute_transforms/construct.py @@ -140,7 +140,7 @@ def spin_spherical_kernel( # using the Fourier decomposition of Wigner d-functions. # - The complexity of this approach is O(L^3LogL). # - This approach is stable for arbitrary abs(spins) <= L. - if sampling.lower() in samples.EQUIANGULAR_SCHEMES: + if samples.is_equiangular_in_theta(sampling): nsamps = _n_sample_wigner_fourier_inverse_fft(sampling.lower(), len(thetas)) delta = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64) @@ -263,12 +263,12 @@ def spin_spherical_kernel_jax( dl = jnp.swapaxes(dl, 0, 1) # North pole singularity - if sampling.lower() in samples.INCLUDES_NORTH_POLE_SCHEMES: + if samples.includes_north_pole(sampling): dl = dl.at[0].set(0) dl = dl.at[0, :, L - 1 - spin].set(1) # South pole singularity - if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: + if samples.includes_south_pole(sampling): dl = dl.at[-1].set(0) dl = dl.at[-1, :, L - 1 + spin].set((-1) ** (jnp.arange(L) - spin)) dl = dl.at[:, : jnp.abs(spin)].multiply(0) @@ -297,7 +297,7 @@ def spin_spherical_kernel_jax( # using the Fourier decomposition of Wigner d-functions. # - The complexity of this approach is O(L^3LogL). # - This approach is stable for arbitrary abs(spins) <= L. - elif sampling.lower() in samples.EQUIANGULAR_SCHEMES: + elif samples.is_equiangular_in_theta(sampling): nsamps = _n_sample_wigner_fourier_inverse_fft(sampling.lower(), len(thetas)) delta = jnp.zeros((2 * L - 1, 2 * L - 1), dtype=jnp.float64) @@ -385,7 +385,7 @@ def wigner_kernel( np.ndarray: Transform kernel for Wigner transform. """ - if mode.lower() == "fft" and sampling.lower() not in samples.EQUIANGULAR_SCHEMES: + if mode.lower() == "fft" and not samples.is_equiangular_in_theta(sampling): raise ValueError( f"Fourier based recursion is not valid for {sampling} sampling." ) @@ -393,7 +393,7 @@ def wigner_kernel( # - Can only use the FFT approach when uniformly sampling in theta. # - FFT approach is only more efficient when N <= L/Log(L) roughly. if mode.lower() == "auto": - if sampling.lower() in samples.EQUIANGULAR_SCHEMES: + if samples.is_equiangular_in_theta(sampling): mode = "fft" if N <= int(L / np.log(L)) else "direct" else: mode = "direct" @@ -518,7 +518,7 @@ def wigner_kernel_jax( jnp.ndarray: Transform kernel for Wigner transform. """ - if mode.lower() == "fft" and sampling.lower() not in samples.EQUIANGULAR_SCHEMES: + if mode.lower() == "fft" and not samples.is_equiangular_in_theta(sampling): raise ValueError( f"Fourier based recursion is not valid for {sampling} sampling." ) @@ -526,7 +526,7 @@ def wigner_kernel_jax( # - Can only use the FFT approach when uniformly sampling in theta. # - FFT approach is only more efficient when N <= L/Log(L) roughly. if mode.lower() == "auto": - if sampling.lower() in samples.EQUIANGULAR_SCHEMES: + if samples.is_equiangular_in_theta(sampling): mode = "fft" if N <= int(L / np.log(L)) else "direct" else: mode = "direct" diff --git a/s2fft/recursions/price_mcewen.py b/s2fft/recursions/price_mcewen.py index e5ea6dbd6..4b769cfdf 100644 --- a/s2fft/recursions/price_mcewen.py +++ b/s2fft/recursions/price_mcewen.py @@ -59,7 +59,7 @@ def generate_precomputes( else: beta = samples.thetas(L, sampling, nside) - if forward and sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: + if forward and samples.includes_both_poles(sampling): beta = beta[1:-1] ntheta = len(beta) # Number of theta samples @@ -167,7 +167,7 @@ def generate_precomputes_jax( else: beta = samples.thetas(L, sampling, nside) - if forward and sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: + if forward and samples.includes_both_poles(sampling): beta = beta[1:-1] else: beta = betas diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index a0c146247..9ff4f35b1 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -1,13 +1,5 @@ import numpy as np -# Sampling schemes which include south pole (θ = π) - handling of pole singularity needs to be special cased -INCLUDES_SOUTH_POLE_SCHEMES = frozenset(("mw", "mwss", "cc")) -# Sampling schemes which include north pole (θ = 0) - handling of pole singularity needs to be special cased -INCLUDES_NORTH_POLE_SCHEMES = frozenset(("mwss", "cc")) -INCLUDES_BOTH_POLES_SCHEMES = INCLUDES_SOUTH_POLE_SCHEMES & INCLUDES_NORTH_POLE_SCHEMES -# Sampling schemes for which the spacing between pairs of co-latitudes (θ) is constant -EQUIANGULAR_SCHEMES = frozenset(("mw", "mwss", "dh", "cc", "f2")) - def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: r""" @@ -146,7 +138,7 @@ def m_offset(L: int, sampling: str) -> int: Args: L (int): Harmonic band-limit. - sampling (str, optional): Sampling scheme. Supported sampling schemes include + sampling (str): Sampling scheme. Supported sampling schemes include {"mw", "mwss", "dh", "gl", "healpix", "cc", "f2"}. Returns: @@ -158,6 +150,85 @@ def m_offset(L: int, sampling: str) -> int: ) +def includes_south_pole(sampling: str) -> bool: + """ + Whether the sampling scheme includes the south-pole (`θ = π`), which is assumed + to be the last θ sample. + + The Wigner-d function recursions used in the latitudinal step are singular at the + poles, with coefficients diverging as `sin θ → 0`, so pole samples cannot be + evaluated within the recursion. The transforms handle the poles separately, with + at `θ = π` only the `m = spin` harmonic coefficients contributing, and these + have a closed form. + + Args: + sampling (str): Sampling scheme. + + Returns: + Boolean value. + + """ + return sampling.lower() in ("mw", "mwss", "cc") + + +def includes_north_pole(sampling: str) -> bool: + """ + Whether the sampling scheme includes the north-pole (`θ = 0`), which is assumed + to be the first θ sample. + + The Wigner-d function recursions used in the latitudinal step are singular at the + poles, with coefficients diverging as `sin θ → 0`, so pole samples cannot be + evaluated within the recursion. The transforms handle the poles separately, with + at `θ = 0` only the `m = -spin` harmonic coefficients contributing, and these + have a closed form. + + Args: + sampling (str): Sampling scheme. + + Returns: + Boolean value. + + """ + return sampling.lower() in ("mwss", "cc") + + +def includes_both_poles(sampling: str) -> bool: + """ + Whether the sampling scheme includes both the south-pole (`θ = π`) and north-pole + (`θ = 0`), which are assumed to be respectively the last and first θ samples. + + The Wigner-d function recursions used in the latitudinal step are singular at the + poles, with coefficients diverging as `sin θ → 0`, so pole samples cannot be + evaluated within the recursion. The transforms handle the poles separately, with + at `θ = 0` / `θ = π` only the `m = -spin` / `m = spin` harmonic coefficients + contributing, and these have a closed form. + + Args: + sampling (str): Sampling scheme. + + Returns: + Boolean value. + + """ + return includes_south_pole(sampling) and includes_north_pole(sampling) + + +def is_equiangular_in_theta(sampling: str) -> bool: + """ + Whether the sampling scheme has θ samples which are uniformly spaced. + This allows FFT-based constructon of the Wigner-d kernels via the Fourier + decomposition of Wigner-d functions and a periodic extension in θ. + + Args: + sampling (str): Sampling scheme. + + Returns: + Boolean value. + + """ + return sampling.lower in ("mw", "mwss", "dh", "cc", "f2") + + def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> tuple[int, int]: r""" Shape of intermediate array, before/after latitudinal step. @@ -675,9 +746,9 @@ def hp_getidx(L: int, el: int, m: int) -> int: def flm_2d_to_1d(flm_2d: np.ndarray, L: int) -> np.ndarray: r""" Convert from 2D indexed harmonic coefficients to 1D indexed coefficients. - + Note: - Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for + Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for e.g. :math:`L = 3`, are as follows. .. math:: @@ -688,7 +759,7 @@ def flm_2d_to_1d(flm_2d: np.ndarray, L: int) -> np.ndarray: 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix} - + .. math:: \text{1D data format}: [flm_{0,0}, flm_{1,-1}, flm_{1,0}, flm_{1,1}, \dots] @@ -722,9 +793,9 @@ def flm_2d_to_1d(flm_2d: np.ndarray, L: int) -> np.ndarray: def flm_1d_to_2d(flm_1d: np.ndarray, L: int) -> np.ndarray: r""" Convert from 1D indexed harmnonic coefficients to 2D indexed coefficients. - + Note: - Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for + Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for e.g. :math:`L = 3`, are as follows. .. math:: @@ -735,7 +806,7 @@ def flm_1d_to_2d(flm_1d: np.ndarray, L: int) -> np.ndarray: 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix} - + .. math:: \text{1D data format}: [flm_{0,0}, flm_{1,-1}, flm_{1,0}, flm_{1,1}, \dots] @@ -770,10 +841,10 @@ def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray: r""" Converts from HEALPix (healpy) indexed harmonic coefficients to 2D indexed coefficients. - + Notes: - HEALPix implicitly assumes conjugate symmetry and thus only stores positive `m` - coefficients. Here we unpack that into harmonic coefficients of an + HEALPix implicitly assumes conjugate symmetry and thus only stores positive `m` + coefficients. Here we unpack that into harmonic coefficients of an explicitly real signal. Warning: @@ -781,7 +852,7 @@ def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray: where `L = lmax + 1`. Note: - Storage conventions for harmonic coefficients :math:`f_{(\ell,m)}`, for + Storage conventions for harmonic coefficients :math:`f_{(\ell,m)}`, for e.g. :math:`L = 3`, are as follows. .. math:: @@ -791,7 +862,7 @@ def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray: 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix} - + .. math:: \text{HEALPix}: [flm_{(0,0)}, \dots, flm_{(2,0)}, flm_{(1,1)}, \dots, flm_{(L-1,1)}, \dots] @@ -826,15 +897,15 @@ def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray: r""" Converts from 2D indexed harmonic coefficients to HEALPix (healpy) indexed coefficients. - + Note: - HEALPix implicitly assumes conjugate symmetry and thus only stores positive `m` - coefficients. So this function discards the negative `m` values. This process - is NOT invertible! See the `healpy api docs `_ + HEALPix implicitly assumes conjugate symmetry and thus only stores positive `m` + coefficients. So this function discards the negative `m` values. This process + is NOT invertible! See the `healpy api docs `_ for details on healpy indexing and lengths. Note: - Storage conventions for harmonic coefficients :math:`f_{(\ell,m)}`, for + Storage conventions for harmonic coefficients :math:`f_{(\ell,m)}`, for e.g. :math:`L = 3`, are as follows. .. math:: @@ -844,7 +915,7 @@ def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray: 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix} - + .. math:: \text{HEALPix}: [flm_{(0,0)}, \dots, flm_{(2,0)}, flm_{(1,1)}, \dots, flm_{(L-1,1)}, \dots] @@ -860,10 +931,10 @@ def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray: flm_2d (np.ndarray): 2D indexed harmonic coefficients. L (int): Harmonic band-limit. - + Returns: np.ndarray: HEALPix indexed harmonic coefficients. - + """ flm_hp = np.zeros(int(L * (L + 1) / 2), dtype=np.complex128) diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py index 2012198d5..c09d04da7 100644 --- a/s2fft/transforms/otf_recursions.py +++ b/s2fft/transforms/otf_recursions.py @@ -412,7 +412,7 @@ def eval_recursion_step( # Remove south pole singularity m_offset = samples.m_offset(L, sampling) - if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: + if samples.includes_south_pole(sampling): ftm = ftm.at[-1].set(0) ftm = ftm.at[-1, L - 1 + spin + m_offset].set( jnp.nansum( @@ -421,7 +421,7 @@ def eval_recursion_step( ) # Remove north pole singularity - if sampling.lower() in samples.INCLUDES_NORTH_POLE_SCHEMES: + if samples.includes_north_pole(sampling): ftm = ftm.at[0].set(0) ftm = ftm.at[0, L - 1 - spin + m_offset].set( jnp.nansum(flm[L_lower:, L - 1 - spin]) @@ -654,11 +654,11 @@ def forward_latitudinal_step_jax( """ # Avoid pole-singularities for MWSS sampling - if sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: + if samples.includes_both_poles(sampling): ftm = ftm_in[1:-1] beta = beta_in[1:-1] - elif sampling.lower() in ( - samples.INCLUDES_SOUTH_POLE_SCHEMES - samples.INCLUDES_NORTH_POLE_SCHEMES + elif samples.includes_south_pole(sampling) and not samples.includes_north_pole( + sampling ): ftm = ftm_in[:-1] beta = beta_in[:-1] @@ -878,12 +878,12 @@ def eval_recursion_step( # Include both pole singularities explicitly m_offset = samples.m_offset(L, sampling) - if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: + if samples.includes_south_pole(sampling): flm = flm.at[L_lower:, L - 1 + spin].add( (-1) ** abs(jnp.arange(L_lower, L) - spin) * ftm_in[-1, L - 1 + spin + m_offset] ) - if sampling.lower() in samples.INCLUDES_NORTH_POLE_SCHEMES: + if samples.includes_north_pole(sampling): flm = flm.at[L_lower:, L - 1 - spin].add(ftm_in[0, L - 1 - spin + m_offset]) return flm diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py index e51fe297e..4774d0117 100644 --- a/s2fft/transforms/spherical.py +++ b/s2fft/transforms/spherical.py @@ -174,13 +174,13 @@ def inverse_numpy( ) # Remove south pole singularity - if sampling.lower() in samples.INCLUDES_SOUTH_POLE_SCHEMES: + if samples.includes_south_pole(sampling): ftm[-1] = 0 ftm[-1, L - 1 + spin + m_offset] = np.nansum( (-1) ** abs(np.arange(L0, L) - spin) * flm[L0:, L - 1 + spin] ) # Remove north pole singularity - if sampling.lower() in samples.INCLUDES_NORTH_POLE_SCHEMES: + if samples.includes_north_pole(sampling): ftm[0] = 0 ftm[0, L - 1 - spin + m_offset] = jnp.nansum(flm[L0:, L - 1 - spin]) @@ -527,7 +527,7 @@ def forward_numpy( ftm[:, m_start_ind + m_offset :] *= phase_shifts # Perform latitudinal wigner-d recursions - if sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: + if samples.includes_both_poles(sampling): flm = otf.forward_latitudinal_step( ftm[1:-1], thetas[1:-1], @@ -545,7 +545,7 @@ def forward_numpy( ) # Include both pole singularities explicitly - if sampling.lower() in samples.INCLUDES_BOTH_POLES_SCHEMES: + if samples.includes_both_poles(sampling): flm[L0:, L - 1 + spin] += (-1) ** abs(np.arange(L0, L) - spin) * ftm[ -1, L - 1 + spin + m_offset ] diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py index c8d69c01a..cd7978022 100644 --- a/tests/test_quadrature.py +++ b/tests/test_quadrature.py @@ -84,7 +84,7 @@ def _min_n_points_exact(rule, polynomial_degree): def _min_n_points(rule): # Rules including both poles need a minimum of 2 points - if rule in samples.INCLUDES_BOTH_POLES_SCHEMES: + if samples.includes_both_poles(rule): return 2 else: return 1 diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py index bb21ae6cb..1a15ac607 100644 --- a/tests/test_spherical_precompute.py +++ b/tests/test_spherical_precompute.py @@ -363,7 +363,7 @@ def test_wigner_kernel_fft_mode_non_equiangular_scheme_raises(sampling): wigner_kernel_function(L, N, sampling=sampling, mode=mode) -@pytest.mark.parametrize("sampling", samples.EQUIANGULAR_SCHEMES) +@pytest.mark.parametrize("sampling", ("mw", "mwss", "dh", "cc", "f2")) @pytest.mark.parametrize("L", [1, 2, 5]) def test_n_sample_wigner_fourier_inverse_fft(sampling, L): if sampling == "cc" and L == 1: From 2b10ab1dfe38ee8476910e445b9ad97a2e550b45 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:31:45 +0100 Subject: [PATCH 55/66] Add clarifying note about L > 1 requirement for CC sampling to docs --- docs/sampling.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/sampling.rst b/docs/sampling.rst index cae3b8b67..233f9f46a 100644 --- a/docs/sampling.rst +++ b/docs/sampling.rst @@ -247,7 +247,8 @@ Sample positions are defined by \varphi_p &= \frac{2\pi p}{2L}, &\quad p\in\lbrace 0,1,...,2L-1\rbrace. Note that here we follow the original definition of the Clenshaw-Curtis quadrature rule -and include samples at both poles. +and include samples at both poles, and as there are :math:`2L-1` :math:`\theta` samples, +we require that :math:`L \geq 2`. An analogous scheme which excludes the pole is available as :ref:`F2 `. Due to the redundancy in samples at the poles, the total number of distinct sites on the sphere used by this sampling scheme is :math:`(2L-3)(2L)+2 = 4L^2 -6L + 2`. From 33cbff1e9209e053a1dfe539c825dcbc86135536 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:32:41 +0100 Subject: [PATCH 56/66] Update stale comment about pole singularity handling --- s2fft/transforms/otf_recursions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py index c09d04da7..544ee916e 100644 --- a/s2fft/transforms/otf_recursions.py +++ b/s2fft/transforms/otf_recursions.py @@ -653,7 +653,7 @@ def forward_latitudinal_step_jax( recover acceleration by the number of devices. """ - # Avoid pole-singularities for MWSS sampling + # Avoid pole-singularities where relevant if samples.includes_both_poles(sampling): ftm = ftm_in[1:-1] beta = beta_in[1:-1] From 23f3bb2611c7c7f5ccae0b001775d0584b3a6197 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:38:49 +0100 Subject: [PATCH 57/66] Fix incorrect testing of function set membership rather than function call value --- s2fft/sampling/s2_samples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 9ff4f35b1..9387bc366 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -226,7 +226,7 @@ def is_equiangular_in_theta(sampling: str) -> bool: Boolean value. """ - return sampling.lower in ("mw", "mwss", "dh", "cc", "f2") + return sampling.lower() in ("mw", "mwss", "dh", "cc", "f2") def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> tuple[int, int]: From a904b580e6b15811a9796a237a3d9102a5010242 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:39:42 +0100 Subject: [PATCH 58/66] Add assertion to check match between test parameters and equiangular property --- tests/test_spherical_precompute.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py index 1a15ac607..6fcca19ce 100644 --- a/tests/test_spherical_precompute.py +++ b/tests/test_spherical_precompute.py @@ -366,6 +366,7 @@ def test_wigner_kernel_fft_mode_non_equiangular_scheme_raises(sampling): @pytest.mark.parametrize("sampling", ("mw", "mwss", "dh", "cc", "f2")) @pytest.mark.parametrize("L", [1, 2, 5]) def test_n_sample_wigner_fourier_inverse_fft(sampling, L): + assert samples.is_equiangular_in_theta(sampling) if sampling == "cc" and L == 1: pytest.skip("sampling='cc' requires L > 1") n_sample = c._n_sample_wigner_fourier_inverse_fft( From 337c6e6d8c967e2de2473f3f1be8b452859d0f57 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:44:38 +0100 Subject: [PATCH 59/66] Use assert rather than exception for check --- s2fft/utils/quadrature.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py index 26c668cbe..71c960a64 100644 --- a/s2fft/utils/quadrature.py +++ b/s2fft/utils/quadrature.py @@ -417,10 +417,9 @@ def quad_weights_cc_theta_only(L: int, xp: ModuleType = np) -> np.ndarray: """ n_theta = samples.ntheta(L, "cc") - # Computing RFFT below with n = n_theta - 1 < 1 will fail so raise - # early in this case - if n_theta < 2: - raise ValueError("Clenshaw-Curtis sampling scheme requires n_theta >= 2") + # Computing RFFT below with n = n_theta - 1 < 1 will fail - n_theta should be >= 3 by + # construction of definition in samples.ntheta so the below should always pass + assert n_theta > 1 n = n_theta - 1 w = _fejer_second_rule_rfft_weights(n, xp) From 07844b787cc3b7e44f3a5b13225f407d9001a065 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:46:37 +0100 Subject: [PATCH 60/66] Docstring typo fixes --- s2fft/sampling/s2_samples.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 9387bc366..f39f577ac 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -194,7 +194,7 @@ def includes_north_pole(sampling: str) -> bool: def includes_both_poles(sampling: str) -> bool: """ - Whether the sampling scheme includes both the south-pole (`θ = π`) and north-pole + Whether the sampling scheme includes both the south-pole (`θ = π`) and north-pole (`θ = 0`), which are assumed to be respectively the last and first θ samples. The Wigner-d function recursions used in the latitudinal step are singular at the @@ -216,7 +216,7 @@ def includes_both_poles(sampling: str) -> bool: def is_equiangular_in_theta(sampling: str) -> bool: """ Whether the sampling scheme has θ samples which are uniformly spaced. - This allows FFT-based constructon of the Wigner-d kernels via the Fourier + This allows FFT-based construction of the Wigner-d kernels via the Fourier decomposition of Wigner-d functions and a periodic extension in θ. Args: From ba1a96f6b0613366ab295586119b52957e953d29 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:55:15 +0100 Subject: [PATCH 61/66] Update ntheta raises docstring section and add check for L >= 1 --- s2fft/sampling/s2_samples.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index f39f577ac..33b923775 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -22,6 +22,10 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: ValueError: Sampling scheme not supported. + ValueError: Using any non-HEALPix sampling scheme with L < 1. + + ValueError: Using Clenshaw-Curtis sampling scheme with L < 2. + Returns: int: Number of :math:`\theta` samples of sampling scheme at given resolution. @@ -30,6 +34,10 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int: raise ValueError( f"Sampling scheme sampling={sampling} with L={L} not supported" ) + elif sampling.lower() != "healpix" and L < 1: + raise ValueError( + f"Sampling scheme sampling={sampling} requires bandlimit L >= 1" + ) if sampling.lower() in ("mw", "gl"): return L From f54ea7be7ad03b569cb313cd40e5bc5b7d62e907 Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 17:56:46 +0100 Subject: [PATCH 62/66] Add test for L < 1 case --- tests/test_samples.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_samples.py b/tests/test_samples.py index 43254fd71..b5e58c524 100644 --- a/tests/test_samples.py +++ b/tests/test_samples.py @@ -177,6 +177,9 @@ def test_samples_exceptions(): with pytest.raises(ValueError): samples.t2theta(t=0, sampling="mw") + with pytest.raises(ValueError, match="L >= 1"): + samples.ntheta(0, sampling="mw") + @pytest.mark.parametrize("L", [0, 1]) def test_cc_with_invalid_L_raises_error(L): From 5af2d9c9896181fe893fc9c4f0680028e900b22f Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Tue, 20 Oct 2026 18:02:11 +0100 Subject: [PATCH 63/66] Revert unrelated trailing whitespace cleanup --- s2fft/sampling/s2_samples.py | 38 ++++++++++++++++++------------------ 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py index 33b923775..ead85d016 100644 --- a/s2fft/sampling/s2_samples.py +++ b/s2fft/sampling/s2_samples.py @@ -754,9 +754,9 @@ def hp_getidx(L: int, el: int, m: int) -> int: def flm_2d_to_1d(flm_2d: np.ndarray, L: int) -> np.ndarray: r""" Convert from 2D indexed harmonic coefficients to 1D indexed coefficients. - + Note: - Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for + Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for e.g. :math:`L = 3`, are as follows. .. math:: @@ -767,7 +767,7 @@ def flm_2d_to_1d(flm_2d: np.ndarray, L: int) -> np.ndarray: 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix} - + .. math:: \text{1D data format}: [flm_{0,0}, flm_{1,-1}, flm_{1,0}, flm_{1,1}, \dots] @@ -801,9 +801,9 @@ def flm_2d_to_1d(flm_2d: np.ndarray, L: int) -> np.ndarray: def flm_1d_to_2d(flm_1d: np.ndarray, L: int) -> np.ndarray: r""" Convert from 1D indexed harmnonic coefficients to 2D indexed coefficients. - + Note: - Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for + Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for e.g. :math:`L = 3`, are as follows. .. math:: @@ -814,7 +814,7 @@ def flm_1d_to_2d(flm_1d: np.ndarray, L: int) -> np.ndarray: 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix} - + .. math:: \text{1D data format}: [flm_{0,0}, flm_{1,-1}, flm_{1,0}, flm_{1,1}, \dots] @@ -849,10 +849,10 @@ def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray: r""" Converts from HEALPix (healpy) indexed harmonic coefficients to 2D indexed coefficients. - + Notes: - HEALPix implicitly assumes conjugate symmetry and thus only stores positive `m` - coefficients. Here we unpack that into harmonic coefficients of an + HEALPix implicitly assumes conjugate symmetry and thus only stores positive `m` + coefficients. Here we unpack that into harmonic coefficients of an explicitly real signal. Warning: @@ -860,7 +860,7 @@ def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray: where `L = lmax + 1`. Note: - Storage conventions for harmonic coefficients :math:`f_{(\ell,m)}`, for + Storage conventions for harmonic coefficients :math:`f_{(\ell,m)}`, for e.g. :math:`L = 3`, are as follows. .. math:: @@ -870,7 +870,7 @@ def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray: 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix} - + .. math:: \text{HEALPix}: [flm_{(0,0)}, \dots, flm_{(2,0)}, flm_{(1,1)}, \dots, flm_{(L-1,1)}, \dots] @@ -905,15 +905,15 @@ def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray: r""" Converts from 2D indexed harmonic coefficients to HEALPix (healpy) indexed coefficients. - + Note: - HEALPix implicitly assumes conjugate symmetry and thus only stores positive `m` - coefficients. So this function discards the negative `m` values. This process - is NOT invertible! See the `healpy api docs `_ + HEALPix implicitly assumes conjugate symmetry and thus only stores positive `m` + coefficients. So this function discards the negative `m` values. This process + is NOT invertible! See the `healpy api docs `_ for details on healpy indexing and lengths. Note: - Storage conventions for harmonic coefficients :math:`f_{(\ell,m)}`, for + Storage conventions for harmonic coefficients :math:`f_{(\ell,m)}`, for e.g. :math:`L = 3`, are as follows. .. math:: @@ -923,7 +923,7 @@ def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray: 0 & flm_{(1,-1)} & flm_{(1,0)} & flm_{(1,1)} & 0 \\ flm_{(2,-2)} & flm_{(2,-1)} & flm_{(2,0)} & flm_{(2,1)} & flm_{(2,2)} \end{bmatrix} - + .. math:: \text{HEALPix}: [flm_{(0,0)}, \dots, flm_{(2,0)}, flm_{(1,1)}, \dots, flm_{(L-1,1)}, \dots] @@ -939,10 +939,10 @@ def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray: flm_2d (np.ndarray): 2D indexed harmonic coefficients. L (int): Harmonic band-limit. - + Returns: np.ndarray: HEALPix indexed harmonic coefficients. - + """ flm_hp = np.zeros(int(L * (L + 1) / 2), dtype=np.complex128) From 8aa6208347726b3396e6af7a3478ac0501e8fdec Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 21 Oct 2026 10:12:10 +0100 Subject: [PATCH 64/66] Correct regex match in exception test --- tests/test_samples.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_samples.py b/tests/test_samples.py index b5e58c524..2cbac81bd 100644 --- a/tests/test_samples.py +++ b/tests/test_samples.py @@ -183,7 +183,7 @@ def test_samples_exceptions(): @pytest.mark.parametrize("L", [0, 1]) def test_cc_with_invalid_L_raises_error(L): - with pytest.raises(ValueError, match="at least L=2"): + with pytest.raises(ValueError, match="bandlimit L >= 1"): samples.ntheta(L, sampling="cc") - with pytest.raises(ValueError, match="at least L=2"): + with pytest.raises(ValueError, match="bandlimit L >= 1"): samples.thetas(L, sampling="cc") From c419e66362682275d643de8ec9d1a95bf64447af Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 21 Oct 2026 14:16:54 +0100 Subject: [PATCH 65/66] Cover both invalid L check errors messages in exception test regex match --- tests/test_samples.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_samples.py b/tests/test_samples.py index 2cbac81bd..97c9bc9fa 100644 --- a/tests/test_samples.py +++ b/tests/test_samples.py @@ -183,7 +183,7 @@ def test_samples_exceptions(): @pytest.mark.parametrize("L", [0, 1]) def test_cc_with_invalid_L_raises_error(L): - with pytest.raises(ValueError, match="bandlimit L >= 1"): + with pytest.raises(ValueError, match="(L >= 1)|(at least L=2)"): samples.ntheta(L, sampling="cc") - with pytest.raises(ValueError, match="bandlimit L >= 1"): + with pytest.raises(ValueError, match="(L >= 1)|(at least L=2)"): samples.thetas(L, sampling="cc") From 098e5cc64623ee89da60035dbab55b4af66ff8de Mon Sep 17 00:00:00 2001 From: Matt Graham Date: Wed, 21 Oct 2026 17:19:29 +0100 Subject: [PATCH 66/66] Update and tighten array-api-compat dependency version specifier --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 6b96a5d8a..2dc3472fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "numpy>=1.20", "jax>=0.5.0", "array-api-compat~=1.15", - "array-api-extra >= 0.10.1", + "array-api-extra~=0.11", ] dynamic = [ "version",