Repository: astro-informatics/s2fft (language: python)

Issue:
Decide on a code style convention and set-up automatic formatting
We don't currently specify the code style convention we want the package code to adhere to. Ideally we should settle on something and document this to make it easy for contributors to know what to target when proposing changes and to ensure we are consistent across the repository to improve readability. This has come up as part of @ASKabalan's contributions in #204.

Currently I think we are mainly following [PEP8 guidelines](https://peps.python.org/pep-0008/) and this is probably the most commonly used convention in open-source Python projects so I would personally suggest adopting this. Alternatives include the [Google Python style guide](https://google.github.io/styleguide/pyguide.html).

Related to this, to simplify keeping to a standard style, it would be good to adopt a specific formatter which enforces our chosen style conventions - as an additional check as part of the CI workflows and/or by setting up [`pre-commit` hooks](https://pre-commit.com/) to automatically run the formatter on any commits. My suggestion for a formatter would be [the Ruff formatter](https://docs.astral.sh/ruff/formatter/) as its very performant so minimizes overhead when formatting and uses the same formatting rules as [Black ](https://github.com/psf/black) which is widely used. 

Tagging @jasonmcewen and @CosmoMatt for their views.


Discussion / hints:
I am ok with any formatter.
[the Ruff formatter](https://docs.astral.sh/ruff/formatter/) seems ok for me

I also add the fact that I chose the c++ formatter to be [google](https://github.com/astro-informatics/s2fft/blob/cufft-primitive/.clang-format) with clang-format, using [LSST guidelines formatting rules](https://developer.lsst.io/cpp/clang-format.html#general)
@jasonmcewen, @CosmoMatt and myself discussed this in a meeting today and there was a consesus for using the [Black code style](https://black.readthedocs.io/en/stable/the_black_code_style/index.html). 
Thanks @matt-graham.  Yes, since we're using Black in other codes and its fairly common for Python we thought we should go with that.  We have been using Black already with s2fft I belive, although perhaps not extensively.  We should indeed try to get this fully set up and specified in any contributor guidelines.

Fix patch:
diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml
new file mode 100644
index 00000000..6197fa5a
--- /dev/null
+++ b/.github/workflows/linting.yml
@@ -0,0 +1,33 @@
+name: Linting
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  linting:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout source
+        uses: actions/checkout@v4
+
+      - name: Cache pre-commit
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/pre-commit
+          key: pre-commit-${{ '{{' }} hashFiles('.pre-commit-config.yaml') {{ '}}' }}
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.x"
+          cache: pip
+          cache-dependency-path: pyproject.toml
+
+      - name: Install dependencies
+        run: python -m pip install pre-commit
+
+      - name: Run pre-commit
+        run: pre-commit run --all-files --color always --verbose
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 00000000..6ca21cc4
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,6 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.6.9
+    hooks:
+      - id: ruff
+      - id: ruff-format
diff --git a/README.md b/README.md
index ac31c3da..93d43b9c 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,10 @@
-[![image](https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml/badge.svg?branch=main)](https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml)
-[![image](https://codecov.io/gh/astro-informatics/s2fft/branch/main/graph/badge.svg?token=7QYAFAAWLE)](https://codecov.io/gh/astro-informatics/s2fft)
-[![image](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
-[![image](https://badge.fury.io/py/s2fft.svg)](https://badge.fury.io/py/s2fft)
-[![image](http://img.shields.io/badge/arXiv-2311.14670-orange.svg?style=flat)](https://arxiv.org/abs/2311.14670)<!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->
+[![Tests status](https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml/badge.svg?branch=main)](https://github.com/astro-informatics/s2fft/actions/workflows/tests.yml)
+[![Linting status](https://github.com/astro-informatics/s2fft/actions/workflows/linting.yml/badge.svg?branch=main)](https://github.com/astro-informatics/s2fft/actions/workflows/linting.yml)
+[![Documentation status](https://github.com/astro-informatics/s2fft/actions/workflows/docs.yml/badge.svg?branch=main)](https://github.com/astro-informatics/s2fft/actions/workflows/docs.yml)
+[![Codecov](https://codecov.io/gh/astro-informatics/s2fft/branch/main/graph/badge.svg?token=7QYAFAAWLE)](https://codecov.io/gh/astro-informatics/s2fft)
+[![MIT License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![PyPI package](https://badge.fury.io/py/s2fft.svg)](https://badge.fury.io/py/s2fft)
+[![arXiv](http://img.shields.io/badge/arXiv-2311.14670-orange.svg?style=flat)](https://arxiv.org/abs/2311.14670)<!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->
 [![All Contributors](https://img.shields.io/badge/all_contributors-11-orange.svg?style=flat-square)](#contributors-)
 <!-- ALL-CONTRIBUTORS-BADGE:END --> 
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/astro-informatics/s2fft/blob/main/notebooks/spherical_harmonic_transform.ipynb)
diff --git a/benchmarks/benchmarking.py b/benchmarks/benchmarking.py
index 946a87e9..5a9496fb 100644
--- a/benchmarks/benchmarking.py
+++ b/benchmarks/benchmarking.py
@@ -60,13 +60,13 @@ def mean(x, n):
 """
 
 import argparse
+import inspect
+import json
+import timeit
 from ast import literal_eval
 from functools import partial
 from itertools import product
 from pathlib import Path
-import json
-import timeit
-import inspect
 
 try:
     import memory_profiler
diff --git a/benchmarks/spherical.py b/benchmarks/spherical.py
index de803266..c0f529e0 100644
--- a/benchmarks/spherical.py
+++ b/benchmarks/spherical.py
@@ -2,11 +2,11 @@
 
 import numpy as np
 import pyssht
+from benchmarking import benchmark, parse_args_collect_and_run_benchmarks, skip
+
 import s2fft
 from s2fft.recursions.price_mcewen import generate_precomputes
 from s2fft.sampling import s2_samples as samples
-from benchmarking import benchmark, skip, parse_args_collect_and_run_benchmarks
-
 
 L_VALUES = [8, 16, 32, 64, 128, 256]
 L_LOWER_VALUES = [0]
diff --git a/benchmarks/wigner.py b/benchmarks/wigner.py
index 51b32d90..717501e3 100644
--- a/benchmarks/wigner.py
+++ b/benchmarks/wigner.py
@@ -1,13 +1,14 @@
 """Benchmarks for Wigner transforms."""
 
 import numpy as np
+from benchmarking import benchmark, parse_args_collect_and_run_benchmarks, skip
+
 import s2fft
 from s2fft.base_transforms import wigner as base_wigner
 from s2fft.recursions.price_mcewen import (
     generate_precomputes_wigner,
     generate_precomputes_wigner_jax,
 )
-from benchmarking import benchmark, skip, parse_args_collect_and_run_benchmarks
 
 L_VALUES = [16, 32, 64, 128, 256]
 N_VALUES = [2]
diff --git a/docs/conf.py b/docs/conf.py
index 14c8b4f0..f7b09146 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 #
 # Configuration file for the Sphinx documentation builder.
 #
diff --git a/notebooks/plotting_functions.py b/notebooks/plotting_functions.py
index 90937ad2..eb0c1abe 100644
--- a/notebooks/plotting_functions.py
+++ b/notebooks/plotting_functions.py
@@ -1,5 +1,6 @@
-import pyvista as pv
 import numpy as np
+import pyvista as pv
+
 from s2fft.sampling import s2_samples as samples
 
 
@@ -19,7 +20,8 @@ def plot_sphere(
     cmap: str = "inferno",
     isnotebook: bool = False,
 ) -> None:
-    """Plots signal (f) on the sphere using PyVista.
+    """
+    Plots signal (f) on the sphere using PyVista.
 
     Args:
         f (np.ndarray): Signal on the sphere.
@@ -36,8 +38,8 @@ def plot_sphere(
 
     Note:
         TODO: Add support for healpix sampling.
-    """
 
+    """
     phis = samples.phis_equiang(L, sampling)
     thetas = samples.thetas(L, sampling)
     xx_bounds = _cell_bounds(np.degrees(phis))
diff --git a/pyproject.toml b/pyproject.toml
index 16b2642a..75f6777b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -99,6 +99,55 @@ filterwarnings = [
     "ignore:FutureWarning",
 ]
 
+[tool.ruff]
+fix = true
+force-exclude = true
+format.exclude = ["*.ipynb"]
+lint.exclude = ["*.ipynb"]
+lint.ignore = [
+    "B023", # function-uses-loop-variable
+    "COM812", # trailing commas (ruff-format recommended)
+    "D100", # undocumented-public-module
+    "D104", # undocumented-public-package
+    "D203", # no-blank-line-before-class
+    "D205", # blank line between summary and description in docstrings
+    "D212", # multi-line-summary-first-line
+    "D401", # non-imperative-mood
+    "D407", # removed dashes lines under sections
+    "D417", # argument description in docstring (unreliable)
+    "ISC001", # simplify implicit str concatenation (ruff-format recommended)
+    "S101", # assert
+]
+lint.per-file-ignores = {"benchmarks*" = [
+    "D", # No docstring checks in benchmarks
+], "tests*" = [
+    "D", # No docstring checks in tests
+    "INP001", # File is part of an implicit namespace package.
+    "S101", # Use of `assert` detected
+], "__init__.py" = [
+    "F401" # unused-import
+]}
+lint.select = [
+    "B",
+    "C90",
+    "D",
+    "E1",
+    "E4",
+    "E7",
+    "E9",
+    "F",
+    "I",
+    "S",
+    "UP",
+]
+lint.isort.known-first-party = [
+    "s2fft"
+]
+lint.mccabe.max-complexity = 18
+lint.pep8-naming.classmethod-decorators = [
+    "classmethod",
+]
+
 [tool.setuptools]
 packages = ["s2fft"]
 
diff --git a/s2fft/__init__.py b/s2fft/__init__.py
index 8dca5a49..c7606c7d 100644
--- a/s2fft/__init__.py
+++ b/s2fft/__init__.py
@@ -1,16 +1,24 @@
+import logging
+
+import jax
+
 from . import logs
-from .transforms import wigner
-from .transforms.spherical import *
 from .recursions.price_mcewen import (
     generate_precomputes,
     generate_precomputes_jax,
     generate_precomputes_wigner,
     generate_precomputes_wigner_jax,
 )
-from .utils.rotation import rotate_flms, generate_rotate_dls
-
-import logging
-import jax
+from .transforms import wigner
+from .transforms.spherical import (
+    forward,
+    forward_jax,
+    forward_numpy,
+    inverse,
+    inverse_jax,
+    inverse_numpy,
+)
+from .utils.rotation import generate_rotate_dls, rotate_flms
 
 if jax.config.read("jax_enable_x64") is False:
     logger = logging.getLogger("s2fft")
diff --git a/s2fft/base_transforms/__init__.py b/s2fft/base_transforms/__init__.py
index 2d67a9ef..09315991 100644
--- a/s2fft/base_transforms/__init__.py
+++ b/s2fft/base_transforms/__init__.py
@@ -1,2 +1,1 @@
-from . import spherical
-from . import wigner
+from . import spherical, wigner
diff --git a/s2fft/base_transforms/spherical.py b/s2fft/base_transforms/spherical.py
index 6bfaf5de..e4070e3c 100644
--- a/s2fft/base_transforms/spherical.py
+++ b/s2fft/base_transforms/spherical.py
@@ -1,9 +1,11 @@
-import numpy as np
 from warnings import warn
+
+import numpy as np
+
+from s2fft import recursions
 from s2fft.sampling import s2_samples as samples
-from s2fft.utils import quadrature, resampling
 from s2fft.utils import healpix_ffts as hp
-from s2fft import recursions
+from s2fft.utils import quadrature, resampling
 
 
 def inverse(
@@ -15,7 +17,8 @@ def inverse(
     reality: bool = False,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Compute inverse spherical harmonic transform.
+    r"""
+    Compute inverse spherical harmonic transform.
 
     Uses a vectorised separation of variables method with np.fft.
 
@@ -41,6 +44,7 @@ def inverse(
 
     Returns:
         np.ndarray: Signal on the sphere.
+
     """
     return _inverse(
         flm,
@@ -64,7 +68,8 @@ def _inverse(
     reality: bool = False,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Compute inverse spherical harmonic transform using a specified method.
+    r"""
+    Compute inverse spherical harmonic transform using a specified method.
 
     Args:
         flm (np.ndarray): Spherical harmonic coefficients.
@@ -92,6 +97,7 @@ def _inverse(
 
     Returns:
         np.ndarray: Signal on the sphere.
+
     """
     assert flm.shape == samples.flm_shape(L)
     assert L > 0
@@ -101,7 +107,8 @@ def _inverse(
         reality = False
         warn(
             "Reality acceleration only supports spin 0 fields. "
-            + "Deferring to complex transform."
+            + "Deferring to complex transform.",
+            stacklevel=2,
         )
 
     thetas = samples.thetas(L, sampling, nside)
@@ -132,7 +139,8 @@ def forward(
     reality: bool = False,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Compute forward spherical harmonic transform.
+    r"""
+    Compute forward spherical harmonic transform.
 
     Uses a vectorised separation of variables method with np.fft.
 
@@ -158,6 +166,7 @@ def forward(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients.
+
     """
     return _forward(
         f,
@@ -181,7 +190,8 @@ def _forward(
     reality: bool = False,
     L_lower: int = 0,
 ):
-    r"""Compute forward spherical harmonic transform using a specified method.
+    r"""
+    Compute forward spherical harmonic transform using a specified method.
 
     Args:
         f (np.ndarray): Signal on the sphere.
@@ -209,6 +219,7 @@ def _forward(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients.
+
     """
     assert f.shape == samples.f_shape(L, sampling, nside)
     assert L > 0
@@ -218,7 +229,8 @@ def _forward(
         reality = False
         warn(
             "Reality acceleration only supports spin 0 fields. "
-            + "Deferring to complex transform."
+            + "Deferring to complex transform.",
+            stacklevel=2,
         )
 
     if sampling.lower() == "mw":
@@ -265,7 +277,8 @@ def _compute_inverse_direct(
     reality: bool,
     L_lower: int,
 ):
-    r"""Compute inverse spherical harmonic transform directly.
+    r"""
+    Compute inverse spherical harmonic transform directly.
 
     Args:
         flm (np.ndarray): Spherical harmonic coefficients.
@@ -290,6 +303,7 @@ def _compute_inverse_direct(
 
     Returns:
         np.ndarray: Signal on the sphere.
+
     """
     if sampling.lower() != "healpix":
         phis_ring = samples.phis_equiang(L, sampling)
@@ -349,7 +363,8 @@ def _compute_inverse_sov(
     reality: bool,
     L_lower: int,
 ):
-    r"""Compute inverse spherical harmonic transform by separation of variables with a
+    r"""
+    Compute inverse spherical harmonic transform by separation of variables with a
         manual Fourier transform.
 
     Args:
@@ -375,6 +390,7 @@ def _compute_inverse_sov(
 
     Returns:
         np.ndarray: Signal on the sphere.
+
     """
     ftm = np.zeros((len(thetas), 2 * L - 1), dtype=np.complex128)
     for t, theta in enumerate(thetas):
@@ -421,7 +437,8 @@ def _compute_inverse_sov_fft(
     reality: bool,
     L_lower: int,
 ):
-    r"""Compute inverse spherical harmonic transform by separation of variables with a
+    r"""
+    Compute inverse spherical harmonic transform by separation of variables with a
         Fast Fourier transform.
 
     Args:
@@ -447,8 +464,8 @@ def _compute_inverse_sov_fft(
 
     Returns:
         np.ndarray: Signal on the sphere.
-    """
 
+    """
     if sampling.lower() == "healpix":
         assert L >= 2 * nside
 
@@ -513,7 +530,8 @@ def _compute_inverse_sov_fft_vectorized(
     reality: bool,
     L_lower: int,
 ):
-    r"""A vectorized function to compute inverse spherical harmonic transform by
+    r"""
+    A vectorized function to compute inverse spherical harmonic transform by
         separation of variables with a manual Fourier transform.
 
     Args:
@@ -539,6 +557,7 @@ def _compute_inverse_sov_fft_vectorized(
 
     Returns:
         np.ndarray: Signal on the sphere.
+
     """
     ftm = np.zeros(samples.ftm_shape(L, sampling, nside), dtype=np.complex128)
     m_offset = 1 if sampling in ["mwss", "healpix"] else 0
@@ -588,7 +607,8 @@ def _compute_forward_direct(
     reality: bool,
     L_lower: int,
 ):
-    r"""Compute forward spherical harmonic transform directly.
+    r"""
+    Compute forward spherical harmonic transform directly.
 
     Args:
         f (np.ndarray): Signal on the sphere.
@@ -615,6 +635,7 @@ def _compute_forward_direct(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients.
+
     """
     flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
 
@@ -677,7 +698,8 @@ def _compute_forward_sov(
     reality: bool,
     L_lower: int,
 ):
-    r"""Compute forward spherical harmonic transform by separation of variables with a
+    r"""
+    Compute forward spherical harmonic transform by separation of variables with a
         manual Fourier transform.
 
     Args:
@@ -705,8 +727,8 @@ def _compute_forward_sov(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients.
-    """
 
+    """
     if sampling.lower() != "healpix":
         phis_ring = samples.phis_equiang(L, sampling)
 
@@ -772,7 +794,8 @@ def _compute_forward_sov_fft(
     reality: bool,
     L_lower: int,
 ):
-    r"""Compute forward spherical harmonic transform by separation of variables with a
+    r"""
+    Compute forward spherical harmonic transform by separation of variables with a
         Fast Fourier transform.
 
     Args:
@@ -800,6 +823,7 @@ def _compute_forward_sov_fft(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients.
+
     """
     flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
     ftm = np.zeros_like(f).astype(np.complex128)
@@ -886,7 +910,8 @@ def _compute_forward_sov_fft_vectorized(
     reality: bool,
     L_lower: int,
 ):
-    r"""A vectorized function to compute forward spherical harmonic transform by
+    r"""
+    A vectorized function to compute forward spherical harmonic transform by
         separation of variables with a manual Fourier transform.
 
     Args:
@@ -914,6 +939,7 @@ def _compute_forward_sov_fft_vectorized(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients.
+
     """
     flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
     ftm = np.zeros_like(f).astype(np.complex128)
diff --git a/s2fft/base_transforms/wigner.py b/s2fft/base_transforms/wigner.py
index b7ed60fe..46bee380 100644
--- a/s2fft/base_transforms/wigner.py
+++ b/s2fft/base_transforms/wigner.py
@@ -1,6 +1,7 @@
 import numpy as np
-from s2fft.sampling import so3_samples as samples
+
 from s2fft.base_transforms import spherical
+from s2fft.sampling import so3_samples as samples
 
 
 def inverse(
@@ -12,7 +13,8 @@ def inverse(
     reality: bool = False,
     nside: int = None,
 ) -> np.ndarray:
-    r"""Compute the inverse Wigner transform, i.e. inverse Fourier transform on
+    r"""
+    Compute the inverse Wigner transform, i.e. inverse Fourier transform on
     :math:`SO(3)`.
 
     Note:
@@ -49,6 +51,7 @@ def inverse(
         np.ndarray:  Signal on the on :math:`SO(3)` with shape :math:`[n_{\gamma},
         n_{\beta}, n_{\alpha}]`, where :math:`n_\xi` denotes the number of samples for
         angle :math:`\xi`.
+
     """
     assert flmn.shape == samples.flmn_shape(L, N)
 
@@ -90,7 +93,8 @@ def forward(
     reality: bool = False,
     nside: int = None,
 ) -> np.ndarray:
-    r"""Compute the forward Wigner transform, i.e. Fourier transform on
+    r"""
+    Compute the forward Wigner transform, i.e. Fourier transform on
     :math:`SO(3)`.
 
     Note:
@@ -127,6 +131,7 @@ def forward(
 
     Returns:
         np.ndarray: Wigner coefficients `flmn` with shape :math:`[2N-1, L, 2L-1]`.
+
     """
     assert f.shape == samples.f_shape(L, N, sampling, nside)
 
diff --git a/s2fft/logs.py b/s2fft/logs.py
index 82b72b8a..a4124dfb 100644
--- a/s2fft/logs.py
+++ b/s2fft/logs.py
@@ -1,13 +1,16 @@
-import os
-import logging.config
 import logging
+import logging.config
+import os
 from pathlib import Path
+
 import yaml
+
 import s2fft
 
 
 def setup_logging(custom_yaml_path: str = None):
-    """Initialise and configure logging.
+    """
+    Initialise and configure logging.
 
     Should be called at the beginning of code to initialise and configure the
     desired logging level. Logging levels can be ints in [0,50] where 10 is
@@ -29,7 +32,7 @@ def setup_logging(custom_yaml_path: str = None):
         path = Path(custom_yaml_path)
     if not path.exists():
         raise ValueError(f"Logging config path {path} does not exist.")
-    with open(path, "rt") as f:
+    with open(path) as f:
         config = yaml.safe_load(f.read())
     if custom_yaml_path is None:
         config["handlers"]["info_file_handler"]["filename"] = "info.log"
@@ -39,7 +42,8 @@ def setup_logging(custom_yaml_path: str = None):
 
 
 def debug_log(message: str):
-    """Log a debug message (e.g. for background logs to assist debugging).
+    """
+    Log a debug message (e.g. for background logs to assist debugging).
 
     Args:
         message (str): Message to log.
@@ -50,7 +54,8 @@ def debug_log(message: str):
 
 
 def warning_log(message: str):
-    """Log a warning (e.g. for internal code warnings such as large dynamic
+    """
+    Log a warning (e.g. for internal code warnings such as large dynamic
     ranges).
 
     Args:
@@ -62,7 +67,8 @@ def warning_log(message: str):
 
 
 def critical_log(message: str):
-    """Log a critical message (e.g. core code failures etc).
+    """
+    Log a critical message (e.g. core code failures etc).
 
     Args:
         message (str): Message to log.
@@ -73,7 +79,8 @@ def critical_log(message: str):
 
 
 def info_log(message: str):
-    """Log an information message (e.g. evidence value printing, run completion
+    """
+    Log an information message (e.g. evidence value printing, run completion
     etc).
 
     Args:
diff --git a/s2fft/precompute_transforms/__init__.py b/s2fft/precompute_transforms/__init__.py
index 876f8b81..28e31847 100644
--- a/s2fft/precompute_transforms/__init__.py
+++ b/s2fft/precompute_transforms/__init__.py
@@ -1,3 +1,1 @@
-from . import spherical
-from . import wigner
-from . import construct
+from . import construct, spherical, wigner
diff --git a/s2fft/precompute_transforms/construct.py b/s2fft/precompute_transforms/construct.py
index 1ca2f2db..f55496c7 100644
--- a/s2fft/precompute_transforms/construct.py
+++ b/s2fft/precompute_transforms/construct.py
@@ -1,14 +1,12 @@
-import jax
-
-jax.config.update("jax_enable_x64", True)
+from warnings import warn
 
-import numpy as np
 import jax.numpy as jnp
+import numpy as np
 import torch
+
+from s2fft import recursions
 from s2fft.sampling import s2_samples as samples
 from s2fft.utils import quadrature, quadrature_jax
-from s2fft import recursions
-from warnings import warn
 
 
 def spin_spherical_kernel(
@@ -20,7 +18,8 @@ def spin_spherical_kernel(
     forward: bool = False,
     using_torch: bool = False,
 ):
-    r"""Precompute the wigner-d kernel for spin-spherical transform. This can be
+    r"""
+    Precompute the wigner-d kernel for spin-spherical transform. This can be
     drastically faster but comes at a :math:`\mathcal{O}(L^3)` memory overhead, making
     it infeasible for :math:`L\geq 512`.
 
@@ -46,12 +45,14 @@ def spin_spherical_kernel(
 
     Returns:
         np.ndarray: Transform kernel for spin-spherical harmonic transform.
+
     """
     if reality and spin != 0:
         reality = False
         warn(
             "Reality acceleration only supports spin 0 fields. "
-            + "Defering to complex transform."
+            + "Defering to complex transform.",
+            stacklevel=2,
         )
     m_start_ind = L - 1 if reality else 0
     m_dim = L if reality else 2 * L - 1
@@ -92,7 +93,8 @@ def spin_spherical_kernel_jax(
     nside: int = None,
     forward: bool = False,
 ):
-    r"""Precompute the wigner-d kernel for spin-spherical transform. This can be
+    r"""
+    Precompute the wigner-d kernel for spin-spherical transform. This can be
     drastically faster but comes at a :math:`\mathcal{O}(L^3)` memory overhead, making
     it infeasible for :math:`L\geq 512`.
 
@@ -116,6 +118,7 @@ def spin_spherical_kernel_jax(
 
     Returns:
         jnp.ndarray: Transform kernel for spin-spherical harmonic transform.
+
     """
     m_start_ind = L - 1 if reality else 0
 
@@ -171,7 +174,8 @@ def wigner_kernel(
     forward: bool = False,
     using_torch: bool = False,
 ):
-    r"""Precompute the wigner-d kernels required for a Wigner transform. This can be
+    r"""
+    Precompute the wigner-d kernels required for a Wigner transform. This can be
     drastically faster but comes at a :math:`\mathcal{O}(NL^3)` memory overhead, making
     it infeasible for :math:`L \geq 512`.
 
@@ -197,6 +201,7 @@ def wigner_kernel(
 
     Returns:
         np.ndarray: Transform kernel for Wigner transform.
+
     """
     n_start_ind = N - 1 if reality else 0
     n_dim = N if reality else 2 * N - 1
@@ -244,7 +249,8 @@ def wigner_kernel_jax(
     nside: int = None,
     forward: bool = False,
 ):
-    r"""Precompute the wigner-d kernels required for a Wigner transform. This can be
+    r"""
+    Precompute the wigner-d kernels required for a Wigner transform. This can be
     drastically faster but comes at a :math:`\mathcal{O}(NL^3)` memory overhead, making
     it infeasible for :math:`L \geq 512`.
 
@@ -268,6 +274,7 @@ def wigner_kernel_jax(
 
     Returns:
         jnp.ndarray: Transform kernel for Wigner transform.
+
     """
     n_start_ind = N - 1 if reality else 0
     n_dim = N if reality else 2 * N - 1
@@ -330,7 +337,8 @@ def healpix_phase_shifts(
     nside: int,
     forward: bool = False,
 ) -> np.ndarray:
-    r"""Generates a phase shift vector for HEALPix for all :math:`\theta` rings.
+    r"""
+    Generates a phase shift vector for HEALPix for all :math:`\theta` rings.
 
     Args:
         L (int, optional): Harmonic band-limit.
@@ -342,10 +350,11 @@ def healpix_phase_shifts(
 
     Returns:
         np.ndarray: Vector of phase shifts with shape :math:`[thetas, 2L-1]`.
+
     """
     thetas = samples.thetas(L, "healpix", nside)
     phase_array = np.zeros((len(thetas), 2 * L - 1), dtype=np.complex128)
-    for t, theta in enumerate(thetas):
+    for t in range(len(thetas)):
         phase_array[t] = samples.ring_phase_shift_hp(L, t, nside, forward)
 
     return phase_array
diff --git a/s2fft/precompute_transforms/spherical.py b/s2fft/precompute_transforms/spherical.py
index 4286b9d8..c6fe4a0d 100644
--- a/s2fft/precompute_transforms/spherical.py
+++ b/s2fft/precompute_transforms/spherical.py
@@ -1,13 +1,14 @@
-from jax import jit
+from functools import partial
+from warnings import warn
 
-import numpy as np
 import jax.numpy as jnp
+import numpy as np
 import torch
+from jax import jit
+
 from s2fft.sampling import s2_samples as samples
-from s2fft.utils import resampling, resampling_jax, resampling_torch
 from s2fft.utils import healpix_ffts as hp
-from functools import partial
-from warnings import warn
+from s2fft.utils import resampling, resampling_jax, resampling_torch
 
 
 def inverse(
@@ -20,7 +21,8 @@ def inverse(
     method: str = "jax",
     nside: int = None,
 ) -> np.ndarray:
-    r"""Compute the inverse spherical harmonic transform via precompute.
+    r"""
+    Compute the inverse spherical harmonic transform via precompute.
 
     Args:
         flm (np.ndarray): Spherical harmonic coefficients.
@@ -51,12 +53,14 @@ def inverse(
 
     Returns:
         np.ndarray: Pixel-space coefficients with shape.
+
     """
     if reality and spin != 0:
         reality = False
         warn(
             "Reality acceleration only supports spin 0 fields. "
-            + "Defering to complex transform."
+            + "Defering to complex transform.",
+            stacklevel=2,
         )
     if method == "numpy":
         return inverse_transform(flm, kernel, L, sampling, reality, spin, nside)
@@ -77,7 +81,8 @@ def inverse_transform(
     spin: int,
     nside: int,
 ) -> np.ndarray:
-    r"""Compute the forward spherical harmonic transform via precompute (vectorized
+    r"""
+    Compute the forward spherical harmonic transform via precompute (vectorized
     implementation).
 
     Args:
@@ -100,6 +105,7 @@ def inverse_transform(
 
     Returns:
         np.ndarray: Pixel-space coefficients.
+
     """
     m_offset = 1 if sampling in ["mwss", "healpix"] else 0
     m_start_ind = L - 1 if reality else 0
@@ -141,7 +147,8 @@ def inverse_transform_jax(
     spin: int,
     nside: int,
 ) -> jnp.ndarray:
-    r"""Compute the inverse spherical harmonic transform via precompute (JAX
+    r"""
+    Compute the inverse spherical harmonic transform via precompute (JAX
     implementation).
 
     Args:
@@ -164,6 +171,7 @@ def inverse_transform_jax(
 
     Returns:
         jnp.ndarray: Pixel-space coefficients with shape.
+
     """
     m_offset = 1 if sampling in ["mwss", "healpix"] else 0
     m_start_ind = L - 1 if reality else 0
@@ -209,7 +217,8 @@ def inverse_transform_torch(
     spin: int,
     nside: int,
 ) -> torch.tensor:
-    r"""Compute the inverse spherical harmonic transform via precompute (Torch
+    r"""
+    Compute the inverse spherical harmonic transform via precompute (Torch
     implementation).
 
     Args:
@@ -232,6 +241,7 @@ def inverse_transform_torch(
 
     Returns:
         torch.tensor: Pixel-space coefficients with shape.
+
     """
     m_offset = 1 if sampling in ["mwss", "healpix"] else 0
     m_start_ind = L - 1 if reality else 0
@@ -286,7 +296,8 @@ def forward(
     method: str = "jax",
     nside: int = None,
 ) -> np.ndarray:
-    r"""Compute the forward spherical harmonic transform via precompute.
+    r"""
+    Compute the forward spherical harmonic transform via precompute.
 
     Args:
         f (np.ndarray): Signal on the sphere.
@@ -317,12 +328,14 @@ def forward(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients.
+
     """
     if reality and spin != 0:
         reality = False
         warn(
             "Reality acceleration only supports spin 0 fields. "
-            + "Defering to complex transform."
+            + "Defering to complex transform.",
+            stacklevel=2,
         )
     if method == "numpy":
         return forward_transform(f, kernel, L, sampling, reality, spin, nside)
@@ -343,7 +356,8 @@ def forward_transform(
     spin: int,
     nside: int,
 ) -> np.ndarray:
-    r"""Compute the forward spherical harmonic transform via precompute (vectorized
+    r"""
+    Compute the forward spherical harmonic transform via precompute (vectorized
     implementation).
 
     Args:
@@ -366,6 +380,7 @@ def forward_transform(
 
     Returns:
         np.ndarray: Pixel-space coefficients.
+
     """
     if sampling.lower() == "mw":
         f = resampling.mw_to_mwss(f, L, spin)
@@ -411,7 +426,8 @@ def forward_transform_jax(
     spin: int,
     nside: int,
 ) -> jnp.ndarray:
-    r"""Compute the forward spherical harmonic tranclearsform via precompute (vectorized
+    r"""
+    Compute the forward spherical harmonic tranclearsform via precompute (vectorized
     implementation).
 
     Args:
@@ -434,6 +450,7 @@ def forward_transform_jax(
 
     Returns:
         jnp.ndarray: Pixel-space coefficients.
+
     """
     if sampling.lower() == "mw":
         f = resampling_jax.mw_to_mwss(f, L, spin)
@@ -483,7 +500,8 @@ def forward_transform_torch(
     spin: int,
     nside: int,
 ) -> torch.tensor:
-    r"""Compute the forward spherical harmonic tranclearsform via precompute (vectorized
+    r"""
+    Compute the forward spherical harmonic tranclearsform via precompute (vectorized
     implementation).
 
     Args:
@@ -506,6 +524,7 @@ def forward_transform_torch(
 
     Returns:
         torch.tensor: Pixel-space coefficients.
+
     """
     if sampling.lower() == "mw":
         f = resampling_torch.mw_to_mwss(f, L, spin)
diff --git a/s2fft/precompute_transforms/wigner.py b/s2fft/precompute_transforms/wigner.py
index 8e266f51..cd3d4dcd 100644
--- a/s2fft/precompute_transforms/wigner.py
+++ b/s2fft/precompute_transforms/wigner.py
@@ -1,12 +1,13 @@
-from jax import jit
+from functools import partial
 
-import numpy as np
 import jax.numpy as jnp
+import numpy as np
 import torch
-from s2fft.utils import resampling, resampling_jax, resampling_torch
-from s2fft.utils import healpix_ffts as hp
+from jax import jit
+
 from s2fft.sampling import so3_samples as samples
-from functools import partial
+from s2fft.utils import healpix_ffts as hp
+from s2fft.utils import resampling, resampling_jax, resampling_torch
 
 
 def inverse(
@@ -19,7 +20,8 @@ def inverse(
     method: str = "jax",
     nside: int = None,
 ) -> np.ndarray:
-    r"""Compute the inverse Wigner transform, i.e. inverse Fourier transform on
+    r"""
+    Compute the inverse Wigner transform, i.e. inverse Fourier transform on
     :math:`SO(3)`.
 
     Args:
@@ -56,6 +58,7 @@ def inverse(
         For a given :math:`\gamma` we thus recover a signal on the sphere indexed by
         :math:`[\theta, \phi]`, i.e. we associate :math:`\beta` with :math:`\theta` and
         :math:`\alpha` with :math:`\phi`.
+
     """
     if method == "numpy":
         return inverse_transform(flmn, kernel, L, N, sampling, reality, nside)
@@ -76,7 +79,8 @@ def inverse_transform(
     reality: bool,
     nside: int,
 ) -> np.ndarray:
-    r"""Compute the inverse Wigner transform, i.e. inverse Fourier transform on
+    r"""
+    Compute the inverse Wigner transform, i.e. inverse Fourier transform on
     :math:`SO(3)`.
 
     Args:
@@ -99,6 +103,7 @@ def inverse_transform(
 
     Returns:
         np.ndarray: Pixel-space coefficients.
+
     """
     m_offset = 1 if sampling in ["mwss", "healpix"] else 0
     n_start_ind = N - 1 if reality else 0
@@ -137,7 +142,8 @@ def inverse_transform_jax(
     reality: bool,
     nside: int,
 ) -> jnp.ndarray:
-    r"""Compute the inverse Wigner transform, i.e. inverse Fourier transform on
+    r"""
+    Compute the inverse Wigner transform, i.e. inverse Fourier transform on
     :math:`SO(3)`.
 
     Args:
@@ -160,6 +166,7 @@ def inverse_transform_jax(
 
     Returns:
         jnp.ndarray: Pixel-space coefficients.
+
     """
     m_offset = 1 if sampling in ["mwss", "healpix"] else 0
     n_start_ind = N - 1 if reality else 0
@@ -214,7 +221,8 @@ def inverse_transform_torch(
     reality: bool,
     nside: int,
 ) -> torch.tensor:
-    r"""Compute the inverse Wigner transform, i.e. inverse Fourier transform on
+    r"""
+    Compute the inverse Wigner transform, i.e. inverse Fourier transform on
     :math:`SO(3)`.
 
     Args:
@@ -237,6 +245,7 @@ def inverse_transform_torch(
 
     Returns:
         torch.tensor: Pixel-space coefficients.
+
     """
     m_offset = 1 if sampling in ["mwss", "healpix"] else 0
     n_start_ind = N - 1 if reality else 0
@@ -291,7 +300,8 @@ def forward(
     method: str = "jax",
     nside: int = None,
 ) -> np.ndarray:
-    r"""Compute the forward Wigner transform, i.e. Fourier transform on
+    r"""
+    Compute the forward Wigner transform, i.e. Fourier transform on
     :math:`SO(3)`.
 
     Args:
@@ -328,6 +338,7 @@ def forward(
         For a given :math:`\gamma` we thus recover a signal on the sphere indexed by
         :math:`[\theta, \phi]`, i.e. we associate :math:`\beta` with :math:`\theta` and
         :math:`\alpha` with :math:`\phi`.
+
     """
     if method == "numpy":
         return forward_transform(f, kernel, L, N, sampling, reality, nside)
@@ -348,7 +359,8 @@ def forward_transform(
     reality: bool,
     nside: int,
 ) -> np.ndarray:
-    r"""Compute the forward Wigner transform, i.e. Fourier transform on
+    r"""
+    Compute the forward Wigner transform, i.e. Fourier transform on
     :math:`SO(3)`.
 
     Args:
@@ -371,6 +383,7 @@ def forward_transform(
 
     Returns:
         np.ndarray: Wigner space coefficients.
+
     """
     n_start_ind = N - 1 if reality else 0
 
@@ -429,7 +442,8 @@ def forward_transform_jax(
     reality: bool,
     nside: int,
 ) -> jnp.ndarray:
-    r"""Compute the forward Wigner transform, i.e. Fourier transform on
+    r"""
+    Compute the forward Wigner transform, i.e. Fourier transform on
     :math:`SO(3)`.
 
     Args:
@@ -452,6 +466,7 @@ def forward_transform_jax(
 
     Returns:
         jnp.ndarray: Wigner space coefficients.
+
     """
     n_start_ind = N - 1 if reality else 0
 
@@ -521,7 +536,8 @@ def forward_transform_torch(
     reality: bool,
     nside: int,
 ) -> torch.tensor:
-    r"""Compute the forward Wigner transform, i.e. Fourier transform on
+    r"""
+    Compute the forward Wigner transform, i.e. Fourier transform on
     :math:`SO(3)`.
 
     Args:
@@ -544,6 +560,7 @@ def forward_transform_torch(
 
     Returns:
         torch.tensor: Wigner space coefficients.
+
     """
     n_start_ind = N - 1 if reality else 0
 
diff --git a/s2fft/recursions/__init__.py b/s2fft/recursions/__init__.py
index 5ed5eddf..b34cff76 100644
--- a/s2fft/recursions/__init__.py
+++ b/s2fft/recursions/__init__.py
@@ -1,6 +1,1 @@
-from . import trapani
-from . import risbo
-from . import risbo_jax
-from . import turok
-from . import turok_jax
-from . import price_mcewen
+from . import price_mcewen, risbo, risbo_jax, trapani, turok, turok_jax
diff --git a/s2fft/recursions/price_mcewen.py b/s2fft/recursions/price_mcewen.py
index 03a83896..8d4f1037 100644
--- a/s2fft/recursions/price_mcewen.py
+++ b/s2fft/recursions/price_mcewen.py
@@ -1,13 +1,13 @@
-import numpy as np
-import jax.numpy as jnp
+import warnings
+from functools import partial
+from typing import List
+
 import jax.lax as lax
+import jax.numpy as jnp
+import numpy as np
 from jax import jit
-from functools import partial
 
 from s2fft.sampling import s2_samples as samples
-from typing import List
-
-import warnings
 
 warnings.filterwarnings("ignore")
 
@@ -20,7 +20,8 @@ def generate_precomputes(
     forward: bool = False,
     L_lower: int = 0,
 ) -> List[np.ndarray]:
-    r"""Compute recursion coefficients with :math:`\mathcal{O}(L^2)` memory overhead.
+    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
     negligible and well worth the acceleration.
 
@@ -46,6 +47,7 @@ def generate_precomputes(
 
     Note:
         TODO: this function should be optimised.
+
     """
     mm = -spin
     L0 = L_lower
@@ -122,7 +124,8 @@ def generate_precomputes_jax(
     L_lower: int = 0,
     betas: jnp.ndarray = None,
 ) -> List[jnp.ndarray]:
-    r"""Compute recursion coefficients with :math:`\mathcal{O}(L^2)` memory overhead.
+    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
     negligible and well worth the acceleration. JAX implementation of
     :func:`~generate_precomputes`.
@@ -148,6 +151,7 @@ def generate_precomputes_jax(
 
     Returns:
         List[jnp.ndarray]: List of precomputed coefficient arrays.
+
     """
     mm = -spin
     L0 = L_lower
@@ -259,7 +263,8 @@ def generate_precomputes_wigner(
     reality: bool = False,
     L_lower: int = 0,
 ) -> List[List[np.ndarray]]:
-    r"""Compute recursion coefficients with :math:`\mathcal{O}(L^2)` memory overhead.
+    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
     negligible and well worth the acceleration. This is a wrapped extension of
     :func:`~generate_precomputes` for the case of multiple spins, i.e. the Wigner
@@ -291,6 +296,7 @@ def generate_precomputes_wigner(
 
     Note:
         TODO: this function should be optimised.
+
     """
     precomps = []
     n_start_ind = 0 if reality else -N + 1
@@ -309,7 +315,8 @@ def generate_precomputes_wigner_jax(
     reality: bool = False,
     L_lower: int = 0,
 ) -> List[List[jnp.ndarray]]:
-    r"""Compute recursion coefficients with :math:`\mathcal{O}(L^2)` memory overhead.
+    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
     negligible and well worth the acceleration. This is a wrapped extension of
     :func:`~generate_precomputes` for the case of multiple spins, i.e. the Wigner
@@ -338,6 +345,7 @@ def generate_precomputes_wigner_jax(
 
     Returns:
         List[List[jnp.ndarray]]: 2N-1 length List of Lists of precomputed coefficient arrays.
+
     """
     lrenorm = []
     vsign = []
@@ -368,7 +376,8 @@ def generate_precomputes_wigner_jax(
 def compute_all_slices(
     beta: np.ndarray, L: int, spin: int, precomps=None
 ) -> np.ndarray:
-    r"""Compute a particular slice :math:`m^{\prime}`, denoted `mm`,
+    r"""
+    Compute a particular slice :math:`m^{\prime}`, denoted `mm`,
     of the complete Wigner-d matrix for all sampled polar angles
     :math:`\beta` and all :math:`\ell` using Price & McEwen recursion.
 
@@ -401,6 +410,7 @@ def compute_all_slices(
 
     Returns:
         np.ndarray: Wigner-d matrix mm slice of dimension :math:`[2L-1, n_{\theta}, n_{\ell}]`.
+
     """
     # Indexing boundaries and constants
     mm = -spin
@@ -497,7 +507,8 @@ def compute_all_slices_jax(
     nside: int = None,
     precomps=None,
 ) -> jnp.ndarray:
-    r"""Compute a particular slice :math:`m^{\prime}`, denoted `mm`,
+    r"""
+    Compute a particular slice :math:`m^{\prime}`, denoted `mm`,
     of the complete Wigner-d matrix for all sampled polar angles
     :math:`\beta` and all :math:`\ell` using Price & McEwen recursion.
 
@@ -530,6 +541,7 @@ def compute_all_slices_jax(
 
     Returns:
         jnp.ndarray: Wigner-d matrix mm slice of dimension :math:`[2L-1, n_{\theta}, n_{\ell}]`.
+
     """
     # Indexing boundaries and constants
     mm = -spin
diff --git a/s2fft/recursions/risbo.py b/s2fft/recursions/risbo.py
index bfb25958..8cc660f9 100644
--- a/s2fft/recursions/risbo.py
+++ b/s2fft/recursions/risbo.py
@@ -2,7 +2,8 @@
 
 
 def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
-    r"""Compute Wigner-d at argument :math:`\beta` for full plane using
+    r"""
+    Compute Wigner-d at argument :math:`\beta` for full plane using
     Risbo recursion.
 
     The Wigner-d plane is computed by recursion over :math:`\ell`.
@@ -20,8 +21,8 @@ def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d for :math:`\ell` and :math:`\beta`, with full plane computed.
-    """
 
+    """
     _arg_checks(dl, beta, L, el)
 
     if el == 0:
@@ -103,7 +104,8 @@ def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
 
 
 def _arg_checks(dl: np.ndarray, beta: float, L: int, el: int):
-    r"""Check arguments of Risbo functions.
+    r"""
+    Check arguments of Risbo functions.
 
     Args:
         dl (np.ndarray): Wigner-d plane of which to check shape.
@@ -111,8 +113,8 @@ def _arg_checks(dl: np.ndarray, beta: float, L: int, el: int):
         L (int): Harmonic band-limit.
 
         el (int): Spherical harmonic degree :math:`\ell`.
-    """
 
+    """
     assert 0 <= el < L  # Should be < not <= once have init routine.
     assert dl.shape[0] == dl.shape[1] == 2 * L - 1
     assert 0 <= beta <= np.pi
diff --git a/s2fft/recursions/risbo_jax.py b/s2fft/recursions/risbo_jax.py
index d254cf4b..4ba41e34 100644
--- a/s2fft/recursions/risbo_jax.py
+++ b/s2fft/recursions/risbo_jax.py
@@ -1,12 +1,14 @@
+from functools import partial
+
 import jax.numpy as jnp
 from jax import jit
-from functools import partial
 
 
 @partial(jit, static_argnums=(1, 2, 3))
 def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
-    r"""Compute Wigner-d at argument :math:`\beta` for full plane using
-    Risbo recursion (JAX implementation)
+    r"""
+    Compute Wigner-d at argument :math:`\beta` for full plane using
+    Risbo recursion (JAX implementation).
 
     The Wigner-d plane is computed by recursion over :math:`\ell` (`el`).
     Thus, for :math:`\ell > 0` the plane must be computed already for
@@ -23,8 +25,8 @@ def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d for `el` and `beta`, with full plane computed.
-    """
 
+    """
     if el == 0:
         dl = dl.at[el + L - 1, el + L - 1].set(1.0)
         return dl
diff --git a/s2fft/recursions/trapani.py b/s2fft/recursions/trapani.py
index 3c9a2af0..d46368c8 100644
--- a/s2fft/recursions/trapani.py
+++ b/s2fft/recursions/trapani.py
@@ -1,11 +1,13 @@
+from functools import partial
+
+import jax.numpy as jnp
 import numpy as np
 from jax import jit, lax
-import jax.numpy as jnp
-from functools import partial
 
 
 def init(dl: np.ndarray, L: int, implementation: str = "vectorized") -> np.ndarray:
-    r"""Initialise Wigner-d at argument :math:`\pi/2` for :math:`\ell=0` for
+    r"""
+    Initialise Wigner-d at argument :math:`\pi/2` for :math:`\ell=0` for
     Trapani & Navaza recursion (multiple implementations).
 
     Args:
@@ -21,8 +23,8 @@ def init(dl: np.ndarray, L: int, implementation: str = "vectorized") -> np.ndarr
 
     Returns:
         np.ndarray: Plane of Wigner-d initialised for :math:`\ell=0`,
-    """
 
+    """
     if implementation.lower() == "loop":
         return init_nonjax(dl, L)
 
@@ -37,7 +39,8 @@ def init(dl: np.ndarray, L: int, implementation: str = "vectorized") -> np.ndarr
 
 
 def init_nonjax(dl: np.ndarray, L: int) -> np.ndarray:
-    r"""Initialise Wigner-d at argument :math:`\pi/2` for :math:`\ell=0` for
+    r"""
+    Initialise Wigner-d at argument :math:`\pi/2` for :math:`\ell=0` for
     Trapani & Navaza recursion (loop-based/vectorized implementation).
 
     See :func:`~init` for further details.
@@ -54,8 +57,8 @@ def init_nonjax(dl: np.ndarray, L: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d initialised for :math:`\ell=0`,
-    """
 
+    """
     el = 0
     dl[el + L - 1, el + L - 1] = 1.0
 
@@ -64,7 +67,8 @@ def init_nonjax(dl: np.ndarray, L: int) -> np.ndarray:
 
 @partial(jit, static_argnums=(1,))
 def init_jax(dl: jnp.ndarray, L: int) -> jnp.ndarray:
-    r"""Initialise Wigner-d at argument :math:`\pi/2` for :math:`\ell=0` for
+    r"""
+    Initialise Wigner-d at argument :math:`\pi/2` for :math:`\ell=0` for
     Trapani & Navaza recursion (JAX implementation).
 
     See :func:`~init` for further details.
@@ -81,8 +85,8 @@ def init_jax(dl: jnp.ndarray, L: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Plane of Wigner-d initialised for :math:`\ell=0`,
-    """
 
+    """
     el = 0
     dl = dl.at[el + L - 1, el + L - 1].set(1.0)
 
@@ -90,7 +94,8 @@ def init_jax(dl: jnp.ndarray, L: int) -> jnp.ndarray:
 
 
 def compute_eighth(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Compute Wigner-d at argument :math:`\pi/2` for eighth of plane using
+    r"""
+    Compute Wigner-d at argument :math:`\pi/2` for eighth of plane using
     Trapani & Navaza recursion.
 
     The Wigner-d plane is computed by recursion over :math:`\ell` (`el`).
@@ -123,7 +128,6 @@ def compute_eighth(dl: np.ndarray, L: int, el: int) -> np.ndarray:
         np.ndarray: Plane of Wigner-d for `el`, with eighth of plane computed.
 
     """
-
     _arg_checks(dl, L, el)
 
     dmm = np.zeros(L)
@@ -171,7 +175,8 @@ def compute_eighth(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 
 def compute_quarter_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Compute Wigner-d at argument :math:`\pi/2` for quarter of plane using
+    r"""
+    Compute Wigner-d at argument :math:`\pi/2` for quarter of plane using
     Trapani & Navaza recursion (vector implementation).
 
     The Wigner-d plane is computed by recursion over :math:`\ell` (`el`).
@@ -199,8 +204,8 @@ def compute_quarter_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d for `el`, with quarter of plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     dmm = np.zeros(L)
@@ -241,7 +246,8 @@ def compute_quarter_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 @partial(jit, static_argnums=(1,))
 def compute_quarter_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
-    r"""Compute Wigner-d at argument :math:`\pi/2` for quarter of plane using
+    r"""
+    Compute Wigner-d at argument :math:`\pi/2` for quarter of plane using
     Trapani & Navaza recursion (JAX implementation).
 
     The Wigner-d plane is computed by recursion over :math:`\ell` (`el`).
@@ -272,8 +278,8 @@ def compute_quarter_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Plane of Wigner-d for `el`, with quarter of plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     dmm = jnp.zeros(L)
@@ -334,7 +340,8 @@ def compute_dl_submatrix_slice(dl_slice_1_dl_slice_2, t1_fact_i_t2_fact_i):
 
 
 def fill_eighth2quarter(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Fill in quarter of Wigner-d plane from eighth.
+    r"""
+    Fill in quarter of Wigner-d plane from eighth.
 
     The Wigner-d plane passed as an argument should be computed for the eighth
     of the plane  :math:`m^\prime <= m < \ell` and :math:`0 <= m^\prime <= \ell`.
@@ -350,8 +357,8 @@ def fill_eighth2quarter(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d for `el`, with quarter of plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     # Diagonal symmetry to fill in quarter.
@@ -365,7 +372,8 @@ def fill_eighth2quarter(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 
 def fill_quarter2half(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Fill in half of Wigner-d plane from quarter.
+    r"""
+    Fill in half of Wigner-d plane from quarter.
 
     The Wigner-d plane passed as an argument should be computed for the quarter
     of the plane :math:`0 <= m, m^\prime <= \ell`.  The
@@ -381,8 +389,8 @@ def fill_quarter2half(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d for `el`, with half of plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     # Symmetry in m to fill in half.
@@ -396,7 +404,8 @@ def fill_quarter2half(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 
 def fill_quarter2half_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Fill in half of Wigner-d plane from quarter (vectorised implementation).
+    r"""
+    Fill in half of Wigner-d plane from quarter (vectorised implementation).
 
     See :func:`~fill_quarter2half` for further details.
 
@@ -412,8 +421,8 @@ def fill_quarter2half_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d for `el`, with half of plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     # Symmetry in m to fill in half.
@@ -430,7 +439,8 @@ def fill_quarter2half_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 @partial(jit, static_argnums=(1,))
 def fill_quarter2half_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
-    r"""Fill in half of Wigner-d plane from quarter (JAX implementation).
+    r"""
+    Fill in half of Wigner-d plane from quarter (JAX implementation).
 
     See :func:`~fill_quarter2half` for further details.
 
@@ -446,8 +456,8 @@ def fill_quarter2half_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Plane of Wigner-d for `el`, with half of plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     # Symmetry in m to fill in half.
@@ -465,7 +475,8 @@ def fill_quarter2half_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
 
 
 def fill_half2full(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Fill in full Wigner-d plane from half.
+    r"""
+    Fill in full Wigner-d plane from half.
 
     The Wigner-d plane passed as an argument should be computed for the half
     of the plane :math:`-\ell <= m <= \ell` and :math:`0 <= m^\prime <= \ell`.
@@ -484,8 +495,8 @@ def fill_half2full(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d for `el`, with full plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     # Symmetry in mm to fill in remaining plane.
@@ -499,7 +510,8 @@ def fill_half2full(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 
 def fill_half2full_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Fill in full Wigner-d plane from half (vectorized implementation).
+    r"""
+    Fill in full Wigner-d plane from half (vectorized implementation).
 
     See :func:`~fill_half2full` for further details.
 
@@ -515,8 +527,8 @@ def fill_half2full_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Plane of Wigner-d for `el`, with full plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     # Symmetry in mm to fill in remaining plane.
@@ -533,7 +545,8 @@ def fill_half2full_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 @partial(jit, static_argnums=(1,))
 def fill_half2full_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
-    r"""Fill in full Wigner-d plane from half (JAX implementation).
+    r"""
+    Fill in full Wigner-d plane from half (JAX implementation).
 
     See :func:`~fill_half2full` for further details.
 
@@ -549,8 +562,8 @@ def fill_half2full_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Plane of Wigner-d for `el`, with full plane computed.
-    """
 
+    """
     _arg_checks(dl, L, el)
 
     # Symmetry in mm to fill in remaining plane.
@@ -570,7 +583,8 @@ def fill_half2full_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
 def compute_full(
     dl: np.ndarray, L: int, el: int, implementation: str = "vectorized"
 ) -> np.ndarray:
-    r"""Compute Wigner-d at argument :math:`\pi/2` for full plane using
+    r"""
+    Compute Wigner-d at argument :math:`\pi/2` for full plane using
     Trapani & Navaza recursion (multiple implementations).
 
     The Wigner-d plane is computed by recursion over :math:`\ell` (`el`).
@@ -605,7 +619,6 @@ def compute_full(
         np.ndarray: Plane of Wigner-d for `el`, with full plane computed.
 
     """
-
     if implementation.lower() == "loop":
         return compute_full_loop(dl, L, el)
 
@@ -620,7 +633,8 @@ def compute_full(
 
 
 def compute_full_loop(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Compute Wigner-d at argument :math:`\pi/2` for full plane using
+    r"""
+    Compute Wigner-d at argument :math:`\pi/2` for full plane using
     Trapani & Navaza recursion (loop-based implementation).
 
     See :func:`~compute_full` for further details.
@@ -642,7 +656,6 @@ def compute_full_loop(dl: np.ndarray, L: int, el: int) -> np.ndarray:
         np.ndarray: Plane of Wigner-d for `el`, with full plane computed.
 
     """
-
     _arg_checks(dl, L, el)
 
     dl = compute_eighth(dl, L, el)
@@ -654,7 +667,8 @@ def compute_full_loop(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 
 def compute_quarter(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Compute Wigner-d at argument :math:`\pi/2` for quarter plane using
+    r"""
+    Compute Wigner-d at argument :math:`\pi/2` for quarter plane using
     Trapani & Navaza recursion.
 
     The Wigner-d plane is computed by recursion over :math:`\ell` (`el`).
@@ -684,7 +698,6 @@ def compute_quarter(dl: np.ndarray, L: int, el: int) -> np.ndarray:
         np.ndarray: Plane of Wigner-d for `el`, with quarter plane computed.
 
     """
-
     _arg_checks(dl, L, el)
 
     dl = compute_eighth(dl, L, el)
@@ -694,7 +707,8 @@ def compute_quarter(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 
 def compute_full_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
-    r"""Compute Wigner-d at argument :math:`\pi/2` for full plane using
+    r"""
+    Compute Wigner-d at argument :math:`\pi/2` for full plane using
     Trapani & Navaza recursion (vectorized implementation).
 
     See :func:`~compute_full` for further details.
@@ -716,7 +730,6 @@ def compute_full_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
         np.ndarray: Plane of Wigner-d for `el`, with full plane computed.
 
     """
-
     _arg_checks(dl, L, el)
 
     dl = compute_quarter_vectorized(dl, L, el)
@@ -728,7 +741,8 @@ def compute_full_vectorized(dl: np.ndarray, L: int, el: int) -> np.ndarray:
 
 @partial(jit, static_argnums=(1,))
 def compute_full_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
-    r"""Compute Wigner-d at argument :math:`\pi/2` for full plane using
+    r"""
+    Compute Wigner-d at argument :math:`\pi/2` for full plane using
     Trapani & Navaza recursion (JAX implementation).
 
     See :func:`~compute_full` for further details.
@@ -760,7 +774,8 @@ def compute_full_jax(dl: jnp.ndarray, L: int, el: int) -> jnp.ndarray:
 
 
 def _arg_checks(dl: np.ndarray, L: int, el: int):
-    """Check arguments of Trapani functions.
+    r"""
+    Check arguments of Trapani functions.
 
     Args:
         dl (np.ndarray): Wigner-d plane of which to check shape.
@@ -768,6 +783,7 @@ def _arg_checks(dl: np.ndarray, L: int, el: int):
         L (int): Harmonic band-limit.
 
         el (int): Spherical harmonic degree :math:`\ell`.
+
     """
 
     # assert 0 < el < L
diff --git a/s2fft/recursions/turok.py b/s2fft/recursions/turok.py
index 7d9ef915..8a2cd034 100644
--- a/s2fft/recursions/turok.py
+++ b/s2fft/recursions/turok.py
@@ -1,9 +1,11 @@
-import numpy as np
 from warnings import warn
 
+import numpy as np
+
 
 def compute_full(beta: float, el: int, L: int) -> np.ndarray:
-    r"""Compute the complete Wigner-d matrix at polar angle :math:`\beta` using
+    r"""
+    Compute the complete Wigner-d matrix at polar angle :math:`\beta` using
     Turok & Bucher recursion.
 
     The Wigner-d plane for a given :math:`\ell` (`el`) and :math:`\beta`
@@ -28,6 +30,7 @@ def compute_full(beta: float, el: int, L: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Wigner-d matrix of dimension [2L-1, 2L-1].
+
     """
     if el >= L:
         raise ValueError(
@@ -41,7 +44,8 @@ def compute_full(beta: float, el: int, L: int) -> np.ndarray:
 def compute_slice(
     beta: float, el: int, L: int, mm: int, positive_m_only: bool = False
 ) -> np.ndarray:
-    r"""Compute a particular slice :math:`m^{\prime}`, denoted `mm`,
+    r"""
+    Compute a particular slice :math:`m^{\prime}`, denoted `mm`,
     of the complete Wigner-d matrix at polar angle :math:`\beta` using Turok & Bucher
     recursion.
 
@@ -82,6 +86,7 @@ def compute_slice(
 
     Returns:
         np.ndarray: Wigner-d matrix mm slice of dimension [2L-1].
+
     """
     if el < mm:
         raise ValueError(f"Wigner-D not valid for l={el} < mm={mm}.")
@@ -95,7 +100,8 @@ def compute_slice(
         positive_m_only = False
         warn(
             "Reality acceleration only supports spin 0 fields. "
-            + "Defering to complex transform."
+            + "Defering to complex transform.",
+            stacklevel=2,
         )
 
     dl = np.zeros(2 * L - 1, dtype=np.float64)
@@ -110,7 +116,8 @@ def compute_quarter_slice(
     mm: int,
     positive_m_only: bool = False,
 ) -> np.ndarray:
-    r"""Compute a single slice at :math:`m^{\prime}` of the Wigner-d matrix evaluated
+    r"""
+    Compute a single slice at :math:`m^{\prime}` of the Wigner-d matrix evaluated
     at :math:`\beta`.
 
     Args:
@@ -129,6 +136,7 @@ def compute_quarter_slice(
 
     Returns:
         np.ndarray: Wigner-d matrix slice of dimension [2L-1] populated only on the mm slice.
+
     """
     # Analytically evaluate singularities
     if np.isclose(beta, 0, atol=1e-8):
@@ -227,8 +235,9 @@ def compute_quarter_slice(
     return dl
 
 
-def compute_quarter(dl: np.ndarray, beta: float, l: int, L: int) -> np.ndarray:
-    """Compute the left quarter triangle of the Wigner-d matrix via Turok & Bucher
+def compute_quarter(dl: np.ndarray, beta: float, ell: int, L: int) -> np.ndarray:  # noqa: C901
+    """
+    Compute the left quarter triangle of the Wigner-d matrix via Turok & Bucher
     recursion.
 
     Args:
@@ -236,32 +245,33 @@ def compute_quarter(dl: np.ndarray, beta: float, l: int, L: int) -> np.ndarray:
 
         beta (float): Polar angle in radians.
 
-        l (int): Harmonic degree of Wigner-d matrix.
+        ell (int): Harmonic degree of Wigner-d matrix.
 
         L (int): Harmonic band-limit.
 
     Returns:
         np.ndarray: Wigner-d matrix of dimension [2L-1, 2L-1] with left quarter
         triangle populated.
+
     """
     # Analytically evaluate singularities
     if np.isclose(beta, 0, atol=1e-8):
-        dl[L - 1 - l : L + l, L - 1 - l : L + l] = np.identity(
-            2 * l + 1, dtype=np.float64
+        dl[L - 1 - ell : L + ell, L - 1 - ell : L + ell] = np.identity(
+            2 * ell + 1, dtype=np.float64
         )
         return dl
 
     if np.isclose(beta, np.pi, atol=1e-8):
-        for m in range(-l, l + 1):
-            dl[L - 1 - m, L - 1 + m] = (-1) ** (l + m)
+        for m in range(-ell, ell + 1):
+            dl[L - 1 - m, L - 1 + m] = (-1) ** (ell + m)
         return dl
 
-    if l == 0:
+    if ell == 0:
         dl[L - 1, L - 1] = 1
         return dl
 
     # Define constants adopted throughout
-    lp1 = 1 - (L - 1 - l)  # Offset for indexing (currently -L < m < L in 2D)
+    lp1 = 1 - (L - 1 - ell)  # Offset for indexing (currently -L < m < L in 2D)
 
     # These constants handle overflow by retrospectively renormalising
     big_const = 1e10
@@ -277,43 +287,43 @@ def compute_quarter(dl: np.ndarray, beta: float, l: int, L: int) -> np.ndarray:
     omc = 1.0 - c
 
     # Vectors with indexing -L < m < L adopted throughout
-    lrenorm = np.zeros(2 * l + 1, dtype=np.float64)
-    cp = np.zeros(2 * l + 1, dtype=np.float64)
-    cpi = np.zeros(2 * l + 1, dtype=np.float64)
-    cp2 = np.zeros(2 * l + 1, dtype=np.float64)
-    log_first_row = np.zeros(2 * l + 1, dtype=np.float64)
-    sign = np.zeros(2 * l + 1, dtype=np.float64)
+    lrenorm = np.zeros(2 * ell + 1, dtype=np.float64)
+    cp = np.zeros(2 * ell + 1, dtype=np.float64)
+    cpi = np.zeros(2 * ell + 1, dtype=np.float64)
+    cp2 = np.zeros(2 * ell + 1, dtype=np.float64)
+    log_first_row = np.zeros(2 * ell + 1, dtype=np.float64)
+    sign = np.zeros(2 * ell + 1, dtype=np.float64)
 
     # Populate vectors for first row
-    log_first_row[0] = 2.0 * l * np.log(np.abs(c2))
+    log_first_row[0] = 2.0 * ell * np.log(np.abs(c2))
     sign[0] = 1.0
 
-    for i in range(2, 2 * l + 2):
-        m = l + 1 - i
-        ratio = np.sqrt((m + l + 1) / (l - m))
+    for i in range(2, 2 * ell + 2):
+        m = ell + 1 - i
+        ratio = np.sqrt((m + ell + 1) / (ell - m))
         log_first_row[i - 1] = log_first_row[i - 2] + np.log(ratio) + np.log(np.abs(t))
         sign[i - 1] = sign[i - 2] * t / np.abs(t)
 
     # Initialising coefficients cp(m)= cplus(l-m).
-    for m in range(1, l + 2):
-        xm = l - m
-        cpi[m - 1] = 2.0 / np.sqrt(l * (l + 1) - xm * (xm + 1))
+    for m in range(1, ell + 2):
+        xm = ell - m
+        cpi[m - 1] = 2.0 / np.sqrt(ell * (ell + 1) - xm * (xm + 1))
         cp[m - 1] = 1.0 / cpi[m - 1]
 
-    for m in range(2, l + 2):
+    for m in range(2, ell + 2):
         cp2[m - 1] = cpi[m - 1] * cp[m - 2]
 
     dl[1 - lp1, 1 - lp1] = 1.0
-    dl[2 * l + 1 - lp1, 1 - lp1] = 1.0
+    dl[2 * ell + 1 - lp1, 1 - lp1] = 1.0
 
     # Use Turok & Bucher recursion to fill from diagonal to horizontal (lower left eight)
-    for index in range(2, l + 2):
+    for index in range(2, ell + 2):
         dl[index - lp1, 1 - lp1] = 1.0
-        lamb = ((l + 1) * omc - index + c) / s
+        lamb = ((ell + 1) * omc - index + c) / s
         dl[index - lp1, 2 - lp1] = lamb * dl[index - lp1, 1 - lp1] * cpi[0]
         if index > 2:
             for m in range(2, index):
-                lamb = ((l + 1) * omc - index + m * c) / s
+                lamb = ((ell + 1) * omc - index + m * c) / s
                 dl[index - lp1, m + 1 - lp1] = (
                     lamb * cpi[m - 1] * dl[index - lp1, m - lp1]
                     - cp2[m - 1] * dl[index - lp1, m - 1 - lp1]
@@ -325,13 +335,13 @@ def compute_quarter(dl: np.ndarray, beta: float, l: int, L: int) -> np.ndarray:
                         dl[index - lp1, im - lp1] = dl[index - lp1, im - lp1] * bigi
 
     # Use Turok & Bucher recursion to fill horizontal to anti-diagonal (upper left eight)
-    for index in range(l + 2, 2 * l + 1):
+    for index in range(ell + 2, 2 * ell + 1):
         dl[index - lp1, 1 - lp1] = 1.0
-        lamb = ((l + 1) * omc - index + c) / s
+        lamb = ((ell + 1) * omc - index + c) / s
         dl[index - lp1, 2 - lp1] = lamb * dl[index - lp1, 1 - lp1] * cpi[0]
-        if index < 2 * l:
-            for m in range(2, 2 * l - index + 2):
-                lamb = ((l + 1) * omc - index + m * c) / s
+        if index < 2 * ell:
+            for m in range(2, 2 * ell - index + 2):
+                lamb = ((ell + 1) * omc - index + m * c) / s
                 dl[index - lp1, m + 1 - lp1] = (
                     lamb * cpi[m - 1] * dl[index - lp1, m - lp1]
                     - cp2[m - 1] * dl[index - lp1, m - 1 - lp1]
@@ -342,60 +352,62 @@ def compute_quarter(dl: np.ndarray, beta: float, l: int, L: int) -> np.ndarray:
                         dl[index - lp1, im - lp1] = dl[index - lp1, im - lp1] * bigi
 
     # Apply renormalisation
-    for i in range(1, l + 2):
+    for i in range(1, ell + 2):
         renorm = sign[i - 1] * np.exp(log_first_row[i - 1] - lrenorm[i - 1])
         for j in range(1, i + 1):
             dl[i - lp1, j - lp1] = dl[i - lp1, j - lp1] * renorm
 
-    for i in range(l + 2, 2 * l + 2):
+    for i in range(ell + 2, 2 * ell + 2):
         renorm = sign[i - 1] * np.exp(log_first_row[i - 1] - lrenorm[i - 1])
-        for j in range(1, 2 * l + 2 - i + 1):
+        for j in range(1, 2 * ell + 2 - i + 1):
             dl[i - lp1, j - lp1] = dl[i - lp1, j - lp1] * renorm
 
     return dl
 
 
-def fill(dl: np.ndarray, l: int, L: int) -> np.ndarray:
-    """Reflects Wigner-d quarter plane to complete full matrix by using symmetry
+def fill(dl: np.ndarray, ell: int, L: int) -> np.ndarray:
+    """
+    Reflects Wigner-d quarter plane to complete full matrix by using symmetry
     properties of the Wigner-d matrices.
 
     Args:
         dl (np.ndarray): Wigner-d matrix to populate by symmetry.
 
-        l (int): Harmonic degree of Wigner-d matrix.
+        ell (int): Harmonic degree of Wigner-d matrix.
 
         L (int): Harmonic band-limit.
 
     Returns:
         np.ndarray: A complete Wigner-d matrix of dimension [2L-1, 2L-1].
+
     """
-    lp1 = 1 - (L - 1 - l)  # Offset for indexing (currently -L < m < L in 2D)
+    lp1 = 1 - (L - 1 - ell)  # Offset for indexing (currently -L < m < L in 2D)
 
     # Reflect across anti-diagonal
-    for i in range(1, l + 1):
-        for j in range(l + 1, 2 * l + 1 - i + 1):
-            dl[2 * l + 2 - i - lp1, 2 * l + 2 - j - lp1] = dl[j - lp1, i - lp1]
+    for i in range(1, ell + 1):
+        for j in range(ell + 1, 2 * ell + 1 - i + 1):
+            dl[2 * ell + 2 - i - lp1, 2 * ell + 2 - j - lp1] = dl[j - lp1, i - lp1]
 
     # Reflect across diagonal
-    for i in range(1, l + 2):
+    for i in range(1, ell + 2):
         sgn = -1
-        for j in range(i + 1, l + 2):
+        for j in range(i + 1, ell + 2):
             dl[i - lp1, j - lp1] = dl[j - lp1, i - lp1] * sgn
             sgn = sgn * (-1)
 
     # Fill right matrix
-    for i in range(l + 2, 2 * l + 2):
+    for i in range(ell + 2, 2 * ell + 2):
         sgn = (-1) ** (i + 1)
 
-        for j in range(1, 2 * l + 2 - i + 1):
+        for j in range(1, 2 * ell + 2 - i + 1):
             dl[j - lp1, i - lp1] = dl[i - lp1, j - lp1] * sgn
             sgn = sgn * (-1)
 
-        for j in range(i, 2 * l + 2):
-            dl[j - lp1, i - lp1] = dl[2 * l + 2 - i - lp1, 2 * l + 2 - j - lp1]
+        for j in range(i, 2 * ell + 2):
+            dl[j - lp1, i - lp1] = dl[2 * ell + 2 - i - lp1, 2 * ell + 2 - j - lp1]
 
-    for i in range(l + 2, 2 * l + 2):
-        for j in range(2 * l + 3 - i, i - 1 + 1):
-            dl[j - lp1, i - lp1] = dl[2 * l + 2 - i - lp1, 2 * l + 2 - j - lp1]
+    for i in range(ell + 2, 2 * ell + 2):
+        for j in range(2 * ell + 3 - i, i - 1 + 1):
+            dl[j - lp1, i - lp1] = dl[2 * ell + 2 - i - lp1, 2 * ell + 2 - j - lp1]
 
     return dl
diff --git a/s2fft/recursions/turok_jax.py b/s2fft/recursions/turok_jax.py
index eec666b0..9101d12a 100644
--- a/s2fft/recursions/turok_jax.py
+++ b/s2fft/recursions/turok_jax.py
@@ -1,13 +1,15 @@
-import numpy as np
+from functools import partial
+
 import jax.lax as lax
-from jax import jit
 import jax.numpy as jnp
-from functools import partial
+import numpy as np
+from jax import jit
 
 
 @partial(jit, static_argnums=(2, 3))
 def compute_slice(beta: float, el: int, L: int, mm: int) -> jnp.ndarray:
-    r"""Compute a particular slice :math:`m^{\prime}`, denoted `mm`,
+    r"""
+    Compute a particular slice :math:`m^{\prime}`, denoted `mm`,
     of the complete Wigner-d matrix at polar angle :math:`\beta` using Turok &
     Bucher recursion.
 
@@ -37,6 +39,7 @@ def compute_slice(beta: float, el: int, L: int, mm: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Wigner-d matrix mm slice of dimension [2L-1].
+
     """
     dl = jnp.zeros(2 * L - 1, dtype=jnp.float64)
     dl = lax.cond(
@@ -63,7 +66,8 @@ def compute_slice(beta: float, el: int, L: int, mm: int) -> jnp.ndarray:
 def _compute_quarter_slice(
     dl: jnp.array, beta: float, el: int, L: int, mm: int
 ) -> jnp.ndarray:
-    r"""Compute a single slice at :math:`m^{\prime}` of the Wigner-d matrix evaluated
+    r"""
+    Compute a single slice at :math:`m^{\prime}` of the Wigner-d matrix evaluated
     at :math:`\beta`.
 
     Args:
@@ -80,6 +84,7 @@ def _compute_quarter_slice(
     Returns:
         jnp.ndarray: Wigner-d matrix slice of dimension [2L-1] populated only on the mm
             slice.
+
     """
     # These constants handle overflow by retrospectively renormalising
     big_const = 1e10
@@ -184,7 +189,8 @@ def renorm_iteration(m, dl_lrenorm):
 
 @partial(jit, static_argnums=(2, 3))
 def _north_pole(dl, el, L, mm) -> jnp.ndarray:
-    r"""Compute Wigner-d matrix for edge case where theta index is located on the
+    r"""
+    Compute Wigner-d matrix for edge case where theta index is located on the
     north pole.
 
     Args:
@@ -198,6 +204,7 @@ def _north_pole(dl, el, L, mm) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Wigner-d matrix slice of dimension [2L-1].
+
     """
     dl = dl.at[:].set(0)
     if mm < 0:
@@ -209,7 +216,8 @@ def _north_pole(dl, el, L, mm) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(2, 3))
 def _south_pole(dl, el, L, mm) -> jnp.ndarray:
-    r"""Compute Wigner-d matrix for edge case where theta index is located on the
+    r"""
+    Compute Wigner-d matrix for edge case where theta index is located on the
     south pole.
 
     Args:
@@ -223,6 +231,7 @@ def _south_pole(dl, el, L, mm) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Wigner-d matrix slice of dimension [2L-1].
+
     """
     dl = dl.at[:].set(0)
     if mm > 0:
@@ -234,7 +243,8 @@ def _south_pole(dl, el, L, mm) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def _el0(dl, L) -> jnp.ndarray:
-    r"""Compute Wigner-d matrix for edge case where the harmonic degree is 0
+    r"""
+    Compute Wigner-d matrix for edge case where the harmonic degree is 0
     (monopole term).
 
     Args:
@@ -244,6 +254,7 @@ def _el0(dl, L) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Wigner-d matrix slice of dimension [2L-1].
+
     """
     dl = dl.at[:].set(0)
     dl = dl.at[-1].set(1)
@@ -252,7 +263,8 @@ def _el0(dl, L) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(2, 3))
 def reindex(dl, el, L, mm) -> jnp.ndarray:
-    r"""Reorders indexing of Wigner-d matrix.
+    r"""
+    Reorders indexing of Wigner-d matrix.
 
     Reindexes the Wigner-d matrix to centre m values around L-1.
     The original indexing is given by
@@ -274,6 +286,7 @@ def reindex(dl, el, L, mm) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Wigner-d matrix slice of dimension [2L-1].
+
     """
     dl = dl.at[: L - 1].set(jnp.roll(dl, L - el - 1)[: L - 1])
     dl = dl.at[L - 1 :].set(jnp.roll(dl, -(L - el - 1))[L - 1 :])
diff --git a/s2fft/sampling/__init__.py b/s2fft/sampling/__init__.py
index 33a52347..3256c340 100644
--- a/s2fft/sampling/__init__.py
+++ b/s2fft/sampling/__init__.py
@@ -1,3 +1,1 @@
-from . import s2_samples
-from . import so3_samples
-from . import reindex
+from . import reindex, s2_samples, so3_samples
diff --git a/s2fft/sampling/reindex.py b/s2fft/sampling/reindex.py
index 8673897d..97d3bb85 100644
--- a/s2fft/sampling/reindex.py
+++ b/s2fft/sampling/reindex.py
@@ -1,11 +1,13 @@
-from jax import jit
-import jax.numpy as jnp
 from functools import partial
 
+import jax.numpy as jnp
+from jax import jit
+
 
 @partial(jit, static_argnums=(1))
 def flm_1d_to_2d_fast(flm_1d: jnp.ndarray, L: int) -> jnp.ndarray:
-    r"""Convert from 1D indexed harmnonic coefficients to 2D indexed coefficients (JAX).    
+    r"""
+    Convert from 1D indexed harmnonic coefficients to 2D indexed coefficients (JAX).
     
     Note:
         Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for 
@@ -31,6 +33,7 @@ def flm_1d_to_2d_fast(flm_1d: jnp.ndarray, L: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: 2D indexed harmonic coefficients.
+
     """
     flm_2d = jnp.zeros((L, 2 * L - 1), dtype=jnp.complex128)
     els = jnp.arange(L)
@@ -43,7 +46,8 @@ def flm_1d_to_2d_fast(flm_1d: jnp.ndarray, L: int) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def flm_2d_to_1d_fast(flm_2d: jnp.ndarray, L: int) -> jnp.ndarray:
-    r"""Convert from 2D indexed harmonic coefficients to 1D indexed coefficients (JAX).
+    r"""
+    Convert from 2D indexed harmonic coefficients to 1D indexed coefficients (JAX).
     
     Note:
         Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for 
@@ -69,6 +73,7 @@ def flm_2d_to_1d_fast(flm_2d: jnp.ndarray, L: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: 1D indexed harmonic coefficients.
+
     """
     flm_1d = jnp.zeros(L**2, dtype=jnp.complex128)
     els = jnp.arange(L)
@@ -81,7 +86,8 @@ def flm_2d_to_1d_fast(flm_2d: jnp.ndarray, L: int) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def flm_hp_to_2d_fast(flm_hp: jnp.ndarray, L: int) -> jnp.ndarray:
-    r"""Converts from HEALPix (healpy) indexed harmonic coefficients to 2D indexed
+    r"""
+    Converts from HEALPix (healpy) indexed harmonic coefficients to 2D indexed
     coefficients (JAX).
     
     Notes:
@@ -119,6 +125,7 @@ def flm_hp_to_2d_fast(flm_hp: jnp.ndarray, L: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: 2D indexed harmonic coefficients.
+
     """
     flm_2d = jnp.zeros((L, 2 * L - 1), dtype=jnp.complex128)
 
@@ -136,7 +143,8 @@ def flm_hp_to_2d_fast(flm_hp: jnp.ndarray, L: int) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def flm_2d_to_hp_fast(flm_2d: jnp.ndarray, L: int) -> jnp.ndarray:
-    r"""Converts from 2D indexed harmonic coefficients to HEALPix (healpy) indexed
+    r"""
+    Converts from 2D indexed harmonic coefficients to HEALPix (healpy) indexed
     coefficients (JAX).
     
     Note:
@@ -175,6 +183,7 @@ def flm_2d_to_hp_fast(flm_2d: jnp.ndarray, L: int) -> jnp.ndarray:
         
     Returns:
         jnp.ndarray: HEALPix indexed harmonic coefficients.
+
     """
     flm_hp = jnp.zeros(int(L * (L + 1) / 2), dtype=jnp.complex128)
 
diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py
index b8f7cfae..06d1996b 100644
--- a/s2fft/sampling/s2_samples.py
+++ b/s2fft/sampling/s2_samples.py
@@ -1,9 +1,11 @@
-import numpy as np
 from typing import Tuple
 
+import numpy as np
+
 
 def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int:
-    r"""Number of :math:`\theta` samples for sampling scheme at specified resolution.
+    r"""
+    Number of :math:`\theta` samples for sampling scheme at specified resolution.
 
     Args:
         L (int, optional): Harmonic band-limit.  Required if sampling not healpix.
@@ -24,8 +26,8 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int:
 
     Returns:
         int: Number of :math:`\theta` samples of sampling scheme at given resolution.
-    """
 
+    """
     if sampling.lower() != "healpix" and L is None:
         raise ValueError(
             f"Sampling scheme sampling={sampling} with L={L} not supported"
@@ -53,7 +55,8 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int:
 
 
 def ntheta_extension(L: int, sampling: str = "mw") -> int:
-    r"""Number of :math:`\theta` samples for MW/MWSS sampling when extended to
+    r"""
+    Number of :math:`\theta` samples for MW/MWSS sampling when extended to
     :math:`2\pi`.
 
     Args:
@@ -69,7 +72,6 @@ def ntheta_extension(L: int, sampling: str = "mw") -> int:
         int: Number of :math:`\theta` samples when extended to :math:`2\pi`.
 
     """
-
     if sampling.lower() == "mw":
         return 2 * L - 1
 
@@ -83,7 +85,8 @@ def ntheta_extension(L: int, sampling: str = "mw") -> int:
 
 
 def nphi_equiang(L: int, sampling: str = "mw") -> int:
-    r"""Number of :math:`\phi` samples for equiangular sampling scheme at specified
+    r"""
+    Number of :math:`\phi` samples for equiangular sampling scheme at specified
     resolution.
 
     Number of samples is independent of :math:`\theta` since equiangular sampling
@@ -102,8 +105,8 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int:
 
     Returns:
         int: Number of :math:`\phi` samples.
-    """
 
+    """
     if sampling.lower() in ["mw", "gl"]:
         return 2 * L - 1
 
@@ -123,7 +126,8 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int:
 
 
 def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> Tuple[int, int]:
-    r"""Shape of intermediate array, before/after latitudinal step.
+    r"""
+    Shape of intermediate array, before/after latitudinal step.
 
     Args:
         L (int): Harmonic band-limit.
@@ -139,8 +143,8 @@ def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> Tuple[int, int
     Returns:
         Tuple[int,int]: Shape of ftm array, i.e. :math:`[n_{\theta}, n_{\phi}]`. Note
         that here "healpix" defaults to :math:`2L = 4nside` phi samples for ftm.
-    """
 
+    """
     if sampling.lower() in ["mwss", "healpix"]:
         return ntheta(L, sampling, nside), 2 * L
 
@@ -154,7 +158,8 @@ def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> Tuple[int, int
 
 
 def nphi_equitorial_band(nside: int) -> int:
-    r"""Number of :math:`\phi` samples within the equitorial band for
+    r"""
+    Number of :math:`\phi` samples within the equitorial band for
     HEALPix sampling scheme.
 
     Args:
@@ -162,12 +167,14 @@ def nphi_equitorial_band(nside: int) -> int:
 
     Returns:
         int: Number of :math:`\phi` samples.
+
     """
     return 4 * nside
 
 
 def nphi_ring(t: int, nside: int = None) -> int:
-    r"""Number of :math:`\phi` samples for HEALPix sampling on given :math:`\theta`
+    r"""
+    Number of :math:`\phi` samples for HEALPix sampling on given :math:`\theta`
     ring.
 
     Args:
@@ -180,8 +187,8 @@ def nphi_ring(t: int, nside: int = None) -> int:
 
     Returns:
         int: Number of :math:`\phi` samples on given :math:`\theta` ring.
-    """
 
+    """
     if (t >= 0) and (t < nside - 1):
         return 4 * (t + 1)
 
@@ -196,7 +203,8 @@ def nphi_ring(t: int, nside: int = None) -> int:
 
 
 def thetas(L: int = None, sampling: str = "mw", nside: int = None) -> np.ndarray:
-    r"""Compute :math:`\theta` samples for given sampling scheme.
+    r"""
+    Compute :math:`\theta` samples for given sampling scheme.
 
     Args:
         L (int, optional): Harmonic band-limit.  Required if sampling not healpix.
@@ -210,6 +218,7 @@ def thetas(L: int = None, sampling: str = "mw", nside: int = None) -> np.ndarray
 
     Returns:
         np.ndarray: Array of :math:`\theta` samples for given sampling scheme.
+
     """
     if sampling.lower() == "gl":
         return np.flip(np.arccos(np.polynomial.legendre.leggauss(L)[0]))
@@ -222,7 +231,8 @@ def thetas(L: int = None, sampling: str = "mw", nside: int = None) -> np.ndarray
 def t2theta(
     t: int, L: int = None, sampling: str = "mw", nside: int = None
 ) -> np.ndarray:
-    r"""Convert index to :math:`\theta` angle for sampling scheme.
+    r"""
+    Convert index to :math:`\theta` angle for sampling scheme.
 
     Args:
         t (int): :math:`\theta` index.
@@ -245,8 +255,8 @@ def t2theta(
 
     Returns:
         np.ndarray: :math:`\theta` angle(s) for passed index or indices.
-    """
 
+    """
     if sampling.lower() != "healpix" and L is None:
         raise ValueError(
             f"Sampling scheme sampling={sampling} with L={L} not supported"
@@ -274,7 +284,8 @@ def t2theta(
 
 
 def _t2theta_healpix(t: int, nside: int) -> np.ndarray:
-    r"""Convert (ring) index to :math:`\theta` angle for HEALPix sampling scheme.
+    r"""
+    Convert (ring) index to :math:`\theta` angle for HEALPix sampling scheme.
 
     Args:
         t (int): :math:`\theta` index.
@@ -283,8 +294,8 @@ def _t2theta_healpix(t: int, nside: int) -> np.ndarray:
 
     Returns:
         np.ndarray: :math:`\theta` angle(s) for passed HEALPix (ring) index or indices.
-    """
 
+    """
     z = np.zeros_like(t)
     z[t < nside - 1] = 1 - (t[t < nside - 1] + 1) ** 2 / (3 * nside**2)
 
@@ -300,7 +311,8 @@ def _t2theta_healpix(t: int, nside: int) -> np.ndarray:
 
 
 def phis_ring(t: int, nside: int) -> np.ndarray:
-    r"""Compute :math:`\phi` samples for given :math:`\theta` HEALPix ring.
+    r"""
+    Compute :math:`\phi` samples for given :math:`\theta` HEALPix ring.
 
     Args:
         t (int): :math:`\theta` index.
@@ -309,15 +321,16 @@ def phis_ring(t: int, nside: int) -> np.ndarray:
 
     Returns:
         np.ndarray: :math:`\phi` angles.
-    """
 
+    """
     p = np.arange(0, nphi_ring(t, nside)).astype(np.float64)
 
     return p2phi_ring(t, p, nside)
 
 
 def p2phi_ring(t: int, p: int, nside: int) -> np.ndarray:
-    r"""Convert index to :math:`\phi` angle for HEALPix for given :math:`\theta` ring.
+    r"""
+    Convert index to :math:`\phi` angle for HEALPix for given :math:`\theta` ring.
 
     Args:
         t (int): :math:`\theta` index of ring.
@@ -328,8 +341,8 @@ def p2phi_ring(t: int, p: int, nside: int) -> np.ndarray:
 
     Returns:
         np.ndarray: :math:`\phi` angle.
-    """
 
+    """
     shift = 1 / 2
     if (t + 1 >= nside) & (t + 1 <= 3 * nside):
         shift *= (t - nside + 2) % 2
@@ -343,7 +356,8 @@ def p2phi_ring(t: int, p: int, nside: int) -> np.ndarray:
 
 
 def phis_equiang(L: int, sampling: str = "mw") -> np.ndarray:
-    r"""Compute :math:`\phi` samples for equiangular sampling scheme.
+    r"""
+    Compute :math:`\phi` samples for equiangular sampling scheme.
 
     Args:
         L (int, optional): Harmonic band-limit.
@@ -353,6 +367,7 @@ def phis_equiang(L: int, sampling: str = "mw") -> np.ndarray:
 
     Returns:
         np.ndarray: Array of :math:`\phi` samples for given sampling scheme.
+
     """
     p = np.arange(0, nphi_equiang(L, sampling))
 
@@ -360,7 +375,8 @@ def phis_equiang(L: int, sampling: str = "mw") -> np.ndarray:
 
 
 def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray:
-    r"""Convert index to :math:`\phi` angle for sampling scheme.
+    r"""
+    Convert index to :math:`\phi` angle for sampling scheme.
 
     Args:
         L (int, optional): Harmonic band-limit.
@@ -377,8 +393,8 @@ def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray:
 
     Returns:
         np.ndarray: :math:`\phi` sample(s) for given sampling scheme.
-    """
 
+    """
     if sampling.lower() in ["mw", "gl"]:
         return 2 * p * np.pi / (2 * L - 1)
 
@@ -402,7 +418,8 @@ def ring_phase_shift_hp(
     forward: bool = False,
     reality: bool = False,
 ) -> np.ndarray:
-    r"""Generates a phase shift vector for HEALPix for a given :math:`\theta` ring.
+    r"""
+    Generates a phase shift vector for HEALPix for a given :math:`\theta` ring.
 
     Args:
         L (int, optional): Harmonic band-limit.
@@ -420,6 +437,7 @@ def ring_phase_shift_hp(
 
     Returns:
         np.ndarray: Vector of phase shifts with shape :math:`[2L-1]`.
+
     """
     phi_offset = p2phi_ring(t, 0, nside)
     sign = -1 if forward else 1
@@ -428,7 +446,8 @@ def ring_phase_shift_hp(
 
 
 def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> Tuple[int]:
-    r"""Shape of spherical signal.
+    r"""
+    Shape of spherical signal.
 
     Args:
         L (int, optional): Harmonic band-limit.
@@ -442,8 +461,8 @@ def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> Tuple[int
     Returns:
         Tuple[int]: Pixel-space array dimensions with shape :math:`[n_{\theta}, n_{\phi}]`.
         Note that "healpix" is instead indexed by a 1D array, with standard conventions.
-    """
 
+    """
     if sampling.lower() != "healpix" and L is None:
         raise ValueError(
             f"Sampling scheme sampling={sampling} with L={L} not supported"
@@ -462,20 +481,22 @@ def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> Tuple[int
 
 
 def flm_shape(L: int) -> Tuple[int, int]:
-    r"""Standard shape of harmonic coefficients.
+    r"""
+    Standard shape of harmonic coefficients.
 
     Args:
         L (int, optional): Harmonic band-limit.
 
     Returns:
         Tuple[int]: Sampling array shape, with indexing :math:`[\ell, m]`.
-    """
 
+    """
     return L, 2 * L - 1
 
 
 def elm2ind(el: int, m: int) -> int:
-    r"""Convert from spherical harmonic 2D indexing of :math:`(\ell,m)` to 1D index.
+    r"""
+    Convert from spherical harmonic 2D indexing of :math:`(\ell,m)` to 1D index.
 
     1D index is defined by `el**2 + el + m`.
 
@@ -489,13 +510,14 @@ def elm2ind(el: int, m: int) -> int:
 
     Returns:
         int: Corresponding 1D index value.
-    """
 
+    """
     return el**2 + el + m
 
 
 def ind2elm(ind: int) -> tuple:
-    r"""Convert from 1D spherical harmonic index to 2D index of :math:`(\ell,m)`.
+    r"""
+    Convert from 1D spherical harmonic index to 2D index of :math:`(\ell,m)`.
 
     Warning:
         Note that 1D storage of spherical harmonic coefficients is *not* the default.
@@ -505,8 +527,8 @@ def ind2elm(ind: int) -> tuple:
 
     Returns:
         tuple: `(el,m)` defining spherical harmonic degree and order.
-    """
 
+    """
     el = np.floor(np.sqrt(ind))
 
     m = ind - el**2 - el
@@ -515,20 +537,22 @@ def ind2elm(ind: int) -> tuple:
 
 
 def ncoeff(L: int) -> int:
-    """Number of spherical harmonic coefficients for given band-limit L.
+    """
+    Number of spherical harmonic coefficients for given band-limit L.
 
     Args:
         L (int, optional): Harmonic band-limit.
 
     Returns:
         int: Number of spherical harmonic coefficients.
-    """
 
+    """
     return elm2ind(L - 1, L - 1) + 1
 
 
 def hp_ang2pix(nside: int, theta: float, phi: float) -> int:
-    r"""Convert angles to HEALPix index for HEALPix ring ordering scheme.
+    r"""
+    Convert angles to HEALPix index for HEALPix ring ordering scheme.
 
     Args:
         nside (int): HEALPix Nside resolution parameter.
@@ -539,15 +563,16 @@ def hp_ang2pix(nside: int, theta: float, phi: float) -> int:
 
     Returns:
         int: HEALPix map index for ring ordering scheme.
-    """
 
+    """
     z = np.cos(theta)
 
     return _hp_zphi2pix(nside, z, phi)
 
 
 def _hp_zphi2pix(nside: int, z: float, phi: float) -> int:
-    r"""Convert angles to HEALPix index for HEALPix ring ordering scheme, using
+    r"""
+    Convert angles to HEALPix index for HEALPix ring ordering scheme, using
     :math:`z=\cos(\theta)`.
 
     Note:
@@ -562,8 +587,8 @@ def _hp_zphi2pix(nside: int, z: float, phi: float) -> int:
 
     Returns:
         int: HEALPix map index for ring ordering scheme.
-    """
 
+    """
     tt = 2 * phi / np.pi
     za = np.abs(z)
     nl2 = int(2 * nside)
@@ -600,7 +625,8 @@ def _hp_zphi2pix(nside: int, z: float, phi: float) -> int:
 
 
 def hp_getidx(L: int, el: int, m: int) -> int:
-    r"""Compute HEALPix harmonic index.
+    r"""
+    Compute HEALPix harmonic index.
 
     Warning:
         Note that the harmonic band-limit `L` differs to the HEALPix :math`\ell_{\text{max}}` convention,
@@ -615,12 +641,14 @@ def hp_getidx(L: int, el: int, m: int) -> int:
 
     Returns:
         int: Corresponding index for RING ordered HEALPix.
+
     """
     return m * (2 * L - 1 - m) // 2 + el
 
 
 def flm_2d_to_1d(flm_2d: np.ndarray, L: int) -> np.ndarray:
-    r"""Convert from 2D indexed harmonic coefficients to 1D indexed coefficients.
+    r"""
+    Convert from 2D indexed harmonic coefficients to 1D indexed coefficients.
     
     Note:
         Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for 
@@ -646,12 +674,13 @@ def flm_2d_to_1d(flm_2d: np.ndarray, L: int) -> np.ndarray:
 
     Returns:
         np.ndarray: 1D indexed harmonic coefficients.
+
     """
     flm_1d = np.zeros(ncoeff(L), dtype=np.complex128)
 
     if len(flm_2d.shape) != 2:
         if len(flm_2d.shape) == 1:
-            raise ValueError(f"Flm is already 1D indexed")
+            raise ValueError("Flm is already 1D indexed")
         else:
             raise ValueError(
                 f"Cannot convert flm of dimension {flm_2d.shape} to 1D indexing"
@@ -665,7 +694,8 @@ 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.    
+    r"""
+    Convert from 1D indexed harmnonic coefficients to 2D indexed coefficients.
     
     Note:
         Storage conventions for harmonic coefficients :math:`flm_{(\ell,m)}`, for 
@@ -691,13 +721,13 @@ def flm_1d_to_2d(flm_1d: np.ndarray, L: int) -> np.ndarray:
 
     Returns:
         np.ndarray: 2D indexed harmonic coefficients.
-    """
 
+    """
     flm_2d = np.zeros(flm_shape(L), dtype=np.complex128)
 
     if len(flm_1d.shape) != 1:
         if len(flm_1d.shape) == 2:
-            raise ValueError(f"Flm is already 2D indexed")
+            raise ValueError("Flm is already 2D indexed")
         else:
             raise ValueError(
                 f"Cannot convert flm of dimension {flm_2d.shape} to 2D indexing"
@@ -711,7 +741,8 @@ def flm_1d_to_2d(flm_1d: np.ndarray, L: int) -> np.ndarray:
 
 
 def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray:
-    r"""Converts from HEALPix (healpy) indexed harmonic coefficients to 2D indexed
+    r"""
+    Converts from HEALPix (healpy) indexed harmonic coefficients to 2D indexed
     coefficients.
     
     Notes:
@@ -749,11 +780,12 @@ def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray:
 
     Returns:
         np.ndarray: 2D indexed harmonic coefficients.
+
     """
     flm_2d = np.zeros(flm_shape(L), dtype=np.complex128)
 
     if len(flm_hp.shape) != 1:
-        raise ValueError(f"Healpix indexed flms are not flat")
+        raise ValueError("Healpix indexed flms are not flat")
 
     for el in range(L):
         flm_2d[el, L - 1 + 0] = flm_hp[hp_getidx(L, el, 0)]
@@ -765,7 +797,8 @@ def flm_hp_to_2d(flm_hp: np.ndarray, L: int) -> np.ndarray:
 
 
 def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray:
-    r"""Converts from 2D indexed harmonic coefficients to HEALPix (healpy) indexed
+    r"""
+    Converts from 2D indexed harmonic coefficients to HEALPix (healpy) indexed
     coefficients.
     
     Note:
@@ -806,11 +839,10 @@ def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray:
         np.ndarray: HEALPix indexed harmonic coefficients.
         
     """
-
     flm_hp = np.zeros(int(L * (L + 1) / 2), dtype=np.complex128)
 
     if len(flm_hp.shape) != 1:
-        raise ValueError(f"HEALPix indexed flms are not flat")
+        raise ValueError("HEALPix indexed flms are not flat")
 
     for el in range(L):
         for m in range(0, el + 1):
@@ -820,7 +852,8 @@ def flm_2d_to_hp(flm_2d: np.ndarray, L: int) -> np.ndarray:
 
 
 def lm2lm_hp(flm: np.ndarray, L: int) -> np.ndarray:
-    r"""Converts from 1D indexed harmonic coefficients to HEALPix (healpy) indexed
+    r"""
+    Converts from 1D indexed harmonic coefficients to HEALPix (healpy) indexed
     coefficients.
 
     Note:
diff --git a/s2fft/sampling/so3_samples.py b/s2fft/sampling/so3_samples.py
index 63a29784..194c2c4d 100644
--- a/s2fft/sampling/so3_samples.py
+++ b/s2fft/sampling/so3_samples.py
@@ -1,12 +1,15 @@
+from typing import Tuple
+
 import numpy as np
+
 from s2fft.sampling import s2_samples as samples
-from typing import Tuple
 
 
 def f_shape(
     L: int, N: int, sampling: str = "mw", nside: int = None
 ) -> Tuple[int, int, int]:
-    r"""Computes the pixel-space sampling shape for signal on the rotation group
+    r"""
+    Computes the pixel-space sampling shape for signal on the rotation group
     :math:`SO(3)`.
 
     Note:
@@ -34,6 +37,7 @@ def f_shape(
     Returns:
         Tuple[int,int,int]: Shape of pixel-space sampling of rotation group
         :math:`SO(3)`.
+
     """
     if sampling in ["mw", "mwss", "dh", "gl"]:
         return _ngamma(N), _nbeta(L, sampling), _nalpha(L, sampling)
@@ -46,7 +50,8 @@ def f_shape(
 
 
 def flmn_shape(L: int, N: int) -> Tuple[int, int, int]:
-    r"""Computes the shape of Wigner coefficients for signal on the rotation group
+    r"""
+    Computes the shape of Wigner coefficients for signal on the rotation group
     :math:`SO(3)`.
 
     Args:
@@ -57,6 +62,7 @@ def flmn_shape(L: int, N: int) -> Tuple[int, int, int]:
     Returns:
         Tuple[int,int,int]: Shape of Wigner space sampling of rotation group
             :math:`SO(3)`.
+
     """
     return 2 * N - 1, L, 2 * L - 1
 
@@ -64,7 +70,8 @@ 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]:
-    r"""Computes the shape of Wigner coefficients for signal on the rotation group
+    r"""
+    Computes the shape of Wigner coefficients for signal on the rotation group
     :math:`SO(3)`.
 
     Args:
@@ -83,8 +90,8 @@ def fnab_shape(
     Returns:
         Tuple[int,int,int]: Shape of Wigner space sampling of rotation group
             :math:`SO(3)`.
-    """
 
+    """
     if sampling.lower() in ["mwss", "healpix"]:
         return _ngamma(N), samples.ntheta(L, sampling, nside), 2 * L
 
@@ -98,7 +105,8 @@ def fnab_shape(
 
 
 def flmn_shape_1d(L: int, N: int) -> int:
-    r"""Computes the number of non-zero Wigner coefficients.
+    r"""
+    Computes the number of non-zero Wigner coefficients.
 
     Args:
         L (int): Harmonic band-limit.
@@ -107,12 +115,14 @@ def flmn_shape_1d(L: int, N: int) -> int:
 
     Returns:
         int: Total number of non-zero Wigner coefficients.
+
     """
     return (2 * N - 1) * L * L
 
 
 def _nalpha(L: int, sampling: str = "mw") -> int:
-    r"""Computes the number of :math:`\alpha` samples.
+    r"""
+    Computes the number of :math:`\alpha` samples.
 
     Args:
         L (int): Harmonic band-limit.
@@ -125,6 +135,7 @@ def _nalpha(L: int, sampling: str = "mw") -> int:
 
     Returns:
         int: Number of :math:`\alpha` samples.
+
     """
     if sampling.lower() in ["mw", "dh", "gl"]:
         return 2 * L - 1
@@ -137,7 +148,8 @@ def _nalpha(L: int, sampling: str = "mw") -> int:
 
 
 def _nbeta(L: int, sampling: str = "mw") -> int:
-    r"""Computes the number of :math:`\beta` samples.
+    r"""
+    Computes the number of :math:`\beta` samples.
 
     Args:
         L (int): Harmonic band-limit.
@@ -150,6 +162,7 @@ def _nbeta(L: int, sampling: str = "mw") -> int:
 
     Returns:
         int: Number of :math:`\beta` samples.
+
     """
     if sampling.lower() in ["mw", "gl"]:
         return L
@@ -165,19 +178,22 @@ def _nbeta(L: int, sampling: str = "mw") -> int:
 
 
 def _ngamma(N: int) -> int:
-    r"""Computes the number of :math:`\gamma` samples.
+    r"""
+    Computes the number of :math:`\gamma` samples.
 
     Args:
         N (int): Directional band-limit.
 
     Returns:
         int: Number of :math:`\gamma` samples, by default :math:`2N-1`.
+
     """
     return 2 * N - 1
 
 
 def elmn2ind(el: int, m: int, n: int, L: int, N: int) -> int:
-    """Convert from Wigner space 3D indexing of :math:`(\ell,m, n)` to 1D index.
+    r"""
+    Convert from Wigner space 3D indexing of :math:`(\ell,m, n)` to 1D index.
 
     Args:
         el (int): Harmonic degree :math:`\ell`.
@@ -192,6 +208,7 @@ def elmn2ind(el: int, m: int, n: int, L: int, N: int) -> int:
 
     Returns:
         int: Corresponding 1D index in Wigner space.
+
     """
     n_offset = (N - 1 + n) * L * L
     el_offset = el * el
@@ -199,7 +216,8 @@ def elmn2ind(el: int, m: int, n: int, L: int, N: int) -> int:
 
 
 def flmn_3d_to_1d(flmn_3d: np.ndarray, L: int, N: int) -> np.ndarray:
-    r"""Convert from 3D indexed Wigner coefficients to 1D indexed coefficients.
+    r"""
+    Convert from 3D indexed Wigner coefficients to 1D indexed coefficients.
 
     Args:
         flm_3d (np.ndarray): 3D indexed Wigner coefficients, index order
@@ -216,11 +234,12 @@ def flmn_3d_to_1d(flmn_3d: np.ndarray, L: int, N: int) -> np.ndarray:
 
     Returns:
         np.ndarray: 1D indexed Wigner coefficients, C flatten index priority :math:`n, \ell, m`.
+
     """
     flmn_1d = np.zeros(flmn_shape_1d(L, N), dtype=np.complex128)
 
     if len(flmn_3d.shape) == 1:
-        raise ValueError(f"flmn is already 1D indexed")
+        raise ValueError("flmn is already 1D indexed")
     elif len(flmn_3d.shape) != 3:
         raise ValueError(
             f"Cannot convert flmn of dimension {flmn_3d.shape} to 1D indexing"
@@ -235,7 +254,8 @@ def flmn_3d_to_1d(flmn_3d: np.ndarray, L: int, N: int) -> np.ndarray:
 
 
 def flmn_1d_to_3d(flmn_1d: np.ndarray, L: int, N: int) -> np.ndarray:
-    r"""Convert from 1D indexed Wigner coefficients to 3D indexed coefficients.
+    r"""
+    Convert from 1D indexed Wigner coefficients to 3D indexed coefficients.
 
     Args:
         flm_1d (np.ndarray): 1D indexed Wigner coefficients, C flatten index priority
@@ -252,11 +272,12 @@ def flmn_1d_to_3d(flmn_1d: np.ndarray, L: int, N: int) -> np.ndarray:
 
     Returns:
         np.ndarray: 3D indexed Wigner coefficients, index order :math:`[\ell, m, n]`.
+
     """
     flmn_3d = np.zeros(flmn_shape(L, N), dtype=np.complex128)
 
     if len(flmn_1d.shape) == 3:
-        raise ValueError(f"Flmn is already 3D indexed")
+        raise ValueError("Flmn is already 3D indexed")
     elif len(flmn_1d.shape) != 1:
         raise ValueError(
             f"Cannot convert flmn of dimension {flmn_1d.shape} to 3D indexing"
diff --git a/s2fft/transforms/__init__.py b/s2fft/transforms/__init__.py
index c15947b4..a7ed3a2b 100644
--- a/s2fft/transforms/__init__.py
+++ b/s2fft/transforms/__init__.py
@@ -1,4 +1,1 @@
-from . import spherical
-from . import wigner
-from . import otf_recursions
-from . import c_backend_spherical
+from . import c_backend_spherical, otf_recursions, spherical, wigner
diff --git a/s2fft/transforms/c_backend_spherical.py b/s2fft/transforms/c_backend_spherical.py
index 3c4d8a25..5b6df8f5 100644
--- a/s2fft/transforms/c_backend_spherical.py
+++ b/s2fft/transforms/c_backend_spherical.py
@@ -1,12 +1,13 @@
-from jax import custom_vjp
-import numpy as np
+import healpy
 import jax.numpy as jnp
-from s2fft.utils import quadrature_jax
-from s2fft.sampling import reindex
+import numpy as np
 
 # C backend functions for which to provide JAX frontend.
 import pyssht
-import healpy
+from jax import custom_vjp
+
+from s2fft.sampling import reindex
+from s2fft.utils import quadrature_jax
 
 
 @custom_vjp
@@ -18,7 +19,8 @@ def ssht_inverse(
     ssht_sampling: int = 0,
     _ssht_backend: int = 1,
 ) -> jnp.ndarray:
-    r"""Compute the inverse spin-spherical harmonic transform (SSHT JAX).
+    r"""
+    Compute the inverse spin-spherical harmonic transform (SSHT JAX).
 
     SSHT is a C library which implements the spin-spherical harmonic transform outlined in
     McEwen & Wiaux 2011 [1]. We make use of their python bindings for which we provide
@@ -49,6 +51,7 @@ def ssht_inverse(
     Note:
         [1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”
             IEEE Transactions on Signal Processing 59 (2011): 5876-5887.
+
     """
     sampling_str = ["MW", "MWSS", "DH", "GL"]
     flm_1d = reindex.flm_2d_to_1d_fast(flm, L)
@@ -73,13 +76,13 @@ def _ssht_inverse_fwd(
     ssht_sampling: int = 0,
     _ssht_backend: int = 1,
 ):
-    """Private function which implements the forward pass for inverse jax_ssht"""
+    """Private function which implements the forward pass for inverse jax_ssht."""
     res = ([], L, spin, reality, ssht_sampling, _ssht_backend)
     return ssht_inverse(flm, L, spin, reality, ssht_sampling, _ssht_backend), res
 
 
 def _ssht_inverse_bwd(res, f):
-    """Private function which implements the backward pass for inverse jax_ssht"""
+    """Private function which implements the backward pass for inverse jax_ssht."""
     _, L, spin, reality, ssht_sampling, _ssht_backend = res
     sampling_str = ["MW", "MWSS", "DH", "GL"]
     _backend = "SSHT" if _ssht_backend == 0 else "ducc0"
@@ -135,7 +138,8 @@ def ssht_forward(
     ssht_sampling: int = 0,
     _ssht_backend: int = 1,
 ) -> jnp.ndarray:
-    r"""Compute the forward spin-spherical harmonic transform (SSHT JAX).
+    r"""
+    Compute the forward spin-spherical harmonic transform (SSHT JAX).
 
     SSHT is a C library which implements the spin-spherical harmonic transform outlined in
     McEwen & Wiaux 2011 [1]. We make use of their python bindings for which we provide
@@ -166,6 +170,7 @@ def ssht_forward(
     Note:
         [1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”
             IEEE Transactions on Signal Processing 59 (2011): 5876-5887.
+
     """
     sampling_str = ["MW", "MWSS", "DH", "GL"]
     _backend = "SSHT" if _ssht_backend == 0 else "ducc0"
@@ -190,13 +195,13 @@ def _ssht_forward_fwd(
     ssht_sampling: int = 0,
     _ssht_backend: int = 1,
 ):
-    """Private function which implements the forward pass for forward jax_ssht"""
+    """Private function which implements the forward pass for forward jax_ssht."""
     res = ([], L, spin, reality, ssht_sampling, _ssht_backend)
     return ssht_forward(f, L, spin, reality, ssht_sampling, _ssht_backend), res
 
 
 def _ssht_forward_bwd(res, flm):
-    """Private function which implements the backward pass for forward jax_ssht"""
+    """Private function which implements the backward pass for forward jax_ssht."""
     _, L, spin, reality, ssht_sampling, _ssht_backend = res
     sampling_str = ["MW", "MWSS", "DH", "GL"]
     _backend = "SSHT" if _ssht_backend == 0 else "ducc0"
@@ -238,7 +243,8 @@ def _ssht_forward_bwd(res, flm):
 
 @custom_vjp
 def healpy_inverse(flm: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
-    r"""Compute the inverse scalar real spherical harmonic transform (HEALPix JAX).
+    r"""
+    Compute the inverse scalar real spherical harmonic transform (HEALPix JAX).
 
     HEALPix is a C++ library which implements the scalar spherical harmonic transform
     outlined in [1]. We make use of their healpy python bindings for which we provide
@@ -260,6 +266,7 @@ def healpy_inverse(flm: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
         [1] Gorski, Krzysztof M., et al. "HEALPix: A framework for high-resolution
         discretization and fast analysis of data distributed on the sphere." The
         Astrophysical Journal 622.2 (2005): 759
+
     """
     flm = reindex.flm_2d_to_hp_fast(flm, L)
     f = jnp.array(healpy.alm2map(np.array(flm), lmax=L - 1, nside=nside))
@@ -267,13 +274,13 @@ def healpy_inverse(flm: jnp.ndarray, L: int, nside: int) -> jnp.ndarray:
 
 
 def _healpy_inverse_fwd(flm: jnp.ndarray, L: int, nside: int):
-    """Private function which implements the forward pass for inverse jax_healpy"""
+    """Private function which implements the forward pass for inverse jax_healpy."""
     res = ([], L, nside)
     return healpy_inverse(flm, L, nside), res
 
 
 def _healpy_inverse_bwd(res, f):
-    """Private function which implements the backward pass for inverse jax_healpy"""
+    """Private function which implements the backward pass for inverse jax_healpy."""
     _, L, nside = res
     f_new = f * (12 * nside**2) / (4 * jnp.pi)
     flm_out = jnp.array(
@@ -292,7 +299,8 @@ def _healpy_inverse_bwd(res, f):
 
 @custom_vjp
 def healpy_forward(f: jnp.ndarray, L: int, nside: int, iter: int = 3) -> jnp.ndarray:
-    r"""Compute the forward scalar spherical harmonic transform (healpy JAX).
+    r"""
+    Compute the forward scalar spherical harmonic transform (healpy JAX).
 
     HEALPix is a C++ library which implements the scalar spherical harmonic transform
     outlined in [1]. We make use of their healpy python bindings for which we provide
@@ -318,19 +326,20 @@ def healpy_forward(f: jnp.ndarray, L: int, nside: int, iter: int = 3) -> jnp.nda
         [1] Gorski, Krzysztof M., et al. "HEALPix: A framework for high-resolution
         discretization and fast analysis of data distributed on the sphere." The
         Astrophysical Journal 622.2 (2005): 759
+
     """
     flm = jnp.array(healpy.map2alm(np.array(f), lmax=L - 1, iter=iter))
     return reindex.flm_hp_to_2d_fast(flm, L)
 
 
 def _healpy_forward_fwd(f: jnp.ndarray, L: int, nside: int, iter: int = 3):
-    """Private function which implements the forward pass for forward jax_healpy"""
+    """Private function which implements the forward pass for forward jax_healpy."""
     res = ([], L, nside, iter)
     return healpy_forward(f, L, nside, iter), res
 
 
 def _healpy_forward_bwd(res, flm):
-    """Private function which implements the backward pass for forward jax_healpy"""
+    """Private function which implements the backward pass for forward jax_healpy."""
     _, L, nside, _ = res
     flm_new = reindex.flm_2d_to_hp_fast(flm, L)
     f = jnp.array(
diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py
index 530c3501..50718eb2 100644
--- a/s2fft/transforms/otf_recursions.py
+++ b/s2fft/transforms/otf_recursions.py
@@ -1,14 +1,16 @@
-import numpy as np
-import jax.numpy as jnp
-import jax.lax as lax
-from jax import jit, pmap, local_device_count
 from functools import partial
 from typing import List
-from s2fft.sampling import s2_samples as samples
+
+import jax.lax as lax
+import jax.numpy as jnp
+import numpy as np
+from jax import jit, local_device_count, pmap
+
 from s2fft.recursions.price_mcewen import (
     generate_precomputes,
     generate_precomputes_jax,
 )
+from s2fft.sampling import s2_samples as samples
 
 
 def inverse_latitudinal_step(
@@ -22,7 +24,8 @@ def inverse_latitudinal_step(
     precomps: List = None,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Evaluate the wigner-d recursion inverse latitundinal step over :math:`\theta`.
+    r"""
+    Evaluate the wigner-d recursion inverse latitundinal step over :math:`\theta`.
     This approach is a heavily engineerd version of the Price & McEwen recursion found in
     :func:`~s2fft.wigner.price_mcewen`, which has at most of :math:`\mathcal{O}(L^2)`
     memory footprint.
@@ -60,6 +63,7 @@ def inverse_latitudinal_step(
 
     Returns:
         np.ndarray: Coefficients ftm with indexing :math:`[\theta, m]`.
+
     """
     mm = -spin
     ntheta = len(beta)
@@ -171,7 +175,8 @@ def inverse_latitudinal_step_jax(
     spmd: bool = False,
     L_lower: int = 0,
 ) -> jnp.ndarray:
-    r"""Evaluate the wigner-d recursion inverse latitundinal step over :math:`\theta`.
+    r"""
+    Evaluate the wigner-d recursion inverse latitundinal step over :math:`\theta`.
     This approach is a heavily engineerd version of the Price & McEwen recursion found in
     :func:`~s2fft.wigner.price_mcewen`, which has at most of :math:`\mathcal{O}(L^2)`
     memory footprint. This is a JAX implementation of :func:`~inverse_latitudinal_step`.
@@ -219,8 +224,8 @@ def inverse_latitudinal_step_jax(
         harmonic bandlimits L this is inefficient as the I/O overhead for communication
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
-    """
 
+    """
     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
@@ -379,7 +384,16 @@ def eval_recursion_step(
                 ).reshape(ntheta, ftm.shape[-1])
 
             else:
-                (ftm, dl_entry, dl_iter, lrenorm, indices, omc, c, s,) = lax.fori_loop(
+                (
+                    ftm,
+                    dl_entry,
+                    dl_iter,
+                    lrenorm,
+                    indices,
+                    omc,
+                    c,
+                    s,
+                ) = lax.fori_loop(
                     2,
                     L - 1 + i,
                     pm_recursion_step,
@@ -417,7 +431,8 @@ def forward_latitudinal_step(
     precomps: List = None,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Evaluate the wigner-d recursion forward latitundinal step over :math:`\theta`.
+    r"""
+    Evaluate the wigner-d recursion forward latitundinal step over :math:`\theta`.
     This approach is a heavily engineerd version of the Price & McEwen recursion found in
     :func:`~s2fft.wigner.price_mcewen`, which has at most of :math:`\mathcal{O}(L^2)`
     memory footprint.
@@ -455,6 +470,7 @@ def forward_latitudinal_step(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients.
+
     """
     mm = -spin  # switch to match convention
     ntheta = len(beta)  # Number of theta samples
@@ -573,7 +589,8 @@ def forward_latitudinal_step_jax(
     spmd: bool = False,
     L_lower: int = 0,
 ) -> jnp.ndarray:
-    r"""Evaluate the wigner-d recursion forward latitundinal step over :math:`\theta`.
+    r"""
+    Evaluate the wigner-d recursion forward latitundinal step over :math:`\theta`.
     This approach is a heavily engineerd version of the Price & McEwen recursion found in
     :func:`~s2fft.wigner.price_mcewen`, which has at most of :math:`\mathcal{O}(L^2)`
     memory footprint. This is a JAX implementation of :func:`~forward_latitudinal_step`.
@@ -621,6 +638,7 @@ def forward_latitudinal_step_jax(
         harmonic bandlimits L this is inefficient as the I/O overhead for communication
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
+
     """
     # Avoid pole-singularities for MWSS sampling
     if sampling.lower() == "mwss":
@@ -809,7 +827,10 @@ def eval_recursion_step(
                 opsdevice = int((L - L_lower) / ndevices)
 
                 flm = flm.at[L_lower:].set(
-                    pmap(eval_recursion_step, in_axes=(0, 1, 2, 2, 1, 1, 1, 1),)(
+                    pmap(
+                        eval_recursion_step,
+                        in_axes=(0, 1, 2, 2, 1, 1, 1, 1),
+                    )(
                         flm[L_lower:].reshape(ndevices, opsdevice, 2 * L - 1),
                         dl_entry.reshape(ntheta, ndevices, opsdevice),
                         dl_iter.reshape(2, ntheta, ndevices, opsdevice),
diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py
index 8af8a677..1eb138ac 100644
--- a/s2fft/transforms/spherical.py
+++ b/s2fft/transforms/spherical.py
@@ -1,19 +1,20 @@
-from jax import jit, custom_vjp
-
-import numpy as np
-import jax.numpy as jnp
 from functools import partial
 from typing import List
+
+import jax.numpy as jnp
+import numpy as np
+from jax import custom_vjp, jit
+
 from s2fft.sampling import s2_samples as samples
+from s2fft.transforms import c_backend_spherical as c_sph
+from s2fft.transforms import otf_recursions as otf
+from s2fft.utils import healpix_ffts as hp
 from s2fft.utils import (
-    resampling,
     quadrature,
-    resampling_jax,
     quadrature_jax,
+    resampling,
+    resampling_jax,
 )
-from s2fft.utils import healpix_ffts as hp
-from s2fft.transforms import otf_recursions as otf
-from s2fft.transforms import c_backend_spherical as c_sph
 
 
 def inverse(
@@ -29,7 +30,8 @@ def inverse(
     L_lower: int = 0,
     _ssht_backend: int = 1,
 ) -> np.ndarray:
-    r"""Wrapper for the inverse spin-spherical harmonic transform.
+    r"""
+    Wrapper for the inverse spin-spherical harmonic transform.
 
     Args:
         flm (np.ndarray): Spherical harmonic coefficients.
@@ -77,8 +79,8 @@ def inverse(
         harmonic bandlimits L this is inefficient as the I/O overhead for communication
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
-    """
 
+    """
     if spin >= 8 and method in ["numpy", "jax"]:
         raise Warning("Recursive transform may provide lower precision beyond spin ~ 8")
 
@@ -113,7 +115,8 @@ def inverse_numpy(
     precomps: List = None,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Compute the inverse spin-spherical harmonic transform (numpy).
+    r"""
+    Compute the inverse spin-spherical harmonic transform (numpy).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -144,6 +147,7 @@ def inverse_numpy(
 
     Returns:
         np.ndarray: Signal on the sphere.
+
     """
     # Define latitudinal sample positions and Fourier offsets
     thetas = samples.thetas(L, sampling, nside)
@@ -209,7 +213,8 @@ def inverse_jax(
     spmd: bool = False,
     L_lower: int = 0,
 ) -> jnp.ndarray:
-    r"""Compute the inverse spin-spherical harmonic transform (JAX).
+    r"""
+    Compute the inverse spin-spherical harmonic transform (JAX).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -251,6 +256,7 @@ def inverse_jax(
         harmonic bandlimits L this is inefficient as the I/O overhead for communication
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
+
     """
     # Define latitudinal sample positions and Fourier offsets
     thetas = samples.thetas(L, sampling, nside)
@@ -337,7 +343,8 @@ def forward(
     iter: int = 3,
     _ssht_backend: int = 1,
 ) -> np.ndarray:
-    r"""Wrapper for the forward spin-spherical harmonic transform.
+    r"""
+    Wrapper for the forward spin-spherical harmonic transform.
 
     Args:
         f (np.ndarray): Signal on the sphere.
@@ -389,8 +396,8 @@ def forward(
         harmonic bandlimits L this is inefficient as the I/O overhead for communication
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
-    """
 
+    """
     if spin >= 8 and method in ["numpy", "jax"]:
         raise Warning("Recursive transform may provide lower precision beyond spin ~ 8")
 
@@ -448,7 +455,8 @@ def forward_numpy(
     precomps: List = None,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Compute the forward spin-spherical harmonic transform (JAX).
+    r"""
+    Compute the forward spin-spherical harmonic transform (JAX).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -479,6 +487,7 @@ def forward_numpy(
 
     Returns:
         np.ndarray: Spherical harmonic coefficients
+
     """
     # Resample mw onto mwss and double resolution of both
     if sampling.lower() == "mw":
@@ -573,7 +582,8 @@ def forward_jax(
     L_lower: int = 0,
     use_healpix_custom_primitive: bool = False,
 ) -> jnp.ndarray:
-    r"""Compute the forward spin-spherical harmonic transform (JAX).
+    r"""
+    Compute the forward spin-spherical harmonic transform (JAX).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -606,7 +616,7 @@ def forward_jax(
         L_lower (int, optional): Harmonic lower-bound. Transform will only be computed
             for :math:`\texttt{L_lower} \leq \ell < \texttt{L}`. Defaults to 0.
 
-        use_healpix_custom_primitive (bool, optional): Whether to use a custom CUDA 
+        use_healpix_custom_primitive (bool, optional): Whether to use a custom CUDA
             primitive for computing HEALPix fast fourier transform when `sampling =
             "healpix"` and running on a cuda compatible gpu device. using a custom
             primitive reduces long compilation times when jit compiling. defaults to
@@ -621,6 +631,7 @@ def forward_jax(
         harmonic bandlimits L this is inefficient as the I/O overhead for communication
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
+
     """
     # Resample mw onto mwss and double resolution of both
     if sampling.lower() == "mw":
diff --git a/s2fft/transforms/wigner.py b/s2fft/transforms/wigner.py
index 5b140780..92090cf8 100644
--- a/s2fft/transforms/wigner.py
+++ b/s2fft/transforms/wigner.py
@@ -1,9 +1,10 @@
-from jax import jit, vmap
-
-import numpy as np
-import jax.numpy as jnp
 from functools import partial
 from typing import List
+
+import jax.numpy as jnp
+import numpy as np
+from jax import jit, vmap
+
 import s2fft
 from s2fft.sampling import so3_samples as samples
 from s2fft.transforms import c_backend_spherical as c_sph
@@ -21,7 +22,8 @@ def inverse(
     L_lower: int = 0,
     _ssht_backend: int = 1,
 ) -> np.ndarray:
-    r"""Wrapper for the inverse Wigner transform, i.e. inverse Fourier transform on
+    r"""
+    Wrapper for the inverse Wigner transform, i.e. inverse Fourier transform on
     :math:`SO(3)`.
 
     Importantly, the convention adopted for storage of f is :math:`[\gamma, \beta,
@@ -77,8 +79,8 @@ def inverse(
     Note:
         [1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”
             IEEE Transactions on Signal Processing 59 (2011): 5876-5887.
-    """
 
+    """
     if N >= 8 and method in ["numpy", "jax"]:
         raise Warning("Recursive transform may provide lower precision beyond N ~ 8")
 
@@ -106,7 +108,8 @@ def inverse_numpy(
     precomps: List = None,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Compute the inverse Wigner transform (numpy).
+    r"""
+    Compute the inverse Wigner transform (numpy).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -145,6 +148,7 @@ def inverse_numpy(
 
     Returns:
         np.ndarray: Signal on the sphere.
+
     """
     if precomps is None:
         precomps = s2fft.generate_precomputes_wigner(
@@ -191,7 +195,8 @@ def inverse_jax(
     precomps: List = None,
     L_lower: int = 0,
 ) -> jnp.ndarray:
-    r"""Compute the inverse Wigner transform (JAX).
+    r"""
+    Compute the inverse Wigner transform (JAX).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -230,6 +235,7 @@ def inverse_jax(
 
     Returns:
         jnp.ndarray: Signal on the sphere.
+
     """
     if precomps is None:
         precomps = s2fft.generate_precomputes_wigner_jax(
@@ -279,7 +285,8 @@ def inverse_jax_ssht(
     reality: bool = False,
     _ssht_backend: int = 1,
 ) -> jnp.ndarray:
-    r"""Compute the inverse Wigner transform (SSHT JAX).
+    r"""
+    Compute the inverse Wigner transform (SSHT JAX).
 
     SSHT is a C library which implements the spin-spherical harmonic transform outlined
     in McEwen & Wiaux 2011 [1]. We make use of their python bindings for which we
@@ -314,6 +321,7 @@ def inverse_jax_ssht(
     Note:
         [1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”
             IEEE Transactions on Signal Processing 59 (2011): 5876-5887.
+
     """
     flmn, fban = _inverse_norm(flmn, L, N, L_lower, sampling)
     fban = _flmn_to_fban(flmn, fban, L, N, sampling, reality, _ssht_backend)
@@ -332,7 +340,8 @@ def forward(
     L_lower: int = 0,
     _ssht_backend: int = 1,
 ) -> np.ndarray:
-    r"""Wrapper for the forward Wigner transform, i.e. Fourier transform on
+    r"""
+    Wrapper for the forward Wigner transform, i.e. Fourier transform on
     :math:`SO(3)`.
 
     Importantly, the convention adopted for storage of f is :math:`[\gamma, \beta,
@@ -387,8 +396,8 @@ def forward(
     Note:
         [1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”
             IEEE Transactions on Signal Processing 59 (2011): 5876-5887.
-    """
 
+    """
     if N >= 8 and method in ["numpy", "jax"]:
         raise Warning("Recursive transform may provide lower precision beyond N ~ 8")
 
@@ -416,7 +425,8 @@ def forward_numpy(
     precomps: List = None,
     L_lower: int = 0,
 ) -> np.ndarray:
-    r"""Compute the forward Wigner transform (numpy).
+    r"""
+    Compute the forward Wigner transform (numpy).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -456,6 +466,7 @@ def forward_numpy(
 
     Returns:
         np.ndarray: Wigner coefficients `flmn` with shape :math:`[2N-1, L, 2L-1]`.
+
     """
     if precomps is None:
         precomps = s2fft.generate_precomputes_wigner(
@@ -509,7 +520,8 @@ def forward_jax(
     precomps: List = None,
     L_lower: int = 0,
 ) -> jnp.ndarray:
-    r"""Compute the forward Wigner transform (JAX).
+    r"""
+    Compute the forward Wigner transform (JAX).
 
     Uses separation of variables and exploits the Price & McEwen recursion for accelerated
     and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this
@@ -549,6 +561,7 @@ def forward_jax(
 
     Returns:
         jnp.ndarray: Wigner coefficients `flmn` with shape :math:`[2N-1, L, 2L-1]`.
+
     """
     if precomps is None:
         precomps = s2fft.generate_precomputes_wigner_jax(
@@ -617,7 +630,8 @@ def forward_jax_ssht(
     reality: bool = False,
     _ssht_backend: int = 1,
 ) -> jnp.ndarray:
-    r"""Compute the forward Wigner transform (SSHT JAX).
+    r"""
+    Compute the forward Wigner transform (SSHT JAX).
 
     SSHT is a C library which implements the spin-spherical harmonic transform outlined
     in McEwen & Wiaux 2011 [1]. We make use of their python bindings for which we
@@ -654,6 +668,7 @@ def forward_jax_ssht(
     Note:
         [1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”
             IEEE Transactions on Signal Processing 59 (2011): 5876-5887.
+
     """
     flmn, fban = _f_to_fban(f, L, N, reality)
     flmn = _fban_to_flmn(flmn, fban, L, N, sampling, reality, _ssht_backend)
@@ -662,7 +677,7 @@ def forward_jax_ssht(
 
 @partial(jit, static_argnums=(1, 2, 3))
 def _f_to_fban(f: jnp.ndarray, L: int, N: int, reality: bool = False) -> jnp.ndarray:
-    """Private function which maps from f to fban (C backend)"""
+    """Private function which maps from f to fban (C backend)."""
     flmn = jnp.zeros(samples.flmn_shape(L, N), dtype=jnp.complex128)
 
     if reality:
@@ -684,7 +699,7 @@ def _fban_to_flmn(
     reality: bool = False,
     _ssht_backend: int = 1,
 ) -> jnp.ndarray:
-    """Private function which maps from fban to flmn (C backend)"""
+    """Private function which maps from fban to flmn (C backend)."""
     ssht_sampling = ["mw", "mwss", "dh", "gl"].index(sampling.lower())
     n_start_ind = 0 if reality else -N + 1
     func = partial(
@@ -705,7 +720,7 @@ def _fban_to_flmn(
 def _reality_and_norm(
     flmn: jnp.ndarray, L: int, N: int, L_lower: int = 0, reality: bool = False
 ) -> jnp.ndarray:
-    """Private function which maps from f to fban (C backend)"""
+    """Private function which maps from f to fban (C backend)."""
     if reality:
         nidx = jnp.arange(1, N)
         sgn = (-1) ** abs(jnp.arange(-L + 1, L))
@@ -739,7 +754,7 @@ def _reality_and_norm(
 def _inverse_norm(
     flmn: jnp.ndarray, L: int, N: int, L_lower: int = 0, sampling: str = "mw"
 ):
-    """Private function which normalised flmn for inverse Wigner (C backend)"""
+    """Private function which normalised flmn for inverse Wigner (C backend)."""
     fban = jnp.zeros(samples.f_shape(L, N, sampling), dtype=jnp.complex128)
 
     flmn = flmn.at[:, L_lower:].set(
@@ -762,7 +777,7 @@ def _flmn_to_fban(
     reality: bool = False,
     _ssht_backend: int = 1,
 ) -> jnp.ndarray:
-    """Private function which maps from flmn to fban (C backend)"""
+    """Private function which maps from flmn to fban (C backend)."""
     ssht_sampling = ["mw", "mwss", "dh", "gl"].index(sampling.lower())
     n_start_ind = 0 if reality else -N + 1
     func = partial(
@@ -781,7 +796,7 @@ def _flmn_to_fban(
 
 @partial(jit, static_argnums=(1, 2, 3))
 def _fban_to_f(fban: jnp.ndarray, L: int, N: int, reality: bool = False) -> jnp.ndarray:
-    """Private function which maps from fban to f (C backend)"""
+    """Private function which maps from fban to f (C backend)."""
     if reality:
         f = jnp.fft.irfft(fban[N - 1 :], 2 * N - 1, axis=-3, norm="forward")
     else:
diff --git a/s2fft/utils/__init__.py b/s2fft/utils/__init__.py
index 53b97926..e71a7bc7 100644
--- a/s2fft/utils/__init__.py
+++ b/s2fft/utils/__init__.py
@@ -1,10 +1,12 @@
-from . import quadrature
-from . import quadrature_jax
-from . import quadrature_torch
-from . import resampling
-from . import resampling_jax
-from . import resampling_torch
-from . import healpix_ffts
-from . import signal_generator
-from . import rotation
-from . import jax_primitive
+from . import (
+    healpix_ffts,
+    jax_primitive,
+    quadrature,
+    quadrature_jax,
+    quadrature_torch,
+    resampling,
+    resampling_jax,
+    resampling_torch,
+    rotation,
+    signal_generator,
+)
diff --git a/s2fft/utils/healpix_ffts.py b/s2fft/utils/healpix_ffts.py
index d55b99ef..075a35ce 100644
--- a/s2fft/utils/healpix_ffts.py
+++ b/s2fft/utils/healpix_ffts.py
@@ -1,23 +1,24 @@
 from functools import partial
 
-import numpy as np
 import jax.numpy as jnp
 import jaxlib.mlir.ir as ir
+import numpy as np
 import torch
 from jax import jit, vmap
-from jax.lib import xla_client
-from jaxlib.hlo_helpers import custom_call
 
 # did not find promote_dtypes_complex outside _src
 from jax._src.numpy.util import promote_dtypes_complex
+from jax.lib import xla_client
+from jaxlib.hlo_helpers import custom_call
+from s2fft_lib import _s2fft
 
-from s2fft.utils.jax_primitive import register_primitive
 from s2fft.sampling import s2_samples as samples
-from s2fft_lib import _s2fft
+from s2fft.utils.jax_primitive import register_primitive
 
 
 def spectral_folding(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
-    """Folds higher frequency Fourier coefficients back onto lower frequency
+    """
+    Folds higher frequency Fourier coefficients back onto lower frequency
     coefficients, i.e. aliasing high frequencies.
 
     Args:
@@ -29,6 +30,7 @@ def spectral_folding(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
 
     Returns:
         np.ndarray: Lower resolution set of aliased Fourier coefficients.
+
     """
     assert nphi <= 2 * L
 
@@ -49,7 +51,8 @@ def spectral_folding(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
 
 
 def spectral_folding_jax(fm: jnp.ndarray, nphi: int, L: int) -> jnp.ndarray:
-    """Folds higher frequency Fourier coefficients back onto lower frequency
+    """
+    Folds higher frequency Fourier coefficients back onto lower frequency
     coefficients, i.e. aliasing high frequencies. JAX specific implementation of
     :func:`~spectral_folding`.
 
@@ -62,6 +65,7 @@ def spectral_folding_jax(fm: jnp.ndarray, nphi: int, L: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Lower resolution set of aliased Fourier coefficients.
+
     """
     slice_start = L - nphi // 2
     slice_stop = slice_start + nphi
@@ -76,7 +80,8 @@ def spectral_folding_jax(fm: jnp.ndarray, nphi: int, L: int) -> jnp.ndarray:
 
 
 def spectral_folding_torch(fm: torch.tensor, nphi: int, L: int) -> torch.tensor:
-    """Folds higher frequency Fourier coefficients back onto lower frequency
+    """
+    Folds higher frequency Fourier coefficients back onto lower frequency
     coefficients, i.e. aliasing high frequencies. Torch specific implementation of
     :func:`~spectral_folding`.
 
@@ -89,6 +94,7 @@ def spectral_folding_torch(fm: torch.tensor, nphi: int, L: int) -> torch.tensor:
 
     Returns:
         torch.tensor: Lower resolution set of aliased Fourier coefficients.
+
     """
     slice_start = L - nphi // 2
     slice_stop = slice_start + nphi
@@ -109,7 +115,8 @@ def spectral_folding_torch(fm: torch.tensor, nphi: int, L: int) -> torch.tensor:
 
 
 def spectral_periodic_extension(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
-    """Extends lower frequency Fourier coefficients onto higher frequency
+    """
+    Extends lower frequency Fourier coefficients onto higher frequency
     coefficients, i.e. imposed periodicity in Fourier space.
 
     Args:
@@ -121,6 +128,7 @@ def spectral_periodic_extension(fm: np.ndarray, nphi: int, L: int) -> np.ndarray
 
     Returns:
         np.ndarray: Higher resolution set of periodic Fourier coefficients.
+
     """
     assert nphi <= 2 * L
 
@@ -142,7 +150,8 @@ def spectral_periodic_extension(fm: np.ndarray, nphi: int, L: int) -> np.ndarray
 
 
 def spectral_periodic_extension_jax(fm: jnp.ndarray, L: int) -> jnp.ndarray:
-    """Extends lower frequency Fourier coefficients onto higher frequency
+    """
+    Extends lower frequency Fourier coefficients onto higher frequency
     coefficients, i.e. imposed periodicity in Fourier space. Based on
     :func:`~spectral_periodic_extension`, modified to be JIT-compilable.
 
@@ -153,6 +162,7 @@ def spectral_periodic_extension_jax(fm: jnp.ndarray, L: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Higher resolution set of periodic Fourier coefficients.
+
     """
     nphi = fm.shape[0]
     return jnp.concatenate(
@@ -165,7 +175,8 @@ def spectral_periodic_extension_jax(fm: jnp.ndarray, L: int) -> jnp.ndarray:
 
 
 def spectral_periodic_extension_torch(fm: torch.tensor, L: int) -> torch.tensor:
-    """Extends lower frequency Fourier coefficients onto higher frequency
+    """
+    Extends lower frequency Fourier coefficients onto higher frequency
     coefficients, i.e. imposed periodicity in Fourier space. Based on
     :func:`~spectral_periodic_extension`.
 
@@ -176,6 +187,7 @@ def spectral_periodic_extension_torch(fm: torch.tensor, L: int) -> torch.tensor:
 
     Returns:
         torch.tensor: Higher resolution set of periodic Fourier coefficients.
+
     """
     nphi = fm.shape[0]
     return torch.concatenate(
@@ -194,7 +206,8 @@ def healpix_fft(
     method: str = "numpy",
     reality: bool = False,
 ) -> np.ndarray:
-    """Wrapper function for the Forward Fast Fourier Transform with spectral
+    """
+    Wrapper function for the Forward Fast Fourier Transform with spectral
     back-projection in the polar regions to manually enforce Fourier periodicity.
 
     Args:
@@ -216,6 +229,7 @@ def healpix_fft(
 
     Returns:
         np.ndarray: Array of Fourier coefficients for all latitudes.
+
     """
     if method.lower() == "numpy":
         return healpix_fft_numpy(f, L, nside, reality)
@@ -230,7 +244,8 @@ def healpix_fft(
 
 
 def healpix_fft_numpy(f: np.ndarray, L: int, nside: int, reality: bool) -> np.ndarray:
-    """Computes the Forward Fast Fourier Transform with spectral back-projection
+    """
+    Computes the Forward Fast Fourier Transform with spectral back-projection
     in the polar regions to manually enforce Fourier periodicity.
 
     Args:
@@ -245,6 +260,7 @@ def healpix_fft_numpy(f: np.ndarray, L: int, nside: int, reality: bool) -> np.nd
 
     Returns:
         np.ndarray: Array of Fourier coefficients for all latitudes.
+
     """
     index = 0
     ftm = np.zeros(samples.ftm_shape(L, "healpix", nside), dtype=np.complex128)
@@ -289,6 +305,7 @@ def healpix_fft_jax(f: jnp.ndarray, L: int, nside: int, reality: bool) -> jnp.nd
 
     Returns:
         jnp.ndarray: Array of Fourier coefficients for all latitudes.
+
     """
 
     def f_chunks_to_ftm_rows(f_chunks, nphi):
@@ -336,7 +353,8 @@ def f_chunks_to_ftm_rows(f_chunks, nphi):
 def healpix_fft_torch(
     f: torch.tensor, L: int, nside: int, reality: bool
 ) -> torch.tensor:
-    """Computes the Forward Fast Fourier Transform with spectral back-projection
+    """
+    Computes the Forward Fast Fourier Transform with spectral back-projection
     in the polar regions to manually enforce Fourier periodicity. Torch specific
     implementation of :func:`~healpix_fft_numpy`.
 
@@ -352,6 +370,7 @@ def healpix_fft_torch(
 
     Returns:
         torch.tensor: Array of Fourier coefficients for all latitudes.
+
     """
     index = 0
     ftm = torch.zeros(samples.ftm_shape(L, "healpix", nside), dtype=torch.complex128)
@@ -383,7 +402,8 @@ def healpix_ifft(
     method: str = "numpy",
     reality: bool = False,
 ) -> np.ndarray:
-    """Wrapper function for the Inverse Fast Fourier Transform with spectral folding
+    """
+    Wrapper function for the Inverse Fast Fourier Transform with spectral folding
     in the polar regions to mitigate aliasing.
 
     Args:
@@ -405,6 +425,7 @@ def healpix_ifft(
 
     Returns:
         np.ndarray: HEALPix pixel-space array.
+
     """
     assert L >= 2 * nside
     if method.lower() == "numpy":
@@ -422,7 +443,8 @@ def healpix_ifft(
 def healpix_ifft_numpy(
     ftm: np.ndarray, L: int, nside: int, reality: bool
 ) -> np.ndarray:
-    """Computes the Inverse Fast Fourier Transform with spectral folding in the polar
+    """
+    Computes the Inverse Fast Fourier Transform with spectral folding in the polar
     regions to mitigate aliasing.
 
     Args:
@@ -437,6 +459,7 @@ def healpix_ifft_numpy(
 
     Returns:
         np.ndarray: HEALPix pixel-space array.
+
     """
     f = np.zeros(samples.f_shape(sampling="healpix", nside=nside), dtype=np.complex128)
     ntheta = ftm.shape[0]
@@ -460,7 +483,8 @@ def healpix_ifft_numpy(
 def healpix_ifft_jax(
     ftm: jnp.ndarray, L: int, nside: int, reality: bool
 ) -> jnp.ndarray:
-    """Computes the Inverse Fast Fourier Transform with spectral folding in the polar
+    """
+    Computes the Inverse Fast Fourier Transform with spectral folding in the polar
     regions to mitigate aliasing, using JAX. JAX specific implementation of
     :func:`~healpix_ifft_numpy`.
 
@@ -476,6 +500,7 @@ def healpix_ifft_jax(
 
     Returns:
         jnp.ndarray: HEALPix pixel-space array.
+
     """
 
     def ftm_rows_to_f_chunks(ftm_rows, nphi):
@@ -509,7 +534,8 @@ def ftm_rows_to_f_chunks(ftm_rows, nphi):
 def healpix_ifft_torch(
     ftm: torch.tensor, L: int, nside: int, reality: bool
 ) -> torch.tensor:
-    """Computes the Inverse Fast Fourier Transform with spectral folding in the polar
+    """
+    Computes the Inverse Fast Fourier Transform with spectral folding in the polar
     regions to mitigate aliasing. Torch specific implementation of
     :func:`~healpix_ifft_numpy`.
 
@@ -525,6 +551,7 @@ def healpix_ifft_torch(
 
     Returns:
         torch.tensor: HEALPix pixel-space array.
+
     """
     f = torch.zeros(
         samples.f_shape(sampling="healpix", nside=nside), dtype=torch.complex128
@@ -547,7 +574,8 @@ def healpix_ifft_torch(
 
 
 def p2phi_rings(t: np.ndarray, nside: int) -> np.ndarray:
-    r"""Convert index to :math:`\phi` angle for HEALPix for all :math:`\theta` rings.
+    r"""
+    Convert index to :math:`\phi` angle for HEALPix for all :math:`\theta` rings.
     Vectorised implementation of :func:`~samples.p2phi_ring`.
 
     Args:
@@ -557,6 +585,7 @@ def p2phi_rings(t: np.ndarray, nside: int) -> np.ndarray:
 
     Returns:
         np.ndarray: :math:`\phi` offset for each ring index.
+
     """
     shift = 1 / 2
     tt = np.zeros_like(t)
@@ -572,7 +601,8 @@ def p2phi_rings(t: np.ndarray, nside: int) -> np.ndarray:
 
 @partial(jit, static_argnums=(1))
 def p2phi_rings_jax(t: jnp.ndarray, nside: int) -> jnp.ndarray:
-    r"""Convert index to :math:`\phi` angle for HEALPix for all :math:`\theta` rings.
+    r"""
+    Convert index to :math:`\phi` angle for HEALPix for all :math:`\theta` rings.
     JAX implementation of :func:`~p2phi_rings`.
 
     Args:
@@ -582,6 +612,7 @@ def p2phi_rings_jax(t: jnp.ndarray, nside: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: :math:`\phi` offset for each ring index.
+
     """
     shift = 1 / 2
     tt = jnp.zeros_like(t)
@@ -598,7 +629,8 @@ def p2phi_rings_jax(t: jnp.ndarray, nside: int) -> jnp.ndarray:
 def ring_phase_shifts_hp(
     L: int, nside: int, forward: bool = False, reality: bool = False
 ) -> np.ndarray:
-    r"""Generates a phase shift vector for HEALPix for all :math:`\theta` rings.
+    r"""
+    Generates a phase shift vector for HEALPix for all :math:`\theta` rings.
 
     Args:
         L (int, optional): Harmonic band-limit.
@@ -614,6 +646,7 @@ def ring_phase_shifts_hp(
 
     Returns:
         np.ndarray: Vector of phase shifts with shape :math:`[n_{\theta}, 2L-1]`.
+
     """
     t = np.arange(samples.ntheta(L, "healpix", nside))
     phi_offsets = p2phi_rings(t, nside)
@@ -627,7 +660,8 @@ def ring_phase_shifts_hp(
 def ring_phase_shifts_hp_jax(
     L: int, nside: int, forward: bool = False, reality: bool = False
 ) -> jnp.ndarray:
-    r"""Generates a phase shift vector for HEALPix for all :math:`\theta` rings. JAX
+    r"""
+    Generates a phase shift vector for HEALPix for all :math:`\theta` rings. JAX
     implementation of :func:`~ring_phase_shifts_hp`.
 
     Args:
@@ -644,6 +678,7 @@ def ring_phase_shifts_hp_jax(
 
     Returns:
         jnp.ndarray: Vector of phase shifts with shape :math:`[n_{\theta}, 2L-1]`.
+
     """
     t = jnp.arange(samples.ntheta(L, "healpix", nside))
     phi_offsets = p2phi_rings_jax(t, nside)
@@ -734,8 +769,9 @@ def _healpix_fft_cuda_lowering(ctx, f, *, L, nside, reality, fft_type, norm):
 def healpix_fft_cuda(
     f: jnp.ndarray, L: int, nside: int, reality: bool, norm: str = "backward"
 ) -> jnp.ndarray:
-    """Healpix FFT JAX implementation using custom CUDA primitive.
-    
+    """
+    Healpix FFT JAX implementation using custom CUDA primitive.
+
     Computes the Forward Fast Fourier Transform with spectral back-projection
     in the polar regions to manually enforce Fourier periodicity.
 
@@ -751,6 +787,7 @@ def healpix_fft_cuda(
 
     Returns:
         jnp.ndarray: Array of Fourier coefficients for all latitudes.
+
     """
     (f,) = promote_dtypes_complex(f)
     return _healpix_fft_cuda_primitive.bind(
@@ -762,8 +799,9 @@ def healpix_fft_cuda(
 def healpix_ifft_cuda(
     ftm: jnp.ndarray, L: int, nside: int, reality: bool, norm: str = "forward"
 ) -> jnp.ndarray:
-    """Healpix IFFT JAX implementation using custom CUDA primitive.
-    
+    """
+    Healpix IFFT JAX implementation using custom CUDA primitive.
+
     Computes the inverse fast Fourier transform with spectral folding in the polar
     regions to mitigate aliasing.
 
@@ -779,6 +817,7 @@ def healpix_ifft_cuda(
 
     Returns:
         jnp.ndarray: HEALPix pixel-space array.
+
     """
     (ftm,) = promote_dtypes_complex(ftm)
     return _healpix_fft_cuda_primitive.bind(
diff --git a/s2fft/utils/jax_primitive.py b/s2fft/utils/jax_primitive.py
index 76aaa945..6aac7c72 100644
--- a/s2fft/utils/jax_primitive.py
+++ b/s2fft/utils/jax_primitive.py
@@ -1,7 +1,8 @@
 from functools import partial
 from typing import Callable, Dict, Optional, Union
+
 from jax import core
-from jax.interpreters import ad, batching, xla, mlir
+from jax.interpreters import ad, batching, mlir, xla
 
 
 def register_primitive(
@@ -13,7 +14,8 @@ def register_primitive(
     jacobian_vector_product: Optional[Callable] = None,
     transpose: Optional[Callable] = None,
 ):
-    """Register a new custom JAX primitive.
+    """
+    Register a new custom JAX primitive.
 
     Args:
         name: Name for primitive.
@@ -29,6 +31,7 @@ def register_primitive(
 
     Returns:
         Registered custom primtive.
+
     """
     primitive = core.Primitive(name)
     primitive.multiple_results = multiple_results
diff --git a/s2fft/utils/quadrature.py b/s2fft/utils/quadrature.py
index 571effcf..5b2e1578 100644
--- a/s2fft/utils/quadrature.py
+++ b/s2fft/utils/quadrature.py
@@ -1,12 +1,14 @@
 import numpy as np
 import numpy.fft as fft
+
 from s2fft.sampling import s2_samples as samples
 
 
 def quad_weights_transform(
     L: int, sampling: str = "mwss", spin: int = 0, nside: int = 0
 ) -> np.ndarray:
-    r"""Compute quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute quadrature weights for :math:`\theta` and :math:`\phi`
     integration *to use in transform* for various sampling schemes.
 
     Quadrature weights to use in transform for MWSS correspond to quadrature weights
@@ -30,8 +32,8 @@ def quad_weights_transform(
         np.ndarray: Quadrature weights *to use in transform* for sampling scheme for
         each :math:`\theta` (weights are identical as :math:`\phi` varies for given
         :math:`\theta`).
-    """
 
+    """
     if sampling.lower() == "mwss":
         return quad_weights_mwss_theta_only(2 * L, spin=0) * 2 * np.pi / (2 * L)
 
@@ -51,7 +53,8 @@ def quad_weights_transform(
 def quad_weights(
     L: int = None, sampling: str = "mw", spin: int = 0, nside: int = None
 ) -> np.ndarray:
-    r"""Compute quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute quadrature weights for :math:`\theta` and :math:`\phi`
     integration for various sampling schemes.
 
     Args:
@@ -72,8 +75,8 @@ def quad_weights(
     Returns:
         np.ndarray: Quadrature weights for sampling scheme for each :math:`\theta`
         (weights are identical as :math:`\phi` varies for given :math:`\theta`).
-    """
 
+    """
     if sampling.lower() == "mw":
         return quad_weights_mw(L, spin)
 
@@ -94,7 +97,8 @@ def quad_weights(
 
 
 def quad_weights_hp(nside: int) -> np.ndarray:
-    r"""Compute HEALPix quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute HEALPix quadrature weights for :math:`\theta` and :math:`\phi`
     integration.
 
     Note:
@@ -108,8 +112,8 @@ def quad_weights_hp(nside: int) -> np.ndarray:
     Returns:
         np.ndarray: Weights computed for each :math:`\theta` (all weights in array are
         identical).
-    """
 
+    """
     npix = 12 * nside**2
     rings = samples.ntheta(sampling="healpix", nside=nside)
     hp_weights = np.zeros(rings, dtype=np.float64)
@@ -119,7 +123,8 @@ def quad_weights_hp(nside: int) -> np.ndarray:
 
 
 def quad_weights_gl(L: int) -> np.ndarray:
-    r"""Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration.
 
     Args:
         L (int): Harmonic band-limit.
@@ -127,6 +132,7 @@ def quad_weights_gl(L: int) -> np.ndarray:
     Returns:
         np.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     x1, x2 = -1.0, 1.0
     ntheta = samples.ntheta(L, "gl")
@@ -156,7 +162,8 @@ def quad_weights_gl(L: int) -> np.ndarray:
 
 
 def quad_weights_dh(L: int) -> np.ndarray:
-    r"""Compute DH quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute DH quadrature weights for :math:`\theta` and :math:`\phi` integration.
 
     Args:
         L (int): Harmonic band-limit.
@@ -164,15 +171,16 @@ def quad_weights_dh(L: int) -> np.ndarray:
     Returns:
         np.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
-    """
 
+    """
     q = quad_weight_dh_theta_only(samples.thetas(L, sampling="dh"), L)
 
     return q * 2 * np.pi / (2 * L - 1)
 
 
 def quad_weight_dh_theta_only(theta: float, L: int) -> float:
-    r"""Compute DH quadrature weight for :math:`\theta` integration (only), for given
+    r"""
+    Compute DH quadrature weight for :math:`\theta` integration (only), for given
     :math:`\theta`.
 
     Args:
@@ -182,8 +190,8 @@ def quad_weight_dh_theta_only(theta: float, L: int) -> float:
 
     Returns:
         float: Weight computed for each :math:`\theta`.
-    """
 
+    """
     w = 0.0
     for k in range(0, L):
         w += np.sin((2 * k + 1) * theta) / (2 * k + 1)
@@ -194,7 +202,8 @@ def quad_weight_dh_theta_only(theta: float, L: int) -> float:
 
 
 def quad_weights_mw(L: int, spin: int = 0) -> np.ndarray:
-    r"""Compute MW quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute MW quadrature weights for :math:`\theta` and :math:`\phi` integration.
 
     Args:
         L (int): Harmonic band-limit.
@@ -204,13 +213,14 @@ def quad_weights_mw(L: int, spin: int = 0) -> np.ndarray:
     Returns:
         np.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
-    """
 
+    """
     return quad_weights_mw_theta_only(L, spin) * 2 * np.pi / (2 * L - 1)
 
 
 def quad_weights_mwss(L: int, spin: int = 0) -> np.ndarray:
-    r"""Compute MWSS quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute MWSS quadrature weights for :math:`\theta` and :math:`\phi` integration.
 
     Args:
         L (int): Harmonic band-limit.
@@ -220,13 +230,14 @@ def quad_weights_mwss(L: int, spin: int = 0) -> np.ndarray:
     Returns:
         np.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
-    """
 
+    """
     return quad_weights_mwss_theta_only(L, spin) * 2 * np.pi / (2 * L)
 
 
 def quad_weights_mwss_theta_only(L: int, spin: int = 0) -> np.ndarray:
-    r"""Compute MWSS quadrature weights for :math:`\theta` integration (only).
+    r"""
+    Compute MWSS quadrature weights for :math:`\theta` integration (only).
 
     Args:
         L (int): Harmonic band-limit.
@@ -235,8 +246,8 @@ def quad_weights_mwss_theta_only(L: int, spin: int = 0) -> np.ndarray:
 
     Returns:
         np.ndarray: Weights computed for each :math:`\theta`.
-    """
 
+    """
     w = np.zeros(2 * L, dtype=np.complex128)
     # Extra negative m, so logically -el-1 <= m <= el.
     for i in range(-(L - 1) + 1, L + 1):
@@ -252,7 +263,8 @@ def quad_weights_mwss_theta_only(L: int, spin: int = 0) -> np.ndarray:
 
 
 def quad_weights_mw_theta_only(L: int, spin: int = 0) -> np.ndarray:
-    r"""Compute MW quadrature weights for :math:`\theta` integration (only).
+    r"""
+    Compute MW quadrature weights for :math:`\theta` integration (only).
 
     Args:
         L (int): Harmonic band-limit.
@@ -261,8 +273,8 @@ def quad_weights_mw_theta_only(L: int, spin: int = 0) -> np.ndarray:
 
     Returns:
         np.ndarray: Weights computed for each :math:`\theta`.
-    """
 
+    """
     w = np.zeros(2 * L - 1, dtype=np.complex128)
     for i in range(-(L - 1), L):
         w[i + L - 1] = mw_weights(i)
@@ -277,7 +289,8 @@ def quad_weights_mw_theta_only(L: int, spin: int = 0) -> np.ndarray:
 
 
 def mw_weights(m: int) -> float:
-    r"""Compute MW weights given as a function of index m.
+    r"""
+    Compute MW weights given as a function of index m.
 
     MW weights are defined by
 
@@ -292,8 +305,8 @@ def mw_weights(m: int) -> float:
 
     Returns:
         float: MW weight.
-    """
 
+    """
     if m == 1:
         return 1j * np.pi / 2
 
diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py
index 5ab7672a..3b18bff5 100644
--- a/s2fft/utils/quadrature_jax.py
+++ b/s2fft/utils/quadrature_jax.py
@@ -1,16 +1,18 @@
+from functools import partial
+
 import jax
+import jax.numpy as jnp
 from jax import jit
 
-from functools import partial
 from s2fft.sampling import s2_samples as samples
-import jax.numpy as jnp
 
 
 @partial(jit, static_argnums=(0, 1, 2))
 def quad_weights_transform(
     L: int, sampling: str = "mwss", nside: int = 0
 ) -> jnp.ndarray:
-    r"""Compute quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute quadrature weights for :math:`\theta` and :math:`\phi`
     integration *to use in transform* for various sampling schemes. JAX implementation of
     :func:`~s2fft.quadrature.quad_weights_transform`.
 
@@ -33,8 +35,8 @@ def quad_weights_transform(
         jnp.ndarray: Quadrature weights *to use in transform* for sampling scheme for
         each :math:`\theta` (weights are identical as :math:`\phi` varies for given
         :math:`\theta`).
-    """
 
+    """
     if sampling.lower() == "mwss":
         return quad_weights_mwss_theta_only(2 * L) * 2 * jnp.pi / (2 * L)
 
@@ -53,7 +55,8 @@ def quad_weights_transform(
 
 @partial(jit, static_argnums=(0, 1, 2))
 def quad_weights(L: int = None, sampling: str = "mw", nside: int = None) -> jnp.ndarray:
-    r"""Compute quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute quadrature weights for :math:`\theta` and :math:`\phi`
     integration for various sampling schemes. JAX implementation of
     :func:`~s2fft.quadrature.quad_weights`.
 
@@ -75,8 +78,8 @@ def quad_weights(L: int = None, sampling: str = "mw", nside: int = None) -> jnp.
     Returns:
         jnp.ndarray: Quadrature weights for sampling scheme for each :math:`\theta`
         (weights are identical as :math:`\phi` varies for given :math:`\theta`).
-    """
 
+    """
     if sampling.lower() == "mw":
         return quad_weights_mw(L)
 
@@ -98,7 +101,8 @@ def quad_weights(L: int = None, sampling: str = "mw", nside: int = None) -> jnp.
 
 @partial(jit, static_argnums=(0))
 def quad_weights_hp(nside: int) -> jnp.ndarray:
-    r"""Compute HEALPix quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute HEALPix quadrature weights for :math:`\theta` and :math:`\phi`
     integration. JAX implementation of :func:`s2fft.quadrature.quad_weights_hp`.
 
     Note:
@@ -112,6 +116,7 @@ def quad_weights_hp(nside: int) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Weights computed for each :math:`\theta` (all weights in array are
         identical).
+
     """
     npix = 12 * nside**2
     rings = samples.ntheta(sampling="healpix", nside=nside)
@@ -120,7 +125,8 @@ def quad_weights_hp(nside: int) -> jnp.ndarray:
 
 @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.
+    r"""
+    Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration.
 
     Args:
         L (int): Harmonic band-limit.
@@ -128,6 +134,7 @@ def quad_weights_gl(L: int) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     x1, x2 = -1.0, 1.0
     ntheta = samples.ntheta(L, "gl")
@@ -169,7 +176,8 @@ def body(arg):
 
 @partial(jit, static_argnums=(0))
 def quad_weights_dh(L: int) -> jnp.ndarray:
-    r"""Compute DH quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute DH quadrature weights for :math:`\theta` and :math:`\phi` integration.
     JAX implementation of :func:`s2fft.quadrature.quad_weights_dh`.
 
     Args:
@@ -178,6 +186,7 @@ def quad_weights_dh(L: int) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     q = quad_weight_dh_theta_only(samples.thetas(L, sampling="dh"), L)
 
@@ -186,7 +195,8 @@ def quad_weights_dh(L: int) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def quad_weight_dh_theta_only(theta: float, L: int) -> float:
-    r"""Compute DH quadrature weight for :math:`\theta` integration (only), for given
+    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`.
 
     Args:
@@ -196,6 +206,7 @@ def quad_weight_dh_theta_only(theta: float, L: int) -> float:
 
     Returns:
         float: Weight computed for each :math:`\theta`.
+
     """
     w = 0.0
     for k in range(0, L):
@@ -208,7 +219,8 @@ def quad_weight_dh_theta_only(theta: float, L: int) -> float:
 
 @partial(jit, static_argnums=(0))
 def quad_weights_mw(L: int) -> jnp.ndarray:
-    r"""Compute MW quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute MW quadrature weights for :math:`\theta` and :math:`\phi` integration.
     JAX implementation of :func:`s2fft.quadrature.quad_weights_mw`.
 
     Args:
@@ -219,13 +231,15 @@ def quad_weights_mw(L: int) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     return quad_weights_mw_theta_only(L) * 2 * jnp.pi / (2 * L - 1)
 
 
 @partial(jit, static_argnums=(0))
 def quad_weights_mwss(L: int) -> jnp.ndarray:
-    r"""Compute MWSS quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute MWSS quadrature weights for :math:`\theta` and :math:`\phi` integration.
     JAX implementation of :func:`s2fft.quadrature.quad_weights_mwss`.
 
     Args:
@@ -236,13 +250,15 @@ def quad_weights_mwss(L: int) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     return quad_weights_mwss_theta_only(L) * 2 * jnp.pi / (2 * L)
 
 
 @partial(jit, static_argnums=(0))
 def quad_weights_mwss_theta_only(L: int) -> jnp.ndarray:
-    r"""Compute MWSS quadrature weights for :math:`\theta` integration (only).
+    r"""
+    Compute MWSS quadrature weights for :math:`\theta` integration (only).
     JAX implementation of :func:`s2fft.quadrature.quad_weights_mwss_theta_only`.
 
     Args:
@@ -252,6 +268,7 @@ def quad_weights_mwss_theta_only(L: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Weights computed for each :math:`\theta`.
+
     """
     w = jnp.zeros(2 * L, dtype=jnp.complex128)
 
@@ -267,7 +284,8 @@ def quad_weights_mwss_theta_only(L: int) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(0))
 def quad_weights_mw_theta_only(L: int) -> jnp.ndarray:
-    r"""Compute MW quadrature weights for :math:`\theta` integration (only).
+    r"""
+    Compute MW quadrature weights for :math:`\theta` integration (only).
     JAX implementation of :func:`s2fft.quadrature.quad_weights_mw_theta_only`.
 
     Args:
@@ -277,6 +295,7 @@ def quad_weights_mw_theta_only(L: int) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Weights computed for each :math:`\theta`.
+
     """
     w = jnp.zeros(2 * L - 1, dtype=jnp.complex128)
     for i in range(-(L - 1), L):
@@ -292,7 +311,8 @@ def quad_weights_mw_theta_only(L: int) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(0))
 def mw_weights(m: int) -> float:
-    r"""Compute MW weights given as a function of index m.
+    r"""
+    Compute MW weights given as a function of index m.
 
     MW weights are defined by
 
@@ -307,6 +327,7 @@ def mw_weights(m: int) -> float:
 
     Returns:
         float: MW weight.
+
     """
     if m == 1:
         return 1j * jnp.pi / 2
diff --git a/s2fft/utils/quadrature_torch.py b/s2fft/utils/quadrature_torch.py
index 8381642c..f669a85b 100644
--- a/s2fft/utils/quadrature_torch.py
+++ b/s2fft/utils/quadrature_torch.py
@@ -1,11 +1,13 @@
 import torch
+
 from s2fft.sampling import s2_samples as samples
 
 
 def quad_weights_transform(
     L: int, sampling: str = "mwss", nside: int = 0
 ) -> torch.tensor:
-    r"""Compute quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute quadrature weights for :math:`\theta` and :math:`\phi`
     integration *to use in transform* for various sampling schemes. Torch implementation of
     :func:`~s2fft.quadrature.quad_weights_transform`.
 
@@ -28,8 +30,8 @@ def quad_weights_transform(
         torch.tensor: Quadrature weights *to use in transform* for sampling scheme for
         each :math:`\theta` (weights are identical as :math:`\phi` varies for given
         :math:`\theta`).
-    """
 
+    """
     if sampling.lower() == "mwss":
         return quad_weights_mwss_theta_only(2 * L) * 2 * torch.pi / (2 * L)
 
@@ -49,7 +51,8 @@ def quad_weights_transform(
 def quad_weights(
     L: int = None, sampling: str = "mw", nside: int = None
 ) -> torch.tensor:
-    r"""Compute quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute quadrature weights for :math:`\theta` and :math:`\phi`
     integration for various sampling schemes. Torch implementation of
     :func:`~s2fft.quadrature.quad_weights`.
 
@@ -71,8 +74,8 @@ def quad_weights(
     Returns:
         torch.tensor: Quadrature weights for sampling scheme for each :math:`\theta`
         (weights are identical as :math:`\phi` varies for given :math:`\theta`).
-    """
 
+    """
     if sampling.lower() == "mw":
         return quad_weights_mw(L)
 
@@ -93,7 +96,8 @@ def quad_weights(
 
 
 def quad_weights_hp(nside: int) -> torch.tensor:
-    r"""Compute HEALPix quadrature weights for :math:`\theta` and :math:`\phi`
+    r"""
+    Compute HEALPix quadrature weights for :math:`\theta` and :math:`\phi`
     integration. Torch implementation of :func:`s2fft.quadrature.quad_weights_hp`.
 
     Note:
@@ -107,6 +111,7 @@ def quad_weights_hp(nside: int) -> torch.tensor:
     Returns:
         torch.tensor: Weights computed for each :math:`\theta` (all weights in array are
         identical).
+
     """
     npix = 12 * nside**2
     rings = samples.ntheta(sampling="healpix", nside=nside)
@@ -114,7 +119,8 @@ def quad_weights_hp(nside: int) -> torch.tensor:
 
 
 def quad_weights_gl(L: int) -> torch.tensor:
-    r"""Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute GL quadrature weights for :math:`\theta` and :math:`\phi` integration.
 
     Args:
         L (int): Harmonic band-limit.
@@ -122,6 +128,7 @@ def quad_weights_gl(L: int) -> torch.tensor:
     Returns:
         np.ndarray: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     x1, x2 = -1.0, 1.0
     ntheta = samples.ntheta(L, "gl")
@@ -151,7 +158,8 @@ def quad_weights_gl(L: int) -> torch.tensor:
 
 
 def quad_weights_dh(L: int) -> torch.tensor:
-    r"""Compute DH quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute DH quadrature weights for :math:`\theta` and :math:`\phi` integration.
     Torch implementation of :func:`s2fft.quadrature.quad_weights_dh`.
 
     Args:
@@ -160,6 +168,7 @@ def quad_weights_dh(L: int) -> torch.tensor:
     Returns:
         torch.tensor: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     q = quad_weight_dh_theta_only(samples.thetas(L, sampling="dh"), L)
 
@@ -167,7 +176,8 @@ def quad_weights_dh(L: int) -> torch.tensor:
 
 
 def quad_weight_dh_theta_only(theta: float, L: int) -> float:
-    r"""Compute DH quadrature weight for :math:`\theta` integration (only), for given
+    r"""
+    Compute DH quadrature weight for :math:`\theta` integration (only), for given
     :math:`\theta`. Torch implementation of :func:`s2fft.quadrature.quad_weights_dh_theta_only`.
 
     Args:
@@ -177,6 +187,7 @@ def quad_weight_dh_theta_only(theta: float, L: int) -> float:
 
     Returns:
         float: Weight computed for each :math:`\theta`.
+
     """
     w = 0.0
     for k in range(0, L):
@@ -188,7 +199,8 @@ def quad_weight_dh_theta_only(theta: float, L: int) -> float:
 
 
 def quad_weights_mw(L: int) -> torch.tensor:
-    r"""Compute MW quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute MW quadrature weights for :math:`\theta` and :math:`\phi` integration.
     Torch implementation of :func:`s2fft.quadrature.quad_weights_mw`.
 
     Args:
@@ -199,12 +211,14 @@ def quad_weights_mw(L: int) -> torch.tensor:
     Returns:
         torch.tensor: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     return quad_weights_mw_theta_only(L) * 2 * torch.pi / (2 * L - 1)
 
 
 def quad_weights_mwss(L: int) -> torch.tensor:
-    r"""Compute MWSS quadrature weights for :math:`\theta` and :math:`\phi` integration.
+    r"""
+    Compute MWSS quadrature weights for :math:`\theta` and :math:`\phi` integration.
     JAX implementation of :func:`s2fft.quadrature.quad_weights_mwss`.
 
     Args:
@@ -215,12 +229,14 @@ def quad_weights_mwss(L: int) -> torch.tensor:
     Returns:
         torch.tensor: Weights computed for each :math:`\theta` (weights are identical
         as :math:`\phi` varies for given :math:`\theta`).
+
     """
     return quad_weights_mwss_theta_only(L) * 2 * torch.pi / (2 * L)
 
 
 def quad_weights_mwss_theta_only(L: int) -> torch.tensor:
-    r"""Compute MWSS quadrature weights for :math:`\theta` integration (only).
+    r"""
+    Compute MWSS quadrature weights for :math:`\theta` integration (only).
     Torch implementation of :func:`s2fft.quadrature.quad_weights_mwss_theta_only`.
 
     Args:
@@ -230,6 +246,7 @@ def quad_weights_mwss_theta_only(L: int) -> torch.tensor:
 
     Returns:
         np.ndarray: Weights computed for each :math:`\theta`.
+
     """
     w = torch.zeros(2 * L, dtype=torch.complex128)
 
@@ -244,7 +261,8 @@ def quad_weights_mwss_theta_only(L: int) -> torch.tensor:
 
 
 def quad_weights_mw_theta_only(L: int) -> torch.tensor:
-    r"""Compute MW quadrature weights for :math:`\theta` integration (only).
+    r"""
+    Compute MW quadrature weights for :math:`\theta` integration (only).
     Torch implementation of :func:`s2fft.quadrature.quad_weights_mw_theta_only`.
 
     Args:
@@ -254,6 +272,7 @@ def quad_weights_mw_theta_only(L: int) -> torch.tensor:
 
     Returns:
         torch.tensor: Weights computed for each :math:`\theta`.
+
     """
     w = torch.zeros(2 * L - 1, dtype=torch.complex128)
     for i in range(-(L - 1), L):
@@ -270,7 +289,8 @@ def quad_weights_mw_theta_only(L: int) -> torch.tensor:
 
 
 def mw_weights(m: int) -> float:
-    r"""Compute MW weights given as a function of index m.
+    r"""
+    Compute MW weights given as a function of index m.
 
     MW weights are defined by
 
@@ -285,6 +305,7 @@ def mw_weights(m: int) -> float:
 
     Returns:
         float: MW weight.
+
     """
     if m == 1:
         return 1j * torch.pi / 2
diff --git a/s2fft/utils/resampling.py b/s2fft/utils/resampling.py
index 606ea3a5..47002f8b 100644
--- a/s2fft/utils/resampling.py
+++ b/s2fft/utils/resampling.py
@@ -1,12 +1,13 @@
 import numpy as np
-import numpy.fft as fft
+
 from s2fft.sampling import s2_samples as samples
 
 
 def periodic_extension(
     f: np.ndarray, L: int, spin: int = 0, sampling: str = "mw"
 ) -> np.ndarray:
-    r"""Perform period extension of MW/MWSS signal on the sphere in harmonic
+    r"""
+    Perform period extension of MW/MWSS signal on the sphere in harmonic
     domain, extending :math:`\theta` domain from :math:`[0,\pi]` to :math:`[0,2\pi]`.
 
     Args:
@@ -25,6 +26,7 @@ def periodic_extension(
     Returns:
         np.ndarray: Signal on the sphere extended to :math:`\theta` domain
         :math:`[0,2\pi]`, in same scheme (MW/MWSS) as input.
+
     """
     ntheta = samples.ntheta(L, sampling)
     nphi = samples.nphi_equiang(L, sampling)
@@ -60,7 +62,8 @@ def periodic_extension(
 
 
 def periodic_extension_spatial_mwss(f: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
-    r"""Perform period extension of MWSS signal on the sphere in spatial domain,
+    r"""
+    Perform period extension of MWSS signal on the sphere in spatial domain,
     extending :math:`\theta` domain from :math:`[0,\pi]` to :math:`[0,2\pi]`.
 
     For the MWSS sampling scheme, it is possible to do the period extension in
@@ -77,6 +80,7 @@ def periodic_extension_spatial_mwss(f: np.ndarray, L: int, spin: int = 0) -> np.
     Returns:
         np.ndarray: Signal on the sphere extended to :math:`\theta` domain
         :math:`[0,2\pi]`, in MWSS sampling scheme.
+
     """
     ntheta = L + 1
     nphi = 2 * L
@@ -98,7 +102,8 @@ def periodic_extension_spatial_mwss(f: np.ndarray, L: int, spin: int = 0) -> np.
 
 
 def upsample_by_two_mwss(f: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
-    r"""Upsample MWSS sampled signal on the sphere defined on domain :math:`[0,\pi]`
+    r"""
+    Upsample MWSS sampled signal on the sphere defined on domain :math:`[0,\pi]`
     by a factor of two.
 
     Upsampling is performed by a periodic extension in :math:`\theta` to
@@ -116,6 +121,7 @@ def upsample_by_two_mwss(f: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
     Returns:
         np.ndarray: Signal on the sphere sampled with MWSS sampling scheme, sampling at
         resolution 2*L.
+
     """
     if f.ndim == 2:
         f = np.expand_dims(f, 0)
@@ -126,7 +132,8 @@ def upsample_by_two_mwss(f: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
 
 
 def upsample_by_two_mwss_ext(f_ext: np.ndarray, L: int) -> np.ndarray:
-    """Upsample an extended MWSS sampled signal on the sphere defined on domain
+    r"""
+    Upsample an extended MWSS sampled signal on the sphere defined on domain
     :math:`[0,2\pi]` by a factor of two.
 
     Upsampling is performed by zero-padding in harmonic space.
@@ -140,6 +147,7 @@ def upsample_by_two_mwss_ext(f_ext: np.ndarray, L: int) -> np.ndarray:
     Returns:
         np.ndarray: Signal on the sphere sampled on extended MWSS sampling scheme on
         domain :math:`[0,2\pi]`, sampling at resolution 2*L.
+
     """
     nphi = 2 * L
     ntheta_ext = 2 * L
@@ -153,7 +161,8 @@ def upsample_by_two_mwss_ext(f_ext: np.ndarray, L: int) -> np.ndarray:
 
 
 def downsample_by_two_mwss(f_ext: np.ndarray, L: int) -> np.ndarray:
-    """Downsample an MWSS sampled signal on the sphere.
+    r"""
+    Downsample an MWSS sampled signal on the sphere.
 
     Can be applied to either MWSS signal sampled on original domain :math:`[0,\pi]`
     or extended domain :math:`[0,2\pi]`.
@@ -173,8 +182,8 @@ def downsample_by_two_mwss(f_ext: np.ndarray, L: int) -> np.ndarray:
     Returns:
         np.ndarray: Signal on the sphere sampled with MWSS sampling scheme at
         resolution L/2.  L must be even so that L/2 is an integer.
-    """
 
+    """
     # Check L is even.
     if L % 2 != 0:
         raise ValueError(f"L must be even (L={L})")
@@ -185,7 +194,8 @@ def downsample_by_two_mwss(f_ext: np.ndarray, L: int) -> np.ndarray:
 
 
 def unextend(f_ext: np.ndarray, L: int, sampling: str = "mw") -> np.ndarray:
-    r"""Unextend MW/MWSS sampled signal from :math:`\theta` domain
+    r"""
+    Unextend MW/MWSS sampled signal from :math:`\theta` domain
     :math:`[0,2\pi]` to :math:`[0,\pi]`.
 
     Args:
@@ -205,8 +215,8 @@ def unextend(f_ext: np.ndarray, L: int, sampling: str = "mw") -> np.ndarray:
     Returns:
         np.ndarray: Signal on the sphere sampled on :math:`\theta` domain
         :math:`[0,\pi]`.
-    """
 
+    """
     if sampling.lower() not in ["mw", "mwss"]:
         raise ValueError(
             "Only mw and mwss supported for periodic extension "
@@ -236,7 +246,8 @@ def unextend(f_ext: np.ndarray, L: int, sampling: str = "mw") -> np.ndarray:
 
 
 def mw_to_mwss_phi(f_mw: np.ndarray, L: int) -> np.ndarray:
-    r"""Convert :math:`\phi` component of signal on the sphere from MW sampling to
+    r"""
+    Convert :math:`\phi` component of signal on the sphere from MW sampling to
     MWSS sampling.
 
     Conversion is performed by zero padding in harmonic space.
@@ -260,6 +271,7 @@ def mw_to_mwss_phi(f_mw: np.ndarray, L: int) -> np.ndarray:
     Returns:
         np.ndarray: Signal on the sphere with MWSS sampling in :math:`\phi` and
         sampling in :math:`\theta` of the input signal.
+
     """
     f_mwss = np.zeros((f_mw.shape[0], L + 1, 2 * L), dtype=np.complex128)
     f_mwss[:, :, 1:] = np.fft.fftshift(
@@ -270,7 +282,8 @@ def mw_to_mwss_phi(f_mw: np.ndarray, L: int) -> np.ndarray:
 
 
 def mw_to_mwss_theta(f_mw: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
-    r"""Convert :math:`\theta` component of signal on the sphere from MW sampling to
+    r"""
+    Convert :math:`\theta` component of signal on the sphere from MW sampling to
     MWSS sampling.
 
     Conversion is performed by first performing a period extension in
@@ -291,6 +304,7 @@ def mw_to_mwss_theta(f_mw: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
     Returns:
         np.ndarray: Signal on the sphere with MWSS sampling in :math:`\theta` and MW
         sampling in :math:`\phi`.
+
     """
     f_mw_ext = periodic_extension(f_mw, L, spin=spin, sampling="mw")
     fmp_mwss_ext = np.zeros((f_mw_ext.shape[0], 2 * L, 2 * L - 1), dtype=np.complex128)
@@ -313,7 +327,8 @@ def mw_to_mwss_theta(f_mw: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
 
 
 def mw_to_mwss(f_mw: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
-    r"""Convert signal on the sphere from MW sampling to MWSS sampling.
+    r"""
+    Convert signal on the sphere from MW sampling to MWSS sampling.
 
     Conversion is performed by first performing a period extension in
     :math:`\theta` to :math:`2\pi`, followed by zero padding in harmonic space.  The
@@ -330,6 +345,7 @@ def mw_to_mwss(f_mw: np.ndarray, L: int, spin: int = 0) -> np.ndarray:
 
     Returns:
         np.ndarray: Signal on the sphere sampled with MWSS sampling.
+
     """
     if f_mw.ndim == 2:
         return np.squeeze(
diff --git a/s2fft/utils/resampling_jax.py b/s2fft/utils/resampling_jax.py
index 976f6d20..cff96906 100644
--- a/s2fft/utils/resampling_jax.py
+++ b/s2fft/utils/resampling_jax.py
@@ -1,13 +1,15 @@
-from jax import jit
+from functools import partial
 
 import jax.numpy as jnp
+from jax import jit
+
 from s2fft.sampling import s2_samples as samples
-from functools import partial
 
 
 @partial(jit, static_argnums=(1))
 def mw_to_mwss(f_mw: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
-    r"""Convert signal on the sphere from MW sampling to MWSS sampling.
+    r"""
+    Convert signal on the sphere from MW sampling to MWSS sampling.
 
     Conversion is performed by first performing a period extension in
     :math:`\theta` to :math:`2\pi`, followed by zero padding in harmonic space.  The
@@ -26,6 +28,7 @@ def mw_to_mwss(f_mw: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
 
     Returns:
         jnp.ndarray: Signal on the sphere sampled with MWSS sampling.
+
     """
     if f_mw.ndim == 2:
         return jnp.squeeze(
@@ -37,7 +40,8 @@ def mw_to_mwss(f_mw: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def mw_to_mwss_theta(f_mw: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
-    r"""Convert :math:`\theta` component of signal on the sphere from MW sampling to
+    r"""
+    Convert :math:`\theta` component of signal on the sphere from MW sampling to
     MWSS sampling.
 
     Conversion is performed by first performing a period extension in
@@ -61,6 +65,7 @@ def mw_to_mwss_theta(f_mw: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Signal on the sphere with MWSS sampling in :math:`\theta` and MW
         sampling in :math:`\phi`.
+
     """
     f_mw_ext = periodic_extension(f_mw, L, spin=spin, sampling="mw")
     fmp_mwss_ext = jnp.zeros(
@@ -93,7 +98,8 @@ def mw_to_mwss_theta(f_mw: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def mw_to_mwss_phi(f_mw: jnp.ndarray, L: int) -> jnp.ndarray:
-    r"""Convert :math:`\phi` component of signal on the sphere from MW sampling to
+    r"""
+    Convert :math:`\phi` component of signal on the sphere from MW sampling to
     MWSS sampling.
 
     Conversion is performed by zero padding in harmonic space.
@@ -120,6 +126,7 @@ def mw_to_mwss_phi(f_mw: jnp.ndarray, L: int) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Signal on the sphere with MWSS sampling in :math:`\phi` and
         sampling in :math:`\theta` of the input signal.
+
     """
     f_mwss = jnp.zeros((f_mw.shape[0], L + 1, 2 * L), dtype=jnp.complex128)
     f_mwss = f_mwss.at[:, :, 1:].set(
@@ -139,7 +146,8 @@ def mw_to_mwss_phi(f_mw: jnp.ndarray, L: int) -> jnp.ndarray:
 def periodic_extension(
     f: jnp.ndarray, L: int, spin: int = 0, sampling: str = "mw"
 ) -> jnp.ndarray:
-    r"""Perform period extension of MW/MWSS signal on the sphere in harmonic
+    r"""
+    Perform period extension of MW/MWSS signal on the sphere in harmonic
     domain, extending :math:`\theta` domain from :math:`[0,\pi]` to :math:`[0,2\pi]`.
     JAX implementation of :func:`~s2fft.resampling.periodic_extension`.
 
@@ -159,6 +167,7 @@ def periodic_extension(
     Returns:
         jnp.ndarray: Signal on the sphere extended to :math:`\theta` domain
         :math:`[0,2\pi]`, in same scheme (MW/MWSS) as input.
+
     """
     ntheta = samples.ntheta(L, sampling)
     nphi = samples.nphi_equiang(L, sampling)
@@ -226,7 +235,8 @@ def periodic_extension(
 
 @partial(jit, static_argnums=(1, 2))
 def unextend(f_ext: jnp.ndarray, L: int, sampling: str = "mw") -> jnp.ndarray:
-    r"""Unextend MW/MWSS sampled signal from :math:`\theta` domain
+    r"""
+    Unextend MW/MWSS sampled signal from :math:`\theta` domain
     :math:`[0,2\pi]` to :math:`[0,\pi]`.
 
     Args:
@@ -246,6 +256,7 @@ def unextend(f_ext: jnp.ndarray, L: int, sampling: str = "mw") -> jnp.ndarray:
     Returns:
         jnp.ndarray: Signal on the sphere sampled on :math:`\theta` domain
         :math:`[0,\pi]`.
+
     """
     if sampling.lower() == "mw":
         return f_ext[:, 0:L, :]
@@ -262,7 +273,8 @@ def unextend(f_ext: jnp.ndarray, L: int, sampling: str = "mw") -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def upsample_by_two_mwss(f: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
-    r"""Upsample MWSS sampled signal on the sphere defined on domain :math:`[0,\pi]`
+    r"""
+    Upsample MWSS sampled signal on the sphere defined on domain :math:`[0,\pi]`
     by a factor of two.
 
     Upsampling is performed by a periodic extension in :math:`\theta` to
@@ -282,6 +294,7 @@ def upsample_by_two_mwss(f: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Signal on the sphere sampled with MWSS sampling scheme, sampling at
         resolution 2*L.
+
     """
     if f.ndim == 2:
         f = jnp.expand_dims(f, 0)
@@ -293,7 +306,8 @@ def upsample_by_two_mwss(f: jnp.ndarray, L: int, spin: int = 0) -> jnp.ndarray:
 
 @partial(jit, static_argnums=(1))
 def upsample_by_two_mwss_ext(f_ext: jnp.ndarray, L: int) -> jnp.ndarray:
-    """Upsample an extended MWSS sampled signal on the sphere defined on domain
+    r"""
+    Upsample an extended MWSS sampled signal on the sphere defined on domain
     :math:`[0,2\pi]` by a factor of two.
 
     Upsampling is performed by zero-padding in harmonic space. JAX implementation of
@@ -308,6 +322,7 @@ def upsample_by_two_mwss_ext(f_ext: jnp.ndarray, L: int) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Signal on the sphere sampled on extended MWSS sampling scheme on
         domain :math:`[0,2\pi]`, sampling at resolution 2*L.
+
     """
     nphi = 2 * L
     ntheta_ext = 2 * L
@@ -332,7 +347,8 @@ def upsample_by_two_mwss_ext(f_ext: jnp.ndarray, L: int) -> jnp.ndarray:
 def periodic_extension_spatial_mwss(
     f: jnp.ndarray, L: int, spin: int = 0
 ) -> jnp.ndarray:
-    r"""Perform period extension of MWSS signal on the sphere in spatial domain,
+    r"""
+    Perform period extension of MWSS signal on the sphere in spatial domain,
     extending :math:`\theta` domain from :math:`[0,\pi]` to :math:`[0,2\pi]`.
 
     For the MWSS sampling scheme, it is possible to do the period extension in
@@ -351,6 +367,7 @@ def periodic_extension_spatial_mwss(
     Returns:
         jnp.ndarray: Signal on the sphere extended to :math:`\theta` domain
         :math:`[0,2\pi]`, in MWSS sampling scheme.
+
     """
     ntheta = L + 1
     nphi = 2 * L
diff --git a/s2fft/utils/resampling_torch.py b/s2fft/utils/resampling_torch.py
index 7f636909..f5a44c97 100644
--- a/s2fft/utils/resampling_torch.py
+++ b/s2fft/utils/resampling_torch.py
@@ -1,9 +1,11 @@
 import torch
+
 from s2fft.sampling import s2_samples as samples
 
 
 def mw_to_mwss(f_mw: torch.tensor, L: int, spin: int = 0) -> torch.tensor:
-    r"""Convert signal on the sphere from MW sampling to MWSS sampling.
+    r"""
+    Convert signal on the sphere from MW sampling to MWSS sampling.
 
     Conversion is performed by first performing a period extension in
     :math:`\theta` to :math:`2\pi`, followed by zero padding in harmonic space.  The
@@ -22,6 +24,7 @@ def mw_to_mwss(f_mw: torch.tensor, L: int, spin: int = 0) -> torch.tensor:
 
     Returns:
         torch.tensor: Signal on the sphere sampled with MWSS sampling.
+
     """
     if f_mw.ndim == 2:
         return torch.squeeze(
@@ -32,7 +35,8 @@ def mw_to_mwss(f_mw: torch.tensor, L: int, spin: int = 0) -> torch.tensor:
 
 
 def mw_to_mwss_theta(f_mw: torch.tensor, L: int, spin: int = 0) -> torch.tensor:
-    r"""Convert :math:`\theta` component of signal on the sphere from MW sampling to
+    r"""
+    Convert :math:`\theta` component of signal on the sphere from MW sampling to
     MWSS sampling.
 
     Conversion is performed by first performing a period extension in
@@ -56,6 +60,7 @@ def mw_to_mwss_theta(f_mw: torch.tensor, L: int, spin: int = 0) -> torch.tensor:
     Returns:
         torch.tensor: Signal on the sphere with MWSS sampling in :math:`\theta` and MW
         sampling in :math:`\phi`.
+
     """
     f_mw_ext = periodic_extension(f_mw, L, spin=spin, sampling="mw")
     fmp_mwss_ext = torch.zeros(
@@ -89,7 +94,8 @@ def mw_to_mwss_theta(f_mw: torch.tensor, L: int, spin: int = 0) -> torch.tensor:
 
 
 def mw_to_mwss_phi(f_mw: torch.tensor, L: int) -> torch.tensor:
-    r"""Convert :math:`\phi` component of signal on the sphere from MW sampling to
+    r"""
+    Convert :math:`\phi` component of signal on the sphere from MW sampling to
     MWSS sampling.
 
     Conversion is performed by zero padding in harmonic space.
@@ -116,6 +122,7 @@ def mw_to_mwss_phi(f_mw: torch.tensor, L: int) -> torch.tensor:
     Returns:
         torch.tensor: Signal on the sphere with MWSS sampling in :math:`\phi` and
         sampling in :math:`\theta` of the input signal.
+
     """
     f_mwss = torch.zeros((f_mw.shape[0], L + 1, 2 * L), dtype=torch.complex128)
     f_mwss[:, :, 1:] = torch.fft.fftshift(
@@ -134,7 +141,8 @@ def mw_to_mwss_phi(f_mw: torch.tensor, L: int) -> torch.tensor:
 def periodic_extension(
     f: torch.tensor, L: int, spin: int = 0, sampling: str = "mw"
 ) -> torch.tensor:
-    r"""Perform period extension of MW/MWSS signal on the sphere in harmonic
+    r"""
+    Perform period extension of MW/MWSS signal on the sphere in harmonic
     domain, extending :math:`\theta` domain from :math:`[0,\pi]` to :math:`[0,2\pi]`.
     Torch implementation of :func:`~s2fft.resampling.periodic_extension`.
 
@@ -154,6 +162,7 @@ def periodic_extension(
     Returns:
         torch.tensor: Signal on the sphere extended to :math:`\theta` domain
         :math:`[0,2\pi]`, in same scheme (MW/MWSS) as input.
+
     """
     ntheta = samples.ntheta(L, sampling)
     nphi = samples.nphi_equiang(L, sampling)
@@ -176,9 +185,11 @@ def periodic_extension(
         ],
         dims=[-2],
     )
-    f_ext[:, L + m_offset : 2 * L - 1 + m_offset, m_offset : 2 * L - 1 + m_offset,] *= (
-        -1
-    ) ** (torch.arange(-(L - 1), L))
+    f_ext[
+        :,
+        L + m_offset : 2 * L - 1 + m_offset,
+        m_offset : 2 * L - 1 + m_offset,
+    ] *= (-1) ** (torch.arange(-(L - 1), L))
     if hasattr(spin, "size"):
         f_ext[
             :,
@@ -213,7 +224,8 @@ def periodic_extension(
 
 
 def unextend(f_ext: torch.tensor, L: int, sampling: str = "mw") -> torch.tensor:
-    r"""Unextend MW/MWSS sampled signal from :math:`\theta` domain
+    r"""
+    Unextend MW/MWSS sampled signal from :math:`\theta` domain
     :math:`[0,2\pi]` to :math:`[0,\pi]`.
 
     Args:
@@ -233,6 +245,7 @@ def unextend(f_ext: torch.tensor, L: int, sampling: str = "mw") -> torch.tensor:
     Returns:
         torch.tensor: Signal on the sphere sampled on :math:`\theta` domain
         :math:`[0,\pi]`.
+
     """
     if sampling.lower() == "mw":
         return f_ext[:, 0:L, :]
@@ -248,7 +261,8 @@ def unextend(f_ext: torch.tensor, L: int, sampling: str = "mw") -> torch.tensor:
 
 
 def upsample_by_two_mwss(f: torch.tensor, L: int, spin: int = 0) -> torch.tensor:
-    r"""Upsample MWSS sampled signal on the sphere defined on domain :math:`[0,\pi]`
+    r"""
+    Upsample MWSS sampled signal on the sphere defined on domain :math:`[0,\pi]`
     by a factor of two.
 
     Upsampling is performed by a periodic extension in :math:`\theta` to
@@ -268,6 +282,7 @@ def upsample_by_two_mwss(f: torch.tensor, L: int, spin: int = 0) -> torch.tensor
     Returns:
         torch.tensor: Signal on the sphere sampled with MWSS sampling scheme, sampling at
         resolution 2*L.
+
     """
     if f.ndim == 2:
         f = torch.unsqueeze(f, 0)
@@ -278,7 +293,8 @@ def upsample_by_two_mwss(f: torch.tensor, L: int, spin: int = 0) -> torch.tensor
 
 
 def upsample_by_two_mwss_ext(f_ext: torch.tensor, L: int) -> torch.tensor:
-    """Upsample an extended MWSS sampled signal on the sphere defined on domain
+    r"""
+    Upsample an extended MWSS sampled signal on the sphere defined on domain
     :math:`[0,2\pi]` by a factor of two.
 
     Upsampling is performed by zero-padding in harmonic space. Torch implementation of
@@ -293,6 +309,7 @@ def upsample_by_two_mwss_ext(f_ext: torch.tensor, L: int) -> torch.tensor:
     Returns:
         torch.tensor: Signal on the sphere sampled on extended MWSS sampling scheme on
         domain :math:`[0,2\pi]`, sampling at resolution 2*L.
+
     """
     nphi = 2 * L
     ntheta_ext = 2 * L
@@ -316,7 +333,8 @@ def upsample_by_two_mwss_ext(f_ext: torch.tensor, L: int) -> torch.tensor:
 def periodic_extension_spatial_mwss(
     f: torch.tensor, L: int, spin: int = 0
 ) -> torch.tensor:
-    r"""Perform period extension of MWSS signal on the sphere in spatial domain,
+    r"""
+    Perform period extension of MWSS signal on the sphere in spatial domain,
     extending :math:`\theta` domain from :math:`[0,\pi]` to :math:`[0,2\pi]`.
 
     For the MWSS sampling scheme, it is possible to do the period extension in
@@ -335,6 +353,7 @@ def periodic_extension_spatial_mwss(
     Returns:
         torch.tensor: Signal on the sphere extended to :math:`\theta` domain
         :math:`[0,2\pi]`, in MWSS sampling scheme.
+
     """
     ntheta = L + 1
     nphi = 2 * L
diff --git a/s2fft/utils/rotation.py b/s2fft/utils/rotation.py
index d61ac10f..e29b4916 100644
--- a/s2fft/utils/rotation.py
+++ b/s2fft/utils/rotation.py
@@ -1,8 +1,9 @@
-import jax.numpy as jnp
-from jax import jit
 from functools import partial
 from typing import Tuple
 
+import jax.numpy as jnp
+from jax import jit
+
 from s2fft.recursions.risbo_jax import compute_full
 
 
@@ -13,7 +14,8 @@ def rotate_flms(
     rotation: Tuple[float, float, float],
     dl_array: jnp.ndarray = None,
 ) -> jnp.ndarray:
-    """Rotates an array of spherical harmonic coefficients by angle rotation.
+    """
+    Rotates an array of spherical harmonic coefficients by angle rotation.
 
     Args:
         flm (jnp.ndarray): Array of spherical harmonic coefficients.
@@ -24,8 +26,8 @@ def rotate_flms(
 
     Returns:
         jnp.ndarray: Rotated spherical harmonic coefficients with shape [L,2L-1].
-    """
 
+    """
     # Split out angles
     alpha = __exp_array(L, rotation[0])
     gamma = __exp_array(L, rotation[2])
@@ -36,7 +38,7 @@ def rotate_flms(
 
     dl = (
         dl_array
-        if dl_array != None
+        if dl_array is not None
         else jnp.zeros((2 * L - 1, 2 * L - 1)).astype(jnp.float64)
     )
 
@@ -68,13 +70,14 @@ def rotate_flms(
 
 @partial(jit, static_argnums=(0, 1))
 def __exp_array(L: int, x: float) -> jnp.ndarray:
-    """Private function to generate rotation arrays for alpha/gamma rotations"""
+    """Private function to generate rotation arrays for alpha/gamma rotations."""
     return jnp.exp(-1j * jnp.arange(-L + 1, L) * x)
 
 
 @partial(jit, static_argnums=(0, 1))
 def generate_rotate_dls(L: int, beta: float) -> jnp.ndarray:
-    """Function which recursively generates the complete plane of reduced
+    """
+    Function which recursively generates the complete plane of reduced
         Wigner d-function coefficients at a given rotation beta.
 
     Args:
@@ -84,6 +87,7 @@ def generate_rotate_dls(L: int, beta: float) -> jnp.ndarray:
     Returns:
         jnp.ndarray: Complete array of [L, 2L-1,2L-1] Wigner d-function coefficients
             for a fixed rotation beta.
+
     """
     dl = jnp.zeros((L, 2 * L - 1, 2 * L - 1)).astype(jnp.float64)
     dl_iter = jnp.zeros((2 * L - 1, 2 * L - 1)).astype(jnp.float64)
diff --git a/s2fft/utils/signal_generator.py b/s2fft/utils/signal_generator.py
index 553fb890..321569c1 100644
--- a/s2fft/utils/signal_generator.py
+++ b/s2fft/utils/signal_generator.py
@@ -1,5 +1,6 @@
 import numpy as np
 import torch
+
 from s2fft.sampling import s2_samples as samples
 from s2fft.sampling import so3_samples as wigner_samples
 
@@ -12,7 +13,8 @@ def generate_flm(
     reality: bool = False,
     using_torch: bool = False,
 ) -> np.ndarray:
-    r"""Generate a 2D set of random harmonic coefficients.
+    r"""
+    Generate a 2D set of random harmonic coefficients.
 
     Note:
         Real signals are explicitly produced from conjugate symmetry.
@@ -60,9 +62,12 @@ def generate_flmn(
     reality: bool = False,
     using_torch: bool = False,
 ) -> np.ndarray:
-    r"""Generate a 3D set of random Wigner coefficients.
+    r"""
+    Generate a 3D set of random Wigner coefficients.
+
     Note:
         Real signals are explicitly produced from conjugate symmetry.
+
     Args:
         rng (Generator): Random number generator.
 
@@ -78,8 +83,8 @@ def generate_flmn(
         using_torch (bool, optional): Desired frontend functionality. Defaults to False.
 
     Returns:
-
         np.ndarray: Random set of Wigner coefficients.
+
     """
     flmn = np.zeros(wigner_samples.flmn_shape(L, N), dtype=np.complex128)
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 2cb60423..dcb126d8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,7 +1,8 @@
 """Collection of shared fixtures"""
+
 from functools import partial
-import pytest
 
+import pytest
 
 DEFAULT_SEED = 8966433580120847635
 
@@ -38,7 +39,6 @@ def rng(seed):
 def flm_generator(rng):
     # Import s2fft (and indirectly numpy) locally to avoid
     # `RuntimeWarning: numpy.ndarray size changed` when importing at module level
-    import s2fft as s2f
     from s2fft.utils import signal_generator
 
     return partial(signal_generator.generate_flm, rng)
diff --git a/tests/test_healpix_ffts.py b/tests/test_healpix_ffts.py
index 6e826000..3ab75d9b 100644
--- a/tests/test_healpix_ffts.py
+++ b/tests/test_healpix_ffts.py
@@ -1,27 +1,29 @@
-import numpy as np
-from numpy.testing import assert_allclose
 import healpy as hp
-import pytest
 import jax
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
 from packaging.version import Version as _Version
 
-if _Version(jax.__version__) < _Version("0.4.32"):
-    from jax.lib.xla_bridge import get_backend
-else:
-    from jax.extend.backend import get_backend
-gpu_available = get_backend().platform == "gpu"
-
-jax.config.update("jax_enable_x64", True)
 from s2fft.sampling import s2_samples as samples
 from s2fft.utils.healpix_ffts import (
+    healpix_fft_cuda,
     healpix_fft_jax,
     healpix_fft_numpy,
-    healpix_fft_cuda,
+    healpix_ifft_cuda,
     healpix_ifft_jax,
     healpix_ifft_numpy,
-    healpix_ifft_cuda,
 )
 
+if _Version(jax.__version__) < _Version("0.4.32"):
+    from jax.lib.xla_bridge import get_backend
+else:
+    from jax.extend.backend import get_backend
+
+gpu_available = get_backend().platform == "gpu"
+
+jax.config.update("jax_enable_x64", True)
+
 nside_to_test = [4, 5]
 reality_to_test = [False, True]
 
diff --git a/tests/test_logs.py b/tests/test_logs.py
index e742d0a3..a44baec9 100644
--- a/tests/test_logs.py
+++ b/tests/test_logs.py
@@ -1,6 +1,7 @@
-import s2fft.logs as lg
 import pytest
 
+import s2fft.logs as lg
+
 
 def test_incorrect_log_yaml_path():
     dir_name = "random/incorrect/filepath/"
diff --git a/tests/test_quadrature.py b/tests/test_quadrature.py
index 93c826e6..a26f2534 100644
--- a/tests/test_quadrature.py
+++ b/tests/test_quadrature.py
@@ -1,11 +1,12 @@
+import numpy as np
+import pytest
 from jax import config
 
-config.update("jax_enable_x64", True)
-import pytest
-import numpy as np
+from s2fft.base_transforms import spherical
 from s2fft.sampling import s2_samples as samples
 from s2fft.utils import quadrature, quadrature_jax, quadrature_torch
-from s2fft.base_transforms import spherical
+
+config.update("jax_enable_x64", True)
 
 
 @pytest.mark.parametrize("L", [5, 6])
@@ -42,8 +43,8 @@ def test_quadrature_mw_weights(flm_generator, L: int, sampling: str, method: str
 def test_quadrature_exceptions():
     L = 10
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         quadrature.quad_weights_transform(L, sampling="foo")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         quadrature.quad_weights(L, sampling="foo")
diff --git a/tests/test_resampling.py b/tests/test_resampling.py
index 87ecaf99..9b03169e 100644
--- a/tests/test_resampling.py
+++ b/tests/test_resampling.py
@@ -1,16 +1,17 @@
-import pytest
 import numpy as np
-from s2fft.utils import resampling
+import pytest
+
 from s2fft.base_transforms import spherical
+from s2fft.utils import resampling
 
 
 def test_periodic_extension_invalid_sampling():
     f_dummy = np.zeros((2, 2), dtype=np.complex128)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         resampling.periodic_extension(f_dummy, L=5, spin=0, sampling="healpix")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         resampling.periodic_extension(f_dummy, L=5, spin=0, sampling="dh")
 
 
@@ -70,21 +71,21 @@ def test_unextend(flm_generator, L: int, sampling: str, spin_reality):
 def test_resampling_exceptions():
     f_dummy = np.zeros((2, 2), dtype=np.complex128)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         L_odd = 3
         resampling.downsample_by_two_mwss(f_dummy, L_odd)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         resampling.unextend(f_dummy, L=5, sampling="healpix")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         resampling.unextend(f_dummy, L=5, sampling="dh")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         # f_ext has wrong shape
         resampling.unextend(f_dummy, L=5, sampling="mw")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         resampling.mw_to_mwss_phi(f_dummy, L=5)
 
 
diff --git a/tests/test_samples.py b/tests/test_samples.py
index da9edf72..a0343cc2 100644
--- a/tests/test_samples.py
+++ b/tests/test_samples.py
@@ -1,11 +1,11 @@
-import pytest
-import numpy as np
+import healpy as hp
 import jax.numpy as jnp
+import numpy as np
 import pyssht as ssht
-import healpy as hp
-from s2fft.sampling import s2_samples as samples
-from s2fft.sampling import reindex
+import pytest
 
+from s2fft.sampling import reindex
+from s2fft.sampling import s2_samples as samples
 
 nside_to_test = [16, 32]
 
@@ -71,7 +71,7 @@ def test_samples_index_conversion(ind: int):
 def test_samples_ncoeff(L: int):
     n = 0
     for el in range(0, L):
-        for m in range(-el, el + 1):
+        for _ in range(-el, el + 1):
             n += 1
 
     assert samples.ncoeff(L) == pytest.approx(n)
@@ -110,38 +110,38 @@ def test_hp_ang2pix(nside: int):
 def test_samples_exceptions():
     L = 10
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.phis_equiang(L, sampling="healpix")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.phis_equiang(L, sampling="foo")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.ntheta(L, sampling="healpix")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.ntheta(sampling="mw")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.ntheta(L, sampling="foo")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.ntheta_extension(L, sampling="foo")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.nphi_equiang(L, sampling="healpix")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.nphi_equiang(L, sampling="foo")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.nphi_ring(-1, nside=2)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.t2theta(t=0, L=L, sampling="healpix")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.t2theta(t=0, L=L, sampling="foo")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.t2theta(t=0, sampling="mw")
diff --git a/tests/test_spherical_base.py b/tests/test_spherical_base.py
index 9c79df75..7499d0b9 100644
--- a/tests/test_spherical_base.py
+++ b/tests/test_spherical_base.py
@@ -1,10 +1,10 @@
-import pytest
+import healpy as hp
 import numpy as np
-from s2fft.base_transforms import spherical
-from s2fft.sampling import s2_samples as samples
 import pyssht as ssht
-import healpy as hp
+import pytest
 
+from s2fft.base_transforms import spherical
+from s2fft.sampling import s2_samples as samples
 
 L_to_test = [6, 7]
 L_lower_to_test = [0, 2]
@@ -140,10 +140,10 @@ def test_healpix_nside_to_L_exceptions(flm_generator, nside: int):
     flm = flm_generator(L=L, reality=True)
     f = spherical._inverse(flm, L, spin, sampling, "direct", nside)
 
-    with pytest.raises(AssertionError) as e:
+    with pytest.raises(AssertionError):
         spherical.forward(f, L, spin, sampling, nside)
 
-    with pytest.raises(AssertionError) as e:
+    with pytest.raises(AssertionError):
         spherical.inverse(flm, L, spin, sampling, nside)
 
 
@@ -155,14 +155,14 @@ def test_L_lower_exception(flm_generator, L: int):
     flm = flm_generator(L=L, reality=False)
     f = spherical.inverse(flm, L, spin, sampling)
 
-    with pytest.raises(AssertionError) as e:
+    with pytest.raises(AssertionError):
         spherical.forward(f, L, spin, sampling, L_lower=-1)
 
-    with pytest.raises(AssertionError) as e:
+    with pytest.raises(AssertionError):
         spherical.forward(f, L, spin, sampling, L_lower=L)
 
-    with pytest.raises(AssertionError) as e:
+    with pytest.raises(AssertionError):
         spherical.inverse(flm, L, spin, sampling, L_lower=-1)
 
-    with pytest.raises(AssertionError) as e:
+    with pytest.raises(AssertionError):
         spherical.inverse(flm, L, spin, sampling, L_lower=L)
diff --git a/tests/test_spherical_custom_grads.py b/tests/test_spherical_custom_grads.py
index fb01e06f..15f0dcd5 100644
--- a/tests/test_spherical_custom_grads.py
+++ b/tests/test_spherical_custom_grads.py
@@ -1,12 +1,12 @@
 import jax
-
-jax.config.update("jax_enable_x64", True)
-import pytest
 import jax.numpy as jnp
+import pytest
 from jax.test_util import check_grads
 
-from s2fft.transforms import spherical
 from s2fft.recursions.price_mcewen import generate_precomputes_jax
+from s2fft.transforms import spherical
+
+jax.config.update("jax_enable_x64", True)
 
 L_to_test = [16]
 L_lower_to_test = [2]
diff --git a/tests/test_spherical_precompute.py b/tests/test_spherical_precompute.py
index de9d5662..453ce1f3 100644
--- a/tests/test_spherical_precompute.py
+++ b/tests/test_spherical_precompute.py
@@ -1,9 +1,10 @@
-import pytest
 import numpy as np
+import pytest
 import torch
-from s2fft.precompute_transforms.spherical import inverse, forward
-from s2fft.precompute_transforms.construct import spin_spherical_kernel
+
 from s2fft.base_transforms import spherical as base
+from s2fft.precompute_transforms.construct import spin_spherical_kernel
+from s2fft.precompute_transforms.spherical import forward, inverse
 
 L_to_test = [6, 7]
 spin_to_test = [-2, 0, 1]
diff --git a/tests/test_spherical_transform.py b/tests/test_spherical_transform.py
index 18e4e179..a4129fba 100644
--- a/tests/test_spherical_transform.py
+++ b/tests/test_spherical_transform.py
@@ -1,14 +1,14 @@
+import healpy as hp
 import jax
-
-jax.config.update("jax_enable_x64", True)
-import pytest
-import pyssht as ssht
 import numpy as np
-import healpy as hp
+import pyssht as ssht
+import pytest
 
+from s2fft.recursions.price_mcewen import generate_precomputes
 from s2fft.sampling import s2_samples as samples
 from s2fft.transforms import spherical
-from s2fft.recursions.price_mcewen import generate_precomputes
+
+jax.config.update("jax_enable_x64", True)
 
 L_to_test = [6, 7]
 L_lower_to_test = [0, 2]
@@ -190,10 +190,10 @@ def test_spin_exceptions(flm_generator):
     flm = flm_generator(L=L, reality=False)
     f = spherical.inverse(flm, L, spin=0, sampling=sampling, method="jax_ssht")
 
-    with pytest.raises(Warning) as e:
+    with pytest.raises(Warning):
         spherical.inverse(flm, L, spin=spin, sampling=sampling, method="jax")
 
-    with pytest.raises(Warning) as e:
+    with pytest.raises(Warning):
         spherical.forward(f, L, spin=spin, sampling=sampling, method="jax")
 
 
@@ -205,10 +205,10 @@ def test_sampling_ssht_backend_exceptions(flm_generator):
     flm = flm_generator(L=L, reality=False)
     f = spherical.inverse(flm, L, 0, nside, sampling, "jax_healpy")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         spherical.inverse(flm, L, 0, nside, sampling, "jax_ssht")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         spherical.forward(f, L, 0, nside, sampling, "jax_ssht")
 
 
@@ -219,16 +219,16 @@ def test_sampling_healpy_backend_exceptions(flm_generator):
     flm = flm_generator(L=L, reality=False)
     f = spherical.inverse(flm, L, 0, None, sampling, "jax_ssht")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         spherical.inverse(flm, L, 0, None, sampling, "jax_healpy")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         spherical.forward(f, L, 0, None, sampling, "jax_healpy")
 
 
 def test_sampling_exceptions(flm_generator):
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         spherical.inverse(None, 0, 0, None, method="incorrect")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         spherical.forward(None, 0, 0, None, method="incorrect")
diff --git a/tests/test_utils.py b/tests/test_utils.py
index e935f916..ced9cda8 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,14 +1,15 @@
 import jax
-
-jax.config.update("jax_enable_x64", True)
-import pytest
-import pyssht as ssht
-import numpy as np
-from s2fft.sampling import s2_samples as samples
-from s2fft.utils.rotation import rotate_flms, generate_rotate_dls
 import jax.numpy as jnp
+import numpy as np
+import pyssht as ssht
+import pytest
 from jax.test_util import check_grads
 
+from s2fft.sampling import s2_samples as samples
+from s2fft.utils.rotation import generate_rotate_dls, rotate_flms
+
+jax.config.update("jax_enable_x64", True)
+
 L_to_test = [6, 8, 10]
 angles_to_test = [np.pi / 2, np.pi / 6]
 
@@ -46,16 +47,16 @@ def test_flm_reindexing_exceptions(flm_generator):
     flm_1d = samples.flm_2d_to_1d(flm_2d, L)
     flm_3d = np.zeros((1, 1, 1))
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.flm_2d_to_1d(flm_1d, L)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.flm_2d_to_1d(flm_3d, L)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.flm_1d_to_2d(flm_2d, L)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         samples.flm_1d_to_2d(flm_3d, L)
 
 
diff --git a/tests/test_wigner_base.py b/tests/test_wigner_base.py
index e1f5b67a..46c01098 100644
--- a/tests/test_wigner_base.py
+++ b/tests/test_wigner_base.py
@@ -1,9 +1,9 @@
+import numpy as np
 import pytest
 import so3
-import numpy as np
-from s2fft.sampling import so3_samples as samples
-from s2fft.base_transforms import wigner
 
+from s2fft.base_transforms import wigner
+from s2fft.sampling import so3_samples as samples
 
 L_to_test = [6, 7]
 N_to_test = [2, 3]
diff --git a/tests/test_wigner_custom_grads.py b/tests/test_wigner_custom_grads.py
index ee217648..334c0532 100644
--- a/tests/test_wigner_custom_grads.py
+++ b/tests/test_wigner_custom_grads.py
@@ -1,12 +1,12 @@
 import jax
-
-jax.config.update("jax_enable_x64", True)
-import pytest
 import jax.numpy as jnp
+import pytest
 from jax.test_util import check_grads
 
-from s2fft.transforms import wigner
 from s2fft.recursions.price_mcewen import generate_precomputes_wigner_jax
+from s2fft.transforms import wigner
+
+jax.config.update("jax_enable_x64", True)
 
 L_to_test = [6]
 N_to_test = [3]
@@ -91,7 +91,6 @@ def test_ssht_c_backend_inverse_wigner_custom_gradients(
     reality: bool,
     _ssht_backend: int,
 ):
-
     if sampling.lower() == "dh" and _ssht_backend == 1:
         pytest.skip("Driscoll Healy ducc0 backend gradient calculation tempremental.")
 
@@ -134,7 +133,6 @@ def test_ssht_c_backend_forward_wigner_custom_gradients(
     reality: bool,
     _ssht_backend: int,
 ):
-
     if sampling.lower() == "dh" and _ssht_backend == 1:
         pytest.skip("Driscoll Healy ducc0 backend gradient calculation tempremental.")
 
diff --git a/tests/test_wigner_precompute.py b/tests/test_wigner_precompute.py
index 43301f2b..17a4faa7 100644
--- a/tests/test_wigner_precompute.py
+++ b/tests/test_wigner_precompute.py
@@ -1,10 +1,10 @@
-import pytest
 import numpy as np
-
+import pytest
 import torch
-from s2fft.precompute_transforms.wigner import inverse, forward
-from s2fft.precompute_transforms.construct import wigner_kernel
+
 from s2fft.base_transforms import wigner as base
+from s2fft.precompute_transforms.construct import wigner_kernel
+from s2fft.precompute_transforms.wigner import forward, inverse
 
 L_to_test = [8, 10]
 N_to_test = [2, 3]
diff --git a/tests/test_wigner_recursions.py b/tests/test_wigner_recursions.py
index 5cf059ce..3c67db7c 100644
--- a/tests/test_wigner_recursions.py
+++ b/tests/test_wigner_recursions.py
@@ -1,12 +1,13 @@
 import jax
-
-jax.config.update("jax_enable_x64", True)
-import pytest
-import numpy as np
 import jax.numpy as jnp
+import numpy as np
+import pyssht as ssht
+import pytest
+
 from s2fft import recursions
 from s2fft.sampling import s2_samples as samples
-import pyssht as ssht
+
+jax.config.update("jax_enable_x64", True)
 
 L_to_test = [6, 7]
 spin_to_test = [-2, 0, 1]
@@ -129,10 +130,10 @@ def test_trapani_interfaces():
             atol=1e-10,
         )
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         recursions.trapani.init(dl_loop, L, implementation="unexpected")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         recursions.trapani.compute_full(dl_jax, L, el, implementation="unexpected")
 
 
@@ -231,7 +232,7 @@ def test_turok_slice_jax_with_ssht(L: int, spin: int, sampling: str):
 
         for el in range(L):
             if el >= np.abs(spin):
-                print("beta {}, el {}, spin {}".format(beta, el, spin))
+                print(f"beta {beta}, el {el}, spin {spin}")
                 dl_turok = recursions.turok_jax.compute_slice(beta, el, L, -spin)
 
                 np.testing.assert_allclose(
@@ -245,11 +246,11 @@ def test_turok_slice_jax_with_ssht(L: int, spin: int, sampling: str):
 def test_turok_exceptions():
     L = 10
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         recursions.turok.compute_full(np.pi / 2, L, L)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         recursions.turok.compute_slice(beta=np.pi / 2, el=L - 1, L=L, mm=L)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         recursions.turok.compute_slice(beta=np.pi / 2, el=L, L=L, mm=0)
diff --git a/tests/test_wigner_samples.py b/tests/test_wigner_samples.py
index 93902051..add8f730 100644
--- a/tests/test_wigner_samples.py
+++ b/tests/test_wigner_samples.py
@@ -1,6 +1,7 @@
+import numpy as np
 import pytest
 import so3
-import numpy as np
+
 from s2fft.sampling import so3_samples as samples
 
 L_to_test = [8, 16]
@@ -59,7 +60,7 @@ def test_elmn2ind(s2fft_to_so3_sampling, L: int, N: int, sampling: str):
 
 @pytest.mark.parametrize("L", L_to_test)
 @pytest.mark.parametrize("N", N_to_test)
-def test_so3_samples(flmn_generator, L: int, N: int):
+def test_so3_samples_2(flmn_generator, L: int, N: int):
     flmn_3D = flmn_generator(L=L, N=N, reality=False)
     flmn_1D = samples.flmn_3d_to_1d(flmn_3D, L, N)
     flmn_3D_test = samples.flmn_1d_to_3d(flmn_1D, L, N)
diff --git a/tests/test_wigner_transform.py b/tests/test_wigner_transform.py
index cb4ea2e1..26f20d71 100644
--- a/tests/test_wigner_transform.py
+++ b/tests/test_wigner_transform.py
@@ -1,16 +1,16 @@
 import jax
-
-jax.config.update("jax_enable_x64", True)
-import pytest
-import numpy as np
 import jax.numpy as jnp
+import numpy as np
+import pytest
 
-from s2fft.transforms import wigner
 from s2fft.base_transforms import wigner as base_wigner
 from s2fft.recursions.price_mcewen import (
     generate_precomputes_wigner,
     generate_precomputes_wigner_jax,
 )
+from s2fft.transforms import wigner
+
+jax.config.update("jax_enable_x64", True)
 
 L_to_test = [6, 7]
 N_to_test = [2]
@@ -93,7 +93,6 @@ def test_forward_wigner_transform(
 def test_ssht_c_backend_inverse_wigner_transform(
     flmn_generator, L: int, N: int, L_lower: int, sampling: str, reality: bool
 ):
-
     flmn = flmn_generator(L=L, N=N, L_lower=L_lower, reality=reality)
     f_check = base_wigner.inverse(flmn, L, N, L_lower, sampling, reality)
 
@@ -111,7 +110,6 @@ def test_ssht_c_backend_inverse_wigner_transform(
 def test_ssht_c_backend_forward_wigner_transform(
     flmn_generator, L: int, N: int, L_lower: int, sampling: str, reality: bool
 ):
-
     flmn = flmn_generator(L=L, N=N, L_lower=L_lower, reality=reality)
     f = base_wigner.inverse(flmn, L, N, L_lower, sampling, reality)
 
@@ -127,10 +125,10 @@ def test_N_exceptions(flmn_generator):
     flmn = flmn_generator(L=L, N=N)
     f = base_wigner.inverse(flmn, L, N)
 
-    with pytest.raises(Warning) as e:
+    with pytest.raises(Warning):
         wigner.inverse(flmn, L, N)
 
-    with pytest.raises(Warning) as e:
+    with pytest.raises(Warning):
         wigner.forward(f, L, N)
 
 
@@ -140,8 +138,8 @@ def test_sampling_ssht_backend_exceptions(flmn_generator):
     flmn = flmn_generator(L=L, N=N)
     f = base_wigner.inverse(flmn, L, N)
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         wigner.inverse(flmn, L, N, sampling="healpix", method="jax_ssht")
 
-    with pytest.raises(ValueError) as e:
+    with pytest.raises(ValueError):
         wigner.forward(f, L, N, sampling="healpix", method="jax_ssht")