Ruff autofixes for type hints with 3.11+ features

Author: matt-graham

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)

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)`.

s2fft/transforms/otf_recursions.py

@@ -1,5 +1,4 @@
 from functools import partial
-from typing import List
 
 import jax.lax as lax
 import jax.numpy as jnp
@@ -21,7 +20,7 @@ def inverse_latitudinal_step(
     nside: int,
     sampling: str = "mw",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     L_lower: int = 0,
 ) -> np.ndarray:
     r"""
@@ -181,7 +180,7 @@ def inverse_latitudinal_step_jax(
     nside: int,
     sampling: str = "mw",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     spmd: bool = False,
     L_lower: int = 0,
 ) -> jnp.ndarray:
@@ -438,7 +437,7 @@ def forward_latitudinal_step(
     nside: int,
     sampling: str = "mw",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     L_lower: int = 0,
 ) -> np.ndarray:
     r"""
@@ -598,7 +597,7 @@ def forward_latitudinal_step_jax(
     nside: int,
     sampling: str = "mw",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     spmd: bool = False,
     L_lower: int = 0,
 ) -> jnp.ndarray:

s2fft/transforms/spherical.py

@@ -1,5 +1,4 @@
 from functools import partial
-from typing import List, Optional
 
 import jax.numpy as jnp
 import numpy as np
@@ -27,7 +26,7 @@ def inverse(
     sampling: str = "mw",
     method: str = "numpy",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     spmd: bool = False,
     L_lower: int = 0,
     _ssht_backend: int = 1,
@@ -117,7 +116,7 @@ def inverse_numpy(
     nside: int = None,
     sampling: str = "mw",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     L_lower: int = 0,
 ) -> np.ndarray:
     r"""
@@ -217,7 +216,7 @@ def inverse_jax(
     nside: int = None,
     sampling: str = "mw",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     spmd: bool = False,
     L_lower: int = 0,
     use_healpix_custom_primitive: bool = False,
@@ -354,14 +353,14 @@ def forward(
     f: np.ndarray,
     L: int,
     spin: int = 0,
-    nside: Optional[int] = None,
+    nside: int | None = None,
     sampling: str = "mw",
     method: str = "numpy",
     reality: bool = False,
-    precomps: Optional[List] = None,
+    precomps: list | None = None,
     spmd: bool = False,
     L_lower: int = 0,
-    iter: Optional[int] = None,
+    iter: int | None = None,
     _ssht_backend: int = 1,
 ) -> np.ndarray:
     r"""
@@ -472,7 +471,7 @@ def forward_numpy(
     nside: int = None,
     sampling: str = "mw",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     L_lower: int = 0,
 ) -> np.ndarray:
     r"""
@@ -597,7 +596,7 @@ def forward_jax(
     nside: int = None,
     sampling: str = "mw",
     reality: bool = False,
-    precomps: List = None,
+    precomps: list = None,
     spmd: bool = False,
     L_lower: int = 0,
     use_healpix_custom_primitive: bool = False,