Update python_requires and test matrix to support Python 3.11+ (#305)

* Update python_requires and test matrix

* Ruff autofixes for type hints with 3.11+ features

* Use miniforge to install pytorch / healpy on MacOS

* Try using conda-pypi to install dependencies on MacOS

* Manually specify dependencies to install with conda

* Fix pytorch conda package name and skip PyPI dependency install on MacOS

* Add tmate step to allow debugging

* Remove tmate and use explicit shell

* Set explicit shell options as default for job + relax NumPy requirement

* Readd upper bound on NumPy version

* Exclude Python 3.13 on MacOS from matrix

Author: matt-graham

.github/workflows/tests.yml

@@ -25,13 +25,18 @@ jobs:
   build:
 
     runs-on: ${{ matrix.os }}
+    defaults:
+      run:
+        shell: bash -el {0}
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
-        os: [ubuntu-latest]
-        include:
+        python-version: ["3.11", "3.12", "3.13"]
+        os: [ubuntu-latest, macos-latest]
+        exclude:
+          # Skip Python 3.13 on MacOS as 1.20<=numpy<2 requirement inherited from so3
+          # requiring numpy<2 cannot be resolved there
           - os: macos-latest
-            python-version: "3.8"
+            python-version: "3.13"
       fail-fast: false
     env:
       CMAKE_POLICY_VERSION_MINIMUM: 3.5
@@ -42,14 +47,34 @@ jobs:
       with:
         fetch-depth: 0
         fetch-tags: true
-
-    - name: Set up Python ${{ matrix.python-version }}
+        
+    - if: matrix.os == 'macos-latest'
+      name: Set up Miniforge on MacOS
+      uses: conda-incubator/setup-miniconda@v3
+      with:
+        miniforge-version: latest
+        python-version: ${{ matrix.python-version }}
+        
+    - if: matrix.os == 'macos-latest'
+      name: Install dependencies with conda on MacOS
+      # Avoid OpenMP runtime incompatibility when using PyPI wheels
+      # by installing torch and healpy using conda
+      # https://github.com/healpy/healpy/issues/1012
+      run: |
+        conda install jax "jax>=0.3.13,<0.6.0" "numpy>=1.20,<2" ducc0 healpy pytorch pytest pytest-cov
+        python -m pip install --upgrade pip
+        pip install --no-deps so3 pyssht
+        pip install --no-deps .
+  
+    - if: matrix.os != 'macos-latest'
+      name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
         cache: pip
 
-    - name: Install dependencies
+    - if: matrix.os != 'macos-latest'
+      name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install .[tests]

benchmarks/benchmarking.py

@@ -252,7 +252,10 @@ def _format_results_entry(results_entry: dict) -> str:
 
 def _dict_product(dicts: dict[str, Iterable[Any]]) -> Iterable[dict[str, Any]]:
     """Generator corresponding to Cartesian product of dictionaries."""
-    return (dict(zip(dicts.keys(), values)) for values in product(*dicts.values()))
+    return (
+        dict(zip(dicts.keys(), values, strict=False))
+        for values in product(*dicts.values())
+    )
 
 
 def _parse_value(value: str) -> Any:

benchmarks/plotting.py

@@ -141,10 +141,10 @@ def plot_results_against_bandlimit(
         squeeze=False,
     )
     axes = axes.T if functions_along_columns else axes
-    for axes_row, function in zip(axes, functions):
+    for axes_row, function in zip(axes, functions, strict=False):
         results = benchmark_results["results"][function]
         l_values = np.array([r["parameters"]["L"] for r in results])
-        for ax, measurement in zip(axes_row, measurements):
+        for ax, measurement in zip(axes_row, measurements, strict=False):
             plot_function, label = _measurement_plot_functions_and_labels[measurement]
             try:
                 plot_function(ax, "L", l_values, results)

pyproject.toml

@@ -2,8 +2,8 @@
 requires = [
     "setuptools",
     "setuptools-scm",
-    "scikit-build-core >=0.4.3", 
-    "nanobind >=1.3.2"
+    "scikit-build-core>=0.4.3", 
+    "nanobind>=1.3.2"
 ]
 build-backend = "scikit_build_core.build"
 
@@ -16,11 +16,9 @@ authors = [
 classifiers = [
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.8",
-    "Programming Language :: Python :: 3.9",
-    "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Operating System :: OS Independent",
     "Intended Audience :: Developers",
     "Intended Audience :: Science/Research",
@@ -38,7 +36,7 @@ keywords = [
 ]
 name = "s2fft"
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.11"
 license.file = "LICENCE.txt"
 urls.homepage = "https://github.com/astro-informatics/s2fft"
 

s2fft/precompute_transforms/construct.py

@@ -1,4 +1,3 @@
-from typing import Tuple
 from warnings import warn
 
 import jax
@@ -612,7 +611,7 @@ def wigner_kernel_jax(
 wigner_kernel_torch = torch_wrapper.wrap_as_torch_function(wigner_kernel_jax)
 
 
-def fourier_wigner_kernel(L: int) -> Tuple[np.ndarray, np.ndarray]:
+def fourier_wigner_kernel(L: int) -> tuple[np.ndarray, np.ndarray]:
     """
     Computes Fourier coefficients of the reduced Wigner d-functions and quadrature
     weights upsampled for the forward Fourier-Wigner transform.
@@ -640,7 +639,7 @@ def fourier_wigner_kernel(L: int) -> Tuple[np.ndarray, np.ndarray]:
     return deltas, w
 
 
-def fourier_wigner_kernel_jax(L: int) -> Tuple[jnp.ndarray, jnp.ndarray]:
+def fourier_wigner_kernel_jax(L: int) -> tuple[jnp.ndarray, jnp.ndarray]:
     """
     Computes Fourier coefficients of the reduced Wigner d-functions and quadrature
     weights upsampled for the forward Fourier-Wigner transform (JAX implementation).

s2fft/precompute_transforms/custom_ops.py

@@ -1,5 +1,4 @@
 from functools import partial
-from typing import Tuple
 
 import jax.numpy as jnp
 import numpy as np
@@ -9,7 +8,7 @@
 def wigner_subset_to_s2(
     flmn: np.ndarray,
     spins: np.ndarray,
-    DW: Tuple[np.ndarray, np.ndarray],
+    DW: tuple[np.ndarray, np.ndarray],
     L: int,
     sampling: str = "mw",
 ) -> np.ndarray:
@@ -91,7 +90,7 @@ def wigner_subset_to_s2(
 def wigner_subset_to_s2_jax(
     flmn: jnp.ndarray,
     spins: jnp.ndarray,
-    DW: Tuple[jnp.ndarray, jnp.ndarray],
+    DW: tuple[jnp.ndarray, jnp.ndarray],
     L: int,
     sampling: str = "mw",
 ) -> jnp.ndarray:
@@ -173,7 +172,7 @@ def wigner_subset_to_s2_jax(
 def so3_to_wigner_subset(
     f: np.ndarray,
     spins: np.ndarray,
-    DW: Tuple[np.ndarray, np.ndarray],
+    DW: tuple[np.ndarray, np.ndarray],
     L: int,
     N: int,
     sampling: str = "mw",
@@ -214,7 +213,7 @@ def so3_to_wigner_subset(
 def so3_to_wigner_subset_jax(
     f: jnp.ndarray,
     spins: jnp.ndarray,
-    DW: Tuple[jnp.ndarray, jnp.ndarray],
+    DW: tuple[jnp.ndarray, jnp.ndarray],
     L: int,
     N: int,
     sampling: str = "mw",
@@ -257,7 +256,7 @@ def so3_to_wigner_subset_jax(
 def s2_to_wigner_subset(
     fs: np.ndarray,
     spins: np.ndarray,
-    DW: Tuple[np.ndarray, np.ndarray],
+    DW: tuple[np.ndarray, np.ndarray],
     L: int,
     sampling: str = "mw",
 ) -> np.ndarray:
@@ -343,7 +342,7 @@ def s2_to_wigner_subset(
 def s2_to_wigner_subset_jax(
     fs: jnp.ndarray,
     spins: jnp.ndarray,
-    DW: Tuple[jnp.ndarray, jnp.ndarray],
+    DW: tuple[jnp.ndarray, jnp.ndarray],
     L: int,
     sampling: str = "mw",
 ) -> jnp.ndarray:

s2fft/precompute_transforms/spherical.py

@@ -1,5 +1,4 @@
 from functools import partial
-from typing import Optional
 from warnings import warn
 
 import jax.numpy as jnp
@@ -21,11 +20,11 @@ def inverse(
     flm: np.ndarray,
     L: int,
     spin: int = 0,
-    kernel: Optional[np.ndarray] = None,
+    kernel: np.ndarray | None = None,
     sampling: str = "mw",
     reality: bool = False,
     method: str = "jax",
-    nside: Optional[int] = None,
+    nside: int | None = None,
 ) -> np.ndarray:
     r"""
     Compute the inverse spherical harmonic transform via precompute.
@@ -228,11 +227,11 @@ def forward(
     f: np.ndarray,
     L: int,
     spin: int = 0,
-    kernel: Optional[np.ndarray] = None,
+    kernel: np.ndarray | None = None,
     sampling: str = "mw",
     reality: bool = False,
     method: str = "jax",
-    nside: Optional[int] = None,
+    nside: int | None = None,
     iter: int = 0,
 ) -> np.ndarray:
     r"""

s2fft/recursions/price_mcewen.py

@@ -1,6 +1,5 @@
 import warnings
 from functools import partial
-from typing import List
 
 import jax.lax as lax
 import jax.numpy as jnp
@@ -19,7 +18,7 @@ def generate_precomputes(
     nside: int = None,
     forward: bool = False,
     L_lower: int = 0,
-) -> List[np.ndarray]:
+) -> list[np.ndarray]:
     r"""
     Compute recursion coefficients with :math:`\mathcal{O}(L^3)` memory overhead.
 
@@ -125,7 +124,7 @@ def generate_precomputes_jax(
     forward: bool = False,
     L_lower: int = 0,
     betas: jnp.ndarray = None,
-) -> List[jnp.ndarray]:
+) -> list[jnp.ndarray]:
     r"""
     Compute recursion coefficients with :math:`\mathcal{O}(L^2)` memory overhead.
     In practice one could compute these on-the-fly but the memory overhead is
@@ -264,7 +263,7 @@ def generate_precomputes_wigner(
     forward: bool = False,
     reality: bool = False,
     L_lower: int = 0,
-) -> List[List[np.ndarray]]:
+) -> list[list[np.ndarray]]:
     r"""
     Compute recursion coefficients with :math:`\mathcal{O}(L^2)` memory overhead.
     In practice one could compute these on-the-fly but the memory overhead is
@@ -316,7 +315,7 @@ def generate_precomputes_wigner_jax(
     forward: bool = False,
     reality: bool = False,
     L_lower: int = 0,
-) -> List[List[jnp.ndarray]]:
+) -> list[list[jnp.ndarray]]:
     r"""
     Compute recursion coefficients with :math:`\mathcal{O}(L^2)` memory overhead.
     In practice one could compute these on-the-fly but the memory overhead is

s2fft/sampling/s2_samples.py

@@ -1,5 +1,3 @@
-from typing import Tuple
-
 import numpy as np
 
 
@@ -125,7 +123,7 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int:
     return 1
 
 
-def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> Tuple[int, int]:
+def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> tuple[int, int]:
     r"""
     Shape of intermediate array, before/after latitudinal step.
 
@@ -445,7 +443,7 @@ def ring_phase_shift_hp(
     return np.exp(sign * 1j * np.arange(m_start_ind, L) * phi_offset)
 
 
-def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> Tuple[int]:
+def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> tuple[int]:
     r"""
     Shape of spherical signal.
 
@@ -480,7 +478,7 @@ def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> Tuple[int
         return ntheta(L, sampling), nphi_equiang(L, sampling)
 
 
-def flm_shape(L: int) -> Tuple[int, int]:
+def flm_shape(L: int) -> tuple[int, int]:
     r"""
     Standard shape of harmonic coefficients.
 

s2fft/sampling/so3_samples.py

@@ -1,13 +1,11 @@
-from typing import Tuple
-
 import numpy as np
 
 from s2fft.sampling import s2_samples as samples
 
 
 def f_shape(
     L: int, N: int, sampling: str = "mw", nside: int = None
-) -> Tuple[int, int, int]:
+) -> tuple[int, int, int]:
     r"""
     Computes the pixel-space sampling shape for signal on the rotation group
     :math:`SO(3)`.
@@ -49,7 +47,7 @@ def f_shape(
         raise ValueError(f"Sampling scheme sampling={sampling} not supported")
 
 
-def flmn_shape(L: int, N: int) -> Tuple[int, int, int]:
+def flmn_shape(L: int, N: int) -> tuple[int, int, int]:
     r"""
     Computes the shape of Wigner coefficients for signal on the rotation group
     :math:`SO(3)`.
@@ -69,7 +67,7 @@ def flmn_shape(L: int, N: int) -> Tuple[int, int, int]:
 
 def fnab_shape(
     L: int, N: int, sampling: str = "mw", nside: int = None
-) -> Tuple[int, int, int]:
+) -> tuple[int, int, int]:
     r"""
     Computes the shape of Wigner coefficients for signal on the rotation group
     :math:`SO(3)`.