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

Issue:
Healpix transforms with FFTs

Add HEALPix support for SOV and Direct transforms
Closes #81 

**This issue**
This adds healpix support for all direct and sov transforms. It also extends support for mw and mwss to forward direct and sov, updating the CI where appropriate throughout.

**Outstanding**
Can't reasonably test spin signals (spin!=0) for HEALPix and functionality for spin signals in healpy is extremely unwieldy.


Discussion / hints:

Fix patch:
diff --git a/.coveragerc b/.coveragerc
index 7b7113eb..26a060c1 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1,4 +1,17 @@
+[paths]
+source =
+    ../s2fft
+    */site-packages/s2fft
+
+[run]
+branch = true
+parallel = true
+source =
+    s2fft
+
 [report]
+show_missing = true
+precision = 2
 omit = 
     *test_*
     *__init__*
\ No newline at end of file
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 5eebf196..fe9cca64 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -24,15 +24,15 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     
-    - name: Install package
+    - name: Install dependencies
       run: |
         sudo apt install pandoc
-        python -m pip install --upgrade pip    
-        bash build_s2fft.sh
-
+        python -m pip install --upgrade pip
+        python -m pip install tox
+  
     - name: Build Documentation
       run: |
-        cd docs && make html
+        tox -e docs
 
     - name: Deploy
       if: github.ref == 'refs/heads/main'
diff --git a/.gitignore b/.gitignore
index c1a73982..7f1f2fc6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,7 +6,7 @@ build/
 _build/
 *~
 .DS_Store
-logs/*.log
+*.log
 .coverage
 .test.dat
 notebooks/.ipynb_checkpoints/
@@ -18,4 +18,5 @@ dist/
 s2fft.egg-info/
 resources/
 .matrices/
-s2fft.code-workspace
\ No newline at end of file
+s2fft.code-workspace
+.tox
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 00000000..ce5bfa9f
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,19 @@
+graft s2fft
+graft tests
+graft requirements
+graft docs
+
+include .pip_readme.rst
+include README.rst
+include pytest.ini
+include LICENCE.txt
+include LICENCE_EXT.txt
+include s2fft/default-logging-config.yaml
+include tox.ini
+
+exclude .coveragerc
+exclude *.log
+recursive-exclude docs/_build *
+recursive-exclude notebooks *
+
+global-exclude *.py[cod] __pycache__ *.so *.dylib .ipynb_checkpoints/** *~
\ No newline at end of file
diff --git a/README.rst b/README.rst
index 7bc903d9..62c9eaf0 100644
--- a/README.rst
+++ b/README.rst
@@ -21,6 +21,73 @@
 This is a very loose framework for people to start python projects from. To get up and running, go through the code carefully replacing ``Project-name`` with your 
 desired project name (check all documents!), don't forget to change the directory s2fft as well! You will also need to update the links in all badges!
 
+Installation
+============
+
+The Python dependencies for the `s2fft` package are listed in the file ``requirements/requirements-core.txt`` and can be installed 
+into the active Python environment using `pip` by running
+
+.. code-block:: bash 
+
+    pip install -r requirements/requirements-core.txt
+    
+from the root directory of the repository.
+    
+To install the `s2fft` package in editable mode in the current Python environment run
+
+.. code-block:: bash
+    
+    pip install -e .
+    
+from the root directory of the repository.
+
+
+Tests
+=====
+
+To run the tests the additional dependencies in the file ``requirements/requirements-test.txt`` will also need to be installed in the active environment.
+
+The tests can then be run directly
+
+.. code-block:: bash
+    
+    pytest tests
+    
+from the root directory of the repository, or if ``tox`` is installed then it can be used to create a temporary virtual environment with the package dependencies and run the tests in that environment using
+
+.. code-block:: bash
+    
+    tox -e py38
+    
+
+Documentation
+=============
+
+To build the documentation the additional dependencies in the file ``requirements/requirements-docs.txt`` will also need to be installed in the active environment. 
+
+Pandoc also needs to be installed to allow building the tutorial notebook documentation - this can be done using ``conda`` by running
+
+.. code-block:: bash
+    
+    conda install -c conda-forge pandoc
+    
+Note that this installs the Pandoc Haskell library and command-line tool rather than [the `pandoc` Python package on PyPI](https://pypi.org/project/pandoc/) which wraps this library.
+
+The HTML documentation can then be built by running
+
+.. code-block:: bash
+    
+    sphinx-build -M html docs docs/_build -Q
+    
+from the root directory of the repository, or if ``tox`` is installed then it can be used to build the HTML documentation by running
+
+.. code-block:: bash
+    
+    tox -e docs
+    
+from the root directory of the repository.
+
+
 Interface
 =========
 
diff --git a/build_s2fft.sh b/build_s2fft.sh
deleted file mode 100644
index 0394bce2..00000000
--- a/build_s2fft.sh
+++ /dev/null
@@ -1,17 +0,0 @@
-# <This script needs to be ran from within <Project-name> root directory>
-
-# Install core and extra requirements
-echo -ne 'Building Dependencies... \r'
-pip install -q -r requirements/requirements-core.txt
-echo -ne 'Building Dependencies... ########           (33%)\r'
-pip install -q -r requirements/requirements-docs.txt
-echo -ne 'Building Dependencies... ###########        (66%)\r'
-pip install -q -r requirements/requirements-tests.txt
-echo -ne 'Building Dependencies... ################   (100%)\r'
-echo -ne '\n'
-
-# Install specific converter for building tutorial documentation
-conda install pandoc=1.19.2.1 -y
-
-# Build the scattering emulator
-pip install -e .
\ No newline at end of file
diff --git a/docs/Makefile b/docs/Makefile
index 3c7ca094..0fe394bc 100644
--- a/docs/Makefile
+++ b/docs/Makefile
@@ -11,9 +11,13 @@ BUILDDIR      = _build
 help:
 	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 
-.PHONY: help Makefile
+.PHONY: help Makefile clean
 
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
 %: Makefile
 	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+	
+
+clean:
+	rm -r _build
diff --git a/requirements/requirements-core.txt b/requirements/requirements-core.txt
index a58b2634..1a2cde35 100644
--- a/requirements/requirements-core.txt
+++ b/requirements/requirements-core.txt
@@ -3,7 +3,5 @@ numpy>=1.20
 colorlog
 pyyaml
 pyssht
-pytest
-# jax[cpu]
 jaxlib==0.3.10
 jax==0.3.13
diff --git a/requirements/requirements-docs.txt b/requirements/requirements-docs.txt
index f9182c95..0aae193c 100644
--- a/requirements/requirements-docs.txt
+++ b/requirements/requirements-docs.txt
@@ -6,4 +6,7 @@ sphinx-tabs==3.2.0
 sphinx-rtd-dark-mode==1.2.4
 sphinxcontrib-bibtex==2.4.1
 sphinx-git==11.0.0
-sphinxcontrib-texfigure==0.1.3
\ No newline at end of file
+sphinxcontrib-texfigure==0.1.3
+# Notebooks
+ipython==7.16.1
+jupyter==1.0.0
\ No newline at end of file
diff --git a/s2fft/__init__.py b/s2fft/__init__.py
index 34b485dd..30274c33 100644
--- a/s2fft/__init__.py
+++ b/s2fft/__init__.py
@@ -3,4 +3,4 @@
 from . import quadrature
 from . import resampling
 from . import transform
-from . import utils
+from . import utils
\ No newline at end of file
diff --git a/logs/logging.yaml b/s2fft/default-logging-config.yaml
similarity index 93%
rename from logs/logging.yaml
rename to s2fft/default-logging-config.yaml
index 68486da7..1c042cc3 100644
--- a/logs/logging.yaml
+++ b/s2fft/default-logging-config.yaml
@@ -29,21 +29,21 @@ handlers:
         class: logging.FileHandler
         level: INFO
         formatter: simple
-        filename: /logs/info.log
+        filename: info.log
         encoding: utf8
 
     debug_file_handler:
         class: logging.FileHandler
         level: DEBUG
         formatter: simple
-        filename: /logs/debug.log
+        filename: debug.log
         encoding: utf8
 
     critical_file_handler:
         class: logging.FileHandler
         level: CRITICAL
         formatter: simple
-        filename: /logs/critical.log
+        filename: critical.log
         encoding: utf8
 
 loggers:
diff --git a/s2fft/experimental/__init__.py b/s2fft/experimental/__init__.py
deleted file mode 100644
index f671acf7..00000000
--- a/s2fft/experimental/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-from . import transforms_healpix
\ No newline at end of file
diff --git a/s2fft/experimental/transforms_healpix.py b/s2fft/experimental/transforms_healpix.py
deleted file mode 100644
index d9f71720..00000000
--- a/s2fft/experimental/transforms_healpix.py
+++ /dev/null
@@ -1,216 +0,0 @@
-from random import sample
-import numpy as np
-import numpy.fft as fft
-import s2fft.samples as samples
-import s2fft.quadrature as quadrature
-import s2fft.resampling as resampling
-import s2fft.wigner as wigner
-
-def inverse_direct_healpix(
-    flm: np.ndarray, L: int, nside: int, spin: int = 0
-) -> np.ndarray:
-
-    # TODO: Check flm shape consistent with L
-
-    f = np.zeros(12 * nside**2, dtype=np.complex128)
-
-    thetas = samples.thetas(L, "healpix", nside=nside)
-    for t, theta in enumerate(thetas):
-
-        for el in range(0, L):
-
-            if el >= np.abs(spin):
-
-                dl = wigner.turok.compute_slice(theta, el, L, -spin)
-
-                elfactor = np.sqrt((2 * el + 1) / (4 * np.pi))
-
-                for m in range(-el, el + 1):
-
-                    for p, phi in enumerate(samples.phis_ring(t, nside)):
-
-                        f[samples.hp_ang2pix(nside, theta, phi)] += (
-                            (-1) ** spin
-                            * elfactor
-                            * np.exp(1j * m * phi)
-                            * dl[m + L - 1]
-                            * flm[el, m + L - 1]
-                        )
-
-    return f
-
-
-def inverse_sov_healpix(
-    flm: np.ndarray, L: int, nside: int, spin: int = 0
-) -> np.ndarray:
-
-    # TODO: Check flm shape consistent with L
-
-    ntheta = samples.ntheta(L, "healpix", nside=nside)
-
-    f = np.zeros(12 * nside**2, dtype=np.complex128)
-
-    thetas = samples.thetas(L, "healpix", nside=nside)
-    ftm = np.zeros((ntheta, 2 * L - 1), dtype=np.complex128)
-    for t, theta in enumerate(thetas):
-
-        for el in range(0, L):
-
-            if el >= np.abs(spin):
-
-                dl = wigner.turok.compute_slice(theta, el, L, -spin)
-
-                elfactor = np.sqrt((2 * el + 1) / (4 * np.pi))
-
-                for m in range(-el, el + 1):
-
-                    ftm[t, m + L - 1] += (
-                        (-1) ** spin
-                        * elfactor
-                        * dl[m + L - 1]
-                        * flm[el, m + L - 1]
-                    )
-
-    for t, theta in enumerate(thetas):
-
-        for p, phi in enumerate(samples.phis_ring(t, nside)):
-
-            for m in range(-(L - 1), L):
-
-                f[samples.hp_ang2pix(nside, theta, phi)] += ftm[t, m + L - 1] * np.exp(
-                    1j * m * phi
-                )
-
-    return f
-
-
-def inverse_sov_fft_healpix(
-    flm: np.ndarray, L: int, nside: int, spin: int = 0
-) -> np.ndarray:
-
-    raise ValueError(f"Healpix sov + fft not yet functional")
-
-    # TODO: Check flm shape consistent with L
-    from scipy.signal import resample
-
-    ntheta = samples.ntheta(L, "healpix", nside)
-
-    f = np.zeros(12 * nside**2, dtype=np.complex128)
-
-    thetas = samples.thetas(L, "healpix", nside)
-
-    dl = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
-
-    ftm = np.zeros((ntheta, 2 * L - 1), dtype=np.complex128)
-
-    for t, theta in enumerate(thetas):
-
-        for el in range(0, L):
-
-            # TODO: only need quarter of dl plane here and elsewhere
-            dl = wigner.risbo.compute_full(dl, theta, L, el)
-
-            if el >= np.abs(spin):
-
-                elfactor = np.sqrt((2 * el + 1) / (4 * np.pi))
-
-                for m in range(-el, el + 1):
-
-                    # See libsharp paper
-                    psi_0_y = samples.p2phi_ring(t, 0, nside)
-
-                    ftm[t, m + L - 1] += (
-                        (-1) ** spin
-                        * elfactor
-                        * dl[m + L - 1, -spin + L - 1]
-                        * flm[el, m + L - 1]
-                    ) * np.exp(1j * m * psi_0_y)
-    index = 0
-    for t, theta in enumerate(thetas):
-        nphi = samples.nphi_ring(t, nside)
-        f_ring = fft.ifft(fft.ifftshift(ftm[t]), norm="forward")
-        f_ring = resample(f_ring, nphi)
-        f[index : index + nphi] = f_ring
-        index += nphi
-
-    return f
-
-
-def forward_direct_healpix(
-    f: np.ndarray, L: int, nside: int, spin: int = 0
-) -> np.ndarray:
-
-    # TODO: Check f shape consistent with L
-
-    flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
-
-    thetas = samples.thetas(L, "healpix", nside)
-
-    weights = quadrature.quad_weights(L, "healpix", nside=nside)
-    for t, theta in enumerate(thetas):
-
-        for el in range(0, L):
-
-            if el >= np.abs(spin):
-
-                dl = wigner.turok.compute_slice(theta, el, L, -spin)
-
-                elfactor = np.sqrt((2 * el + 1) / (4 * np.pi))
-
-                for m in range(-el, el + 1):
-
-                    for p, phi in enumerate(samples.phis_ring(t, nside)):
-
-                        flm[el, m + L - 1] += (
-                            weights[t]
-                            * (-1) ** spin
-                            * elfactor
-                            * np.exp(-1j * m * phi)
-                            * dl[m + L - 1]
-                            * f[samples.hp_ang2pix(nside, theta, phi)]
-                        )
-
-    return flm
-
-
-def forward_sov_healpix(f: np.ndarray, L: int, nside: int, spin: int = 0) -> np.ndarray:
-
-    # TODO: Check f shape consistent with L
-
-    flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
-
-    thetas = samples.thetas(L, "healpix", nside)
-    ntheta = samples.ntheta(L, "healpix", nside)
-    ftm = np.zeros((ntheta, 2 * L - 1), dtype=np.complex128)
-    for t, theta in enumerate(thetas):
-
-        for m in range(-(L - 1), L):
-
-            for p, phi in enumerate(samples.phis_ring(t, nside)):
-
-                ftm[t, m + L - 1] += (
-                    np.exp(-1j * m * phi) * f[samples.hp_ang2pix(nside, theta, phi)]
-                )
-
-    weights = quadrature.quad_weights(L, "healpix", nside=nside)
-    for t, theta in enumerate(thetas):
-
-        for el in range(0, L):
-
-            if el >= np.abs(spin):
-
-                dl = wigner.turok.compute_slice(theta, el, L, -spin)
-
-                elfactor = np.sqrt((2 * el + 1) / (4 * np.pi))
-
-                for m in range(-el, el + 1):
-
-                    flm[el, m + L - 1] += (
-                        weights[t]
-                        * (-1) ** spin
-                        * elfactor
-                        * dl[m + L - 1]
-                        * ftm[t, m + L - 1]
-                    )
-
-    return flm
\ No newline at end of file
diff --git a/s2fft/logs.py b/s2fft/logs.py
index 2f902b1f..82b72b8a 100644
--- a/s2fft/logs.py
+++ b/s2fft/logs.py
@@ -1,12 +1,12 @@
 import os
 import logging.config
 import logging
+from pathlib import Path
 import yaml
 import s2fft
-import colorlog
 
 
-def setup_logging(custom_yaml_path: str = None, default_level: int = logging.DEBUG):
+def setup_logging(custom_yaml_path: str = None):
     """Initialise and configure logging.
 
     Should be called at the beginning of code to initialise and configure the
@@ -17,47 +17,25 @@ def setup_logging(custom_yaml_path: str = None, default_level: int = logging.DEB
         custom_yaml_path (string): Complete pathname of desired yaml logging
             configuration. If empty will provide default logging config.
 
-        default_level (int, optional): Logging level at which to configure.
-
     Raises:
         ValueError: Raised if logging.yaml is not in ./logs/ directory.
 
     """
-    if custom_yaml_path == None:
-        path = os.path.join(
-            os.path.dirname(os.path.dirname(os.path.realpath(s2fft.__file__)))
-            + "/logs/logging.yaml"
-        )
-    if custom_yaml_path != None:
-        path = custom_yaml_path
-    value = os.getenv("LOG_CFG", None)
-    if value:
-        path = value
-    if os.path.exists(path):
-        with open(path, "rt") as f:
-            config = yaml.safe_load(f.read())
-        if custom_yaml_path == None:
-            config["handlers"]["info_file_handler"]["filename"] = os.path.join(
-                os.path.dirname(os.path.dirname(os.path.realpath(s2fft.__file__)))
-                + "/logs/info.log"
-            )
-            config["handlers"]["debug_file_handler"]["filename"] = os.path.join(
-                os.path.dirname(os.path.dirname(os.path.realpath(s2fft.__file__)))
-                + "/logs/debug.log"
-            )
-            config["handlers"]["critical_file_handler"]["filename"] = os.path.join(
-                os.path.dirname(os.path.dirname(os.path.realpath(s2fft.__file__)))
-                + "/logs/critical.log"
-            )
-            config["handlers"]["info_file_handler"]["filename"] = os.path.join(
-                os.path.dirname(os.path.dirname(os.path.realpath(s2fft.__file__)))
-                + "/logs/info.log"
-            )
-        logging.config.dictConfig(config)
+    if "LOG_CFG" in os.environ:
+        path = Path(os.environ["LOG_CFG"])
+    elif custom_yaml_path is None:
+        path = Path(s2fft.__file__).parent / "default-logging-config.yaml"
     else:
-        logging.basicConfig(level=default_level)
-        raise ValueError("Logging config pathway incorrect.")
-        critical_log("Using custom config from {}".format(path))
+        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:
+        config = yaml.safe_load(f.read())
+    if custom_yaml_path is None:
+        config["handlers"]["info_file_handler"]["filename"] = "info.log"
+        config["handlers"]["debug_file_handler"]["filename"] = "debug.log"
+        config["handlers"]["critical_file_handler"]["filename"] = "critical.log"
+    logging.config.dictConfig(config)
 
 
 def debug_log(message: str):
diff --git a/s2fft/quadrature.py b/s2fft/quadrature.py
index 066c2ab5..7809c1a7 100644
--- a/s2fft/quadrature.py
+++ b/s2fft/quadrature.py
@@ -3,7 +3,9 @@
 import s2fft.samples as samples
 
 
-def quad_weights_transform(L: int, sampling: str = "mwss", spin: int = 0) -> np.ndarray:
+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`
     integration *to use in transform* for various sampling schemes.
 
@@ -14,10 +16,13 @@ def quad_weights_transform(L: int, sampling: str = "mwss", spin: int = 0) -> np.
         L (int): Harmonic band-limit.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mwss", "dh"}.  Defaults to "mwss".
+            {"mwss", "dh", "healpix}.  Defaults to "mwss".
 
         spin (int, optional): Harmonic spin. Defaults to 0.
 
+        nside (int, optional): HEALPix Nside resolution parameter.  Only required
+            if sampling="healpix".  Defaults to None.
+
     Raises:
         ValueError: Invalid sampling scheme.
 
@@ -33,6 +38,9 @@ def quad_weights_transform(L: int, sampling: str = "mwss", spin: int = 0) -> np.
     elif sampling.lower() == "dh":
         return quad_weights_dh(L)
 
+    elif sampling.lower() == "healpix":
+        return quad_weights_hp(nside)
+
     else:
         raise ValueError(f"Sampling scheme sampling={sampling} not supported")
 
@@ -97,7 +105,7 @@ def quad_weights_hp(nside: int) -> np.ndarray:
     """
 
     npix = 12 * nside**2
-    rings = samples.ntheta(0, "healpix", nside)
+    rings = samples.ntheta(sampling="healpix", nside=nside)
     hp_weights = np.zeros(rings, dtype=np.float64)
     hp_weights[:] = 4 * np.pi / npix
 
diff --git a/s2fft/resampling.py b/s2fft/resampling.py
index dbe6ce8c..13bba47d 100644
--- a/s2fft/resampling.py
+++ b/s2fft/resampling.py
@@ -376,3 +376,126 @@ def mw_to_mwss(f_mw: np.ndarray, L, spin: int = 0) -> np.ndarray:
     # phi samples there).
 
     return mw_to_mwss_phi(mw_to_mwss_theta(f_mw, L, spin), L)
+
+
+def spectral_folding(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
+    """Folds higher frequency Fourier coefficients back onto lower frequency
+    coefficients, i.e. aliasing high frequencies.
+
+    Args:
+        fm (np.ndarray): Slice of Fourier coefficients corresponding to ring at latitute t.
+
+        nphi (int): Total number of pixel space phi samples for latitude t.
+
+        L (int): Harmonic band-limit.
+
+    Returns:
+        np.ndarray: Lower resolution set of aliased Fourier coefficients.
+    """
+    assert nphi <= 2 * L
+
+    slice_start = L - nphi // 2
+    slice_stop = slice_start + nphi
+    ftm_slice = fm[slice_start:slice_stop]
+
+    idx = 1
+    while slice_start - idx >= 0:
+        ftm_slice[-idx % nphi] += fm[slice_start - idx]
+        idx += 1
+    idx = 0
+    while slice_stop + idx < len(fm):
+        ftm_slice[idx % nphi] += fm[slice_stop + idx]
+        idx += 1
+
+    return ftm_slice
+
+
+def spectral_periodic_extension(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
+    """Extends lower frequency Fourier coefficients onto higher frequency
+    coefficients, i.e. imposed periodicity in Fourier space.
+
+    Args:
+        fm (np.ndarray): Slice of Fourier coefficients corresponding to ring at latitute t.
+
+        nphi (int): Total number of pixel space phi samples for latitude t.
+
+        L (int): Harmonic band-limit.
+
+    Returns:
+        np.ndarray: Higher resolution set of periodic Fourier coefficients.
+    """
+    assert nphi <= 2 * L
+
+    slice_start = L - nphi // 2
+    slice_stop = slice_start + nphi
+    fm_full = np.zeros(2 * L, dtype=np.complex128)
+    fm_full[slice_start:slice_stop] = fm
+
+    idx = 1
+    while slice_start - idx >= 0:
+        fm_full[slice_start - idx] = fm[-idx % nphi]
+        idx += 1
+    idx = 0
+    while slice_stop + idx < len(fm_full):
+        fm_full[slice_stop + idx] = fm[idx % nphi]
+        idx += 1
+
+    return fm_full
+
+
+def healpix_fft(f: np.ndarray, L: int, nside: int) -> np.ndarray:
+    """Computes the Forward Fast Fourier Transform with spectral back-projection
+    in the polar regions to manually enforce Fourier periodicity.
+
+    Args:
+        f (np.ndarray): HEALPix pixel-space array.
+
+        L (int): Harmonic band-limit.
+
+        nside (int): HEALPix Nside resolution parameter.
+
+    Returns:
+        np.ndarray: Array of Fourier coefficients for all latitudes.
+    """
+    assert L >= 2 * nside
+
+    index = 0
+    ftm = np.zeros(samples.ftm_shape(L, "healpix", nside), dtype=np.complex128)
+    ntheta = ftm.shape[0]
+    for t in range(ntheta):
+        nphi = samples.nphi_ring(t, nside)
+        fm_chunk = fft.fftshift(fft.fft(f[index : index + nphi], norm="backward"))
+        ftm[t] = (
+            fm_chunk
+            if nphi == 2 * L
+            else spectral_periodic_extension(fm_chunk, nphi, L)
+        )
+        index += nphi
+    return ftm
+
+
+def healpix_ifft(ftm: np.ndarray, L: int, nside: int) -> np.ndarray:
+    """Computes the Inverse Fast Fourier Transform with spectral folding in the polar
+    regions to mitigate aliasing.
+
+    Args:
+        ftm (np.ndarray): Array of Fourier coefficients for all latitudes.
+
+        L (int): Harmonic band-limit.
+
+        nside (int): HEALPix Nside resolution parameter.
+
+    Returns:
+        np.ndarray: HEALPix pixel-space array.
+    """
+    assert L >= 2 * nside
+
+    f = np.zeros(samples.f_shape(sampling="healpix", nside=nside), dtype=np.complex128)
+    ntheta = ftm.shape[0]
+    index = 0
+    for t in range(ntheta):
+        nphi = samples.nphi_ring(t, nside)
+        fm_chunk = ftm[t] if nphi == 2 * L else spectral_folding(ftm[t], nphi, L)
+        f[index : index + nphi] = fft.ifft(fft.ifftshift(fm_chunk), norm="forward")
+        index += nphi
+    return f
diff --git a/s2fft/samples.py b/s2fft/samples.py
index 22fd2141..1ce3766a 100644
--- a/s2fft/samples.py
+++ b/s2fft/samples.py
@@ -134,6 +134,52 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int:
     return 1
 
 
+def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> int:
+    r"""stuff and things
+
+    Args:
+        L (int): Harmonic band-limit.
+
+        sampling (str, optional): Sampling scheme.  Supported sampling schemes include
+            {"mw", "mwss", "dh"}.  Defaults to "mw".
+
+        nside (int, optional): HEALPix Nside resolution parameter.
+
+    Raises:
+        ValueError: Unknown sampling scheme.
+
+    Returns:
+        int:
+    """
+
+    if sampling.lower() in ["mwss", "healpix"]:
+
+        return ntheta(L, sampling, nside), 2 * L
+
+    elif sampling.lower() in ["mw", "dh"]:
+
+        return ntheta(L, sampling, nside), 2 * L - 1
+
+    else:
+
+        raise ValueError(f"Sampling scheme sampling={sampling} not supported")
+
+    return 1
+
+
+def nphi_equitorial_band(nside: int) -> int:
+    r"""Number of :math:`\phi` samples within the equitorial band for
+    HEALPix sampling scheme.
+
+    Args:
+        nside (int, optional): HEALPix Nside resolution parameter.
+
+    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`
     ring.
@@ -363,7 +409,7 @@ def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray:
 
     elif sampling.lower() == "healpix":
 
-        raise ValueError(f"Sampling scheme sampling={sampling} not implemented")
+        raise ValueError(f"Sampling scheme sampling={sampling} not supported")
 
     else:
 
@@ -398,7 +444,7 @@ def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> tuple:
 
     if sampling.lower() == "healpix":
 
-        return 1, 12 * nside**2
+        return (12 * nside**2,)
 
     else:
 
diff --git a/s2fft/transform.py b/s2fft/transform.py
index 8f62e970..8260d0bd 100644
--- a/s2fft/transform.py
+++ b/s2fft/transform.py
@@ -1,7 +1,6 @@
-from random import sample
 import numpy as np
 import jax.numpy as jnp
-import jax.lax as lax 
+import jax.lax as lax
 from jax import jit
 from functools import partial
 
@@ -14,7 +13,7 @@
 
 
 def inverse_direct(
-    flm: np.ndarray, L: int, spin: int = 0, sampling: str = "mw"
+    flm: np.ndarray, L: int, spin: int = 0, sampling: str = "mw", nside: int = None
 ) -> np.ndarray:
     """Compute inverse spherical harmonic transform by direct method.
 
@@ -29,7 +28,10 @@ def inverse_direct(
         spin (int, optional): Harmonic spin. Defaults to 0.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+
+        nside (int, optional): HEALPix Nside resolution parameter.  Only required
+            if sampling="healpix".  Defaults to None.
 
     Returns:
         np.ndarray: Signal on the sphere.
@@ -38,12 +40,12 @@ def inverse_direct(
     assert flm.shape == samples.flm_shape(L)
     assert 0 <= spin < L
 
-    ntheta = samples.ntheta(L, sampling)
-    nphi = samples.nphi_equiang(L, sampling)
-    f = np.zeros((ntheta, nphi), dtype=np.complex128)
+    if sampling.lower() != "healpix":
 
-    thetas = samples.thetas(L, sampling)
-    phis_equiang = samples.phis_equiang(L, sampling)
+        phis_ring = samples.phis_equiang(L, sampling)
+
+    thetas = samples.thetas(L, sampling, nside)
+    f = np.zeros(samples.f_shape(L, sampling, nside), dtype=np.complex128)
 
     for t, theta in enumerate(thetas):
 
@@ -57,9 +59,18 @@ def inverse_direct(
 
                 for m in range(-el, el + 1):
 
-                    for p, phi in enumerate(phis_equiang):
+                    if sampling.lower() == "healpix":
+                        phis_ring = samples.phis_ring(t, nside)
+
+                    for p, phi in enumerate(phis_ring):
+
+                        if sampling.lower() != "healpix":
+                            entry = (t, p)
+
+                        else:
+                            entry = samples.hp_ang2pix(nside, theta, phi)
 
-                        f[t, p] += (
+                        f[entry] += (
                             (-1) ** spin
                             * elfactor
                             * np.exp(1j * m * phi)
@@ -71,7 +82,7 @@ def inverse_direct(
 
 
 def inverse_sov(
-    flm: np.ndarray, L: int, spin: int = 0, sampling: str = "mw"
+    flm: np.ndarray, L: int, spin: int = 0, sampling: str = "mw", nside: int = None
 ) -> np.ndarray:
     """Compute inverse spherical harmonic transform by separate of variables method
     (without FFTs).
@@ -87,7 +98,10 @@ def inverse_sov(
         spin (int, optional): Harmonic spin. Defaults to 0.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+
+        nside (int, optional): HEALPix Nside resolution parameter.  Only required
+            if sampling="healpix".  Defaults to None.
 
     Returns:
         np.ndarray: Signal on the sphere.
@@ -96,14 +110,15 @@ def inverse_sov(
     assert flm.shape == samples.flm_shape(L)
     assert 0 <= spin < L
 
-    ntheta = samples.ntheta(L, sampling)
-    nphi = samples.nphi_equiang(L, sampling)
-    f = np.zeros((ntheta, nphi), dtype=np.complex128)
+    ntheta = samples.ntheta(L, sampling, nside=nside)
+    thetas = samples.thetas(L, sampling, nside=nside)
 
-    thetas = samples.thetas(L, sampling)
-    phis_equiang = samples.phis_equiang(L, sampling)
+    if sampling.lower() != "healpix":
+        phis_ring = samples.phis_equiang(L, sampling)
 
     ftm = np.zeros((ntheta, 2 * L - 1), dtype=np.complex128)
+    f = np.zeros(samples.f_shape(L, sampling, nside), dtype=np.complex128)
+
     for t, theta in enumerate(thetas):
 
         for el in range(0, L):
@@ -122,17 +137,26 @@ def inverse_sov(
 
     for t, theta in enumerate(thetas):
 
-        for p, phi in enumerate(phis_equiang):
+        if sampling.lower() == "healpix":
+            phis_ring = samples.phis_ring(t, nside)
+
+        for p, phi in enumerate(phis_ring):
 
             for m in range(-(L - 1), L):
 
-                f[t, p] += ftm[t, m + L - 1] * np.exp(1j * m * phi)
+                if sampling.lower() != "healpix":
+                    entry = (t, p)
+
+                else:
+                    entry = samples.hp_ang2pix(nside, theta, phi)
+
+                f[entry] += ftm[t, m + L - 1] * np.exp(1j * m * phi)
 
     return f
 
 
 def inverse_sov_fft(
-    flm: np.ndarray, L: int, spin: int = 0, sampling: str = "mw"
+    flm: np.ndarray, L: int, spin: int = 0, sampling: str = "mw", nside: int = None
 ) -> np.ndarray:
     """Compute inverse spherical harmonic transform by separate of variables method
     with FFTs.
@@ -145,7 +169,10 @@ def inverse_sov_fft(
         spin (int, optional): Harmonic spin. Defaults to 0.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+
+        nside (int, optional): HEALPix Nside resolution parameter.  Only required
+            if sampling="healpix".  Defaults to None.
 
     Returns:
         np.ndarray: Signal on the sphere.
@@ -153,17 +180,17 @@ def inverse_sov_fft(
 
     assert flm.shape == samples.flm_shape(L)
     assert 0 <= spin < L
+    if sampling.lower() == "healpix":
+        assert L >= 2 * nside
 
-    ntheta = samples.ntheta(L, sampling)
-    nphi = samples.nphi_equiang(L, sampling)
-    f = np.zeros((ntheta, nphi), dtype=np.complex128)
+    thetas = samples.thetas(L, sampling, nside)
 
-    thetas = samples.thetas(L, sampling)
-    phis_equiang = samples.phis_equiang(L, sampling)
+    ftm = np.zeros(samples.ftm_shape(L, sampling, nside), dtype=np.complex128)
+    m_offset = 1 if sampling in ["mwss", "healpix"] else 0
 
-    ftm = np.zeros((ntheta, nphi), dtype=np.complex128)
-    m_offset = 1 if sampling == "mwss" else 0
     for t, theta in enumerate(thetas):
+        if sampling.lower() == "healpix":
+            phi_ring_offset = samples.p2phi_ring(t, 0, nside)
 
         for el in range(0, L):
 
@@ -175,11 +202,20 @@ def inverse_sov_fft(
 
                 for m in range(-el, el + 1):
 
+                    phase_shift = (
+                        np.exp(1j * m * phi_ring_offset)
+                        if sampling.lower() == "healpix"
+                        else 1
+                    )
+
                     ftm[t, m + L - 1 + m_offset] += (
                         (-1) ** spin * elfactor * dl[m + L - 1] * flm[el, m + L - 1]
-                    )
+                    ) * phase_shift
 
-    f = fft.ifft(fft.ifftshift(ftm, axes=1), axis=1, norm="forward")
+    if sampling.lower() == "healpix":
+        f = resampling.healpix_ifft(ftm, L, nside)
+    else:
+        f = fft.ifft(fft.ifftshift(ftm, axes=1), axis=1, norm="forward")
 
     return f
 
@@ -235,7 +271,11 @@ def inverse_sov_fft_vectorized(
 
 
 def forward_direct(
-    f: np.ndarray, L: int, spin: int = 0, sampling: str = "mw"
+    f: np.ndarray,
+    L: int,
+    spin: int = 0,
+    sampling: str = "mw",
+    nside: int = None,
 ) -> np.ndarray:
     """Compute forward spherical harmonic transform by direct method.
 
@@ -250,30 +290,37 @@ def forward_direct(
         spin (int, optional): Harmonic spin. Defaults to 0.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
 
-    Raises:
-        ValueError: Only DH sampling supported at present.
+        nside (int, optional): HEALPix Nside resolution parameter.  Only required
+            if sampling="healpix".  Defaults to None.
 
     Returns:
-        np.ndarray: Spherical harmonic coefficients
+        np.ndarray: Spherical harmonic coefficients.
     """
-
-    assert f.shape == samples.f_shape(L, sampling)
+    assert f.shape == samples.f_shape(L, sampling, nside)
     assert 0 <= spin < L
 
-    if sampling.lower() != "dh":
+    flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
 
-        raise ValueError(
-            f"Sampling scheme sampling={sampling} not implement (only DH supported at present)"
-        )
+    if sampling.lower() == "mw":
+        f = resampling.mw_to_mwss(f, L, spin)
 
-    flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
+    if sampling.lower() in ["mw", "mwss"]:
+        sampling = "mwss"
+        f = resampling.upsample_by_two_mwss(f, L, spin)
+        thetas = samples.thetas(2 * L, sampling)
 
-    thetas = samples.thetas(L, sampling)
-    phis_equiang = samples.phis_equiang(L, sampling)
+    else:
+        thetas = samples.thetas(L, sampling, nside)
+
+    # Don't need to include spin in weights (even for spin signals)
+    # since accounted for already in periodic extension and upsampling.
+    weights = quadrature.quad_weights_transform(L, sampling, 0, nside)
+
+    if sampling.lower() != "healpix":
+        phis_ring = samples.phis_equiang(L, sampling)
 
-    weights = quadrature.quad_weights(L, sampling)
     for t, theta in enumerate(thetas):
 
         for el in range(0, L):
@@ -286,7 +333,15 @@ def forward_direct(
 
                 for m in range(-el, el + 1):
 
-                    for p, phi in enumerate(phis_equiang):
+                    if sampling.lower() == "healpix":
+                        phis_ring = samples.phis_ring(t, nside)
+
+                    for p, phi in enumerate(phis_ring):
+
+                        if sampling.lower() != "healpix":
+                            entry = (t, p)
+                        else:
+                            entry = samples.hp_ang2pix(nside, theta, phi)
 
                         flm[el, m + L - 1] += (
                             weights[t]
@@ -294,14 +349,18 @@ def forward_direct(
                             * elfactor
                             * np.exp(-1j * m * phi)
                             * dl[m + L - 1]
-                            * f[t, p]
+                            * f[entry]
                         )
 
     return flm
 
 
 def forward_sov(
-    f: np.ndarray, L: int, spin: int = 0, sampling: str = "mw"
+    f: np.ndarray,
+    L: int,
+    spin: int = 0,
+    sampling: str = "mw",
+    nside: int = None,
 ) -> np.ndarray:
     """Compute forward spherical harmonic transform by separate of variables method
     (without FFTs).
@@ -317,40 +376,56 @@ def forward_sov(
         spin (int, optional): Harmonic spin. Defaults to 0.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+
+        nside (int, optional): HEALPix Nside resolution parameter.  Only required
+            if sampling="healpix".  Defaults to None.
 
-    Raises:
-        ValueError: Only DH sampling supported at present.
 
     Returns:
-        np.ndarray: Spherical harmonic coefficients
+        np.ndarray: Spherical harmonic coefficients.
     """
 
-    assert f.shape == samples.f_shape(L, sampling)
+    assert f.shape == samples.f_shape(L, sampling, nside)
     assert 0 <= spin < L
 
-    if sampling.lower() != "dh":
+    if sampling.lower() == "mw":
+        f = resampling.mw_to_mwss(f, L, spin)
+
+    if sampling.lower() in ["mw", "mwss"]:
+        sampling = "mwss"
+        f = resampling.upsample_by_two_mwss(f, L, spin)
+        thetas = samples.thetas(2 * L, sampling)
 
-        raise ValueError(
-            f"Sampling scheme sampling={sampling} not implement (only DH supported at present)"
-        )
+    else:
+        thetas = samples.thetas(L, sampling, nside)
 
     flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
 
-    thetas = samples.thetas(L, sampling)
-    phis_equiang = samples.phis_equiang(L, sampling)
+    if sampling.lower() != "healpix":
+        phis_ring = samples.phis_equiang(L, sampling)
 
-    ntheta = samples.ntheta(L, sampling)
-    ftm = np.zeros((ntheta, 2 * L - 1), dtype=np.complex128)
+    ftm = np.zeros((len(thetas), 2 * L - 1), dtype=np.complex128)
     for t, theta in enumerate(thetas):
 
         for m in range(-(L - 1), L):
 
-            for p, phi in enumerate(phis_equiang):
+            if sampling.lower() == "healpix":
+                phis_ring = samples.phis_ring(t, nside)
 
-                ftm[t, m + L - 1] += np.exp(-1j * m * phi) * f[t, p]
+            for p, phi in enumerate(phis_ring):
+
+                if sampling.lower() != "healpix":
+                    entry = (t, p)
+                else:
+                    entry = samples.hp_ang2pix(nside, theta, phi)
+
+                ftm[t, m + L - 1] += np.exp(-1j * m * phi) * f[entry]
+
+    # Don't need to include spin in weights (even for spin signals)
+    # since accounted for already in periodic extension and upsampling.
+    weights = quadrature.quad_weights_transform(L, sampling, 0, nside)
 
-    weights = quadrature.quad_weights(L, sampling)
     for t, theta in enumerate(thetas):
 
         for el in range(0, L):
@@ -375,7 +450,7 @@ def forward_sov(
 
 
 def forward_sov_fft(
-    f: np.ndarray, L: int, spin: int = 0, sampling: str = "mw"
+    f: np.ndarray, L: int, spin: int = 0, sampling: str = "mw", nside: int = None
 ) -> np.ndarray:
     """Compute forward spherical harmonic transform by separate of variables method
     with FFTs.
@@ -388,14 +463,16 @@ def forward_sov_fft(
         spin (int, optional): Harmonic spin. Defaults to 0.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
 
     Returns:
         np.ndarray: Spherical harmonic coefficients
     """
 
-    assert f.shape == samples.f_shape(L, sampling)
+    assert f.shape == samples.f_shape(L, sampling, nside)
     assert 0 <= spin < L
+    if sampling.lower() == "healpix":
+        assert L >= 2 * nside
 
     if sampling.lower() == "mw":
         f = resampling.mw_to_mwss(f, L, spin)
@@ -405,17 +482,23 @@ def forward_sov_fft(
         f = resampling.upsample_by_two_mwss(f, L, spin)
         thetas = samples.thetas(2 * L, sampling)
     else:
-        thetas = samples.thetas(L, sampling)
+        thetas = samples.thetas(L, sampling, nside)
 
     flm = np.zeros(samples.flm_shape(L), dtype=np.complex128)
 
-    ftm = fft.fftshift(fft.fft(f, axis=1, norm="backward"), axes=1)
+    if sampling.lower() == "healpix":
+        ftm = resampling.healpix_fft(f, L, nside)
+    else:
+        ftm = fft.fftshift(fft.fft(f, axis=1, norm="backward"), axes=1)
 
     # Don't need to include spin in weights (even for spin signals)
     # since accounted for already in periodic extension and upsampling.
-    weights = quadrature.quad_weights_transform(L, sampling, spin=0)
-    m_offset = 1 if sampling == "mwss" else 0
+    weights = quadrature.quad_weights_transform(L, sampling, 0, nside)
+    m_offset = 1 if sampling in ["mwss", "healpix"] else 0
+
     for t, theta in enumerate(thetas):
+        if sampling.lower() == "healpix":
+            phi_ring_offset = samples.p2phi_ring(t, 0, nside)
 
         for el in range(0, L):
 
@@ -427,16 +510,23 @@ def forward_sov_fft(
 
                 for m in range(-el, el + 1):
 
+                    phase_shift = (
+                        np.exp(-1j * m * phi_ring_offset)
+                        if sampling.lower() == "healpix"
+                        else 1
+                    )
+
                     flm[el, m + L - 1] += (
                         weights[t]
                         * (-1) ** spin
                         * elfactor
                         * dl[m + L - 1]
                         * ftm[t, m + L - 1 + m_offset]
-                    )
+                    ) * phase_shift
 
     return flm
 
+
 def forward_sov_fft_vectorized(
     f: np.ndarray, L: int, spin: int = 0, sampling: str = "mw"
 ) -> np.ndarray:
diff --git a/s2fft/wigner/risbo.py b/s2fft/wigner/risbo.py
index 435d6b97..c2ffd13b 100644
--- a/s2fft/wigner/risbo.py
+++ b/s2fft/wigner/risbo.py
@@ -1,5 +1,4 @@
 import numpy as np
-import logs
 
 
 def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
diff --git a/s2fft/wigner/turok_jax.py b/s2fft/wigner/turok_jax.py
index 56f80306..6098e78e 100644
--- a/s2fft/wigner/turok_jax.py
+++ b/s2fft/wigner/turok_jax.py
@@ -95,44 +95,45 @@ def _compute_quarter_slice(
     omc = 1.0 - c
 
     # Indexing boundaries
-    half_slices = jnp.zeros(2, dtype=jnp.int64)
-    half_slices = half_slices.at[0].set(el + mm + 1)
-    half_slices = half_slices.at[1].set(el - mm + 1)
-
-    lims = jnp.zeros(2, dtype=jnp.int64)
-    lims = lims.at[0].set(0)
-    lims = lims.at[1].set(-1)
+    half_slices = jnp.array([el + mm + 1, el - mm + 1], dtype=jnp.int64)
+    lims = np.array([0, -1], dtype=np.int64)
 
     # Vectors with indexing -L < m < L adopted throughout
-    lrenorm = jnp.zeros(2, dtype=jnp.float64)
-    sign = jnp.zeros(2, dtype=jnp.float64)
-    cpi = jnp.zeros(L + 1, dtype=jnp.float64)
-    cp2 = jnp.zeros(L + 1, dtype=jnp.float64)
-    log_first_row = jnp.zeros(2 * L + 1, dtype=jnp.float64)
 
     # Populate vectors for first row
-    log_first_row = log_first_row.at[0].set(2.0 * el * jnp.log(jnp.abs(c2)))
-
-    for i in range(2, L + 1 + np.abs(mm)):
-        ratio = (2 * el + 2 - i) / (i - 1)
-        log_first_row = log_first_row.at[i - 1].set(
-            log_first_row[i - 2] + jnp.log(ratio) / 2 + lt
+    def log_first_row_iteration(log_first_row_i_minus_1, i):
+        ratio = (2 * el + 2 - i - 1) / i
+        log_first_row_i = log_first_row_i_minus_1 + jnp.log(ratio) / 2 + lt
+        return log_first_row_i, log_first_row_i
+
+    log_first_row_0 = 2 * el * jnp.log(abs(c2))
+    _, log_first_row_1_to_L_plus_abs_mm = lax.scan(
+        log_first_row_iteration, log_first_row_0, np.arange(1, L + abs(mm))
+    )
+    log_first_row = jnp.concatenate(
+        (
+            jnp.atleast_1d(log_first_row_0),
+            log_first_row_1_to_L_plus_abs_mm,
+            jnp.zeros(L + 1 - abs(mm)),
         )
+    )
 
-    sign = sign.at[0].set((t / jnp.abs(t)) ** ((half_slices[0] - 1) % 2))
-    sign = sign.at[1].set((t / jnp.abs(t)) ** ((half_slices[1] - 1) % 2))
+    sign = (t / abs(t)) ** ((half_slices - 1) % 2)
 
     # Initialising coefficients cp(m)= cplus(l-m).
-    cpi = cpi.at[0].set(2.0 / jnp.sqrt(2 * el))
-    for m in range(2, L + 1):
-        cpi = cpi.at[m - 1].set(2.0 / jnp.sqrt(m * (2 * el + 1 - m)))
-        cp2 = cp2.at[m - 1].set(cpi[m - 1] / cpi[m - 2])
+    m = jnp.arange(L)
+    cpi = jnp.concatenate((2 / ((m + 1) * (2 * el - m)) ** 0.5, jnp.zeros(1)))
+    cp2 = jnp.concatenate((jnp.zeros(1), cpi[1:] / cpi[:-1]))
 
     # Use Turok & Bucher recursion to evaluate a single half row
     # Then evaluate the negative half row and reflect using
     # Wigner-d symmetry relation.
 
     em = jnp.arange(1, L + 1)
+    lrenorm = jnp.zeros(2, dtype=jnp.float64)
+
+    # Static array of indices for first dimension of dl array
+    indices = jnp.arange(2 * L - 1)
 
     for i in range(2):
         sgn = (-1) ** (i)
@@ -142,24 +143,32 @@ def _compute_quarter_slice(
         lamb = ((el + 1) * omc - half_slices[i] + c) / s
         dl = dl.at[lims[i] + sgn].set(lamb * dl[lims[i]] * cpi[0])
 
-        for m in range(2, L):
+        def renorm_iteration(m, dl_lrenorm):
+            dl, lrenorm = dl_lrenorm
             lamb = ((el + 1) * omc - half_slices[i] + m * c) / s
             dl = dl.at[lims[i] + sgn * m].set(
                 lamb * cpi[m - 1] * dl[lims[i] + sgn * (m - 1)]
                 - cp2[m - 1] * dl[lims[i] + sgn * (m - 2)]
             )
+            condition = dl[lims[i] + sgn * m] > big_const
             lrenorm = lax.cond(
-                dl[lims[i] + sgn * m] > big_const,
-                lambda x: _increment_normalisation(x, i, lbig),
-                lambda x: x,
-                lrenorm,
+                condition, lambda x: x.at[i].add(-lbig), lambda x: x, lrenorm
             )
             dl = lax.cond(
-                dl[lims[i] + sgn * m] > big_const,
-                lambda x: _renormalise(x, lims[i], sgn, m, bigi),
+                condition,
+                # multiply first m elements (if i == 0) or last m elements (if i == 1)
+                # of dl array by bigi - use jnp.where rather than directly updating
+                # array using 'in-place' update such as
+                #     dl.at[lims[i]:lims[i] + sgn * (m + 1):sgn].multiply(bigi)
+                # to avoid non-static array slice (due to m dependence) that will raise
+                # an IndexError exception when used with lax.fori_loop
+                lambda x: jnp.where((indices < (m + 1))[::sgn], bigi * x, x),
                 lambda x: x,
-                dl,
+                dl
             )
+            return dl, lrenorm
+
+        dl, lrenorm = lax.fori_loop(2, L, renorm_iteration, (dl, lrenorm))
 
         # Apply renormalisation
         renorm = sign[i] * jnp.exp(log_first_row[half_slices[i] - 1] - lrenorm[i])
@@ -173,48 +182,6 @@ def _compute_quarter_slice(
     return jnp.nan_to_num(dl, neginf=0, posinf=0)
 
 
-@partial(jit, static_argnums=(3))
-def _renormalise(dl, lim, sgn, m, bigi) -> jnp.ndarray:
-    r"""Renormalises the recursion of the Wigner-d matrix to avoid over/underflow.
-
-    Args:
-        dl (np.ndarray): Wigner-d matrix to populate (shape: 2L-1).
-
-        lim (float): Starting index of half-line under consideration.
-
-        sgn (int): Direction that the half-line is computed.
-
-        m (int): recursion iterant at which stabilisation is needed.
-
-        bigi (float): The inverse of a very large number used to avoid over/underflows.
-
-    Returns:
-        jnp.ndarray: Renormalised Wigner-d matrix slice of dimension [2L-1].
-    """
-    for im in range(m + 1):
-        dl = dl.at[lim + sgn * im].multiply(bigi)
-    return dl
-
-
-@partial(jit)
-def _increment_normalisation(lrenorm, i, lbig) -> jnp.ndarray:
-    r"""Increments the renormalisation vector to be removed after computation.
-
-    Args:
-        lrenorm (np.ndarray): Log-renormalisation vector (shape: 2L-1).
-
-        i (float): Which half-line recursion is being considered.
-
-        lbig (float): The logarithm of a very large number used to avoid
-            over/underflows.
-
-    Returns:
-        jnp.ndarray: Renormalised Wigner-d matrix slice of dimension [2L-1].
-    """
-    lrenorm = lrenorm.at[i].set(lrenorm[i] - lbig)
-    return 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
@@ -287,13 +254,13 @@ def _el0(dl, L) -> jnp.ndarray:
 def reindex(dl, el, L) -> jnp.ndarray:
     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 
-    :math:`[-m \rightarrow -1, \dots, 0 \rightarrow m]` and the 
-    resulting indexing is given by 
-    :math:`[\dots, -m \rightarrow 0 \rightarrow m, \dots]`, where 
-    :math:`\dots` represents entries in which the values should be 
-    ignored. These extra entries are necessary to ensure :func:`~compute_slice` 
+    Reindexes the Wigner-d matrix to centre m values around L-1.
+    The original indexing is given by
+    :math:`[-m \rightarrow -1, \dots, 0 \rightarrow m]` and the
+    resulting indexing is given by
+    :math:`[\dots, -m \rightarrow 0 \rightarrow m, \dots]`, where
+    :math:`\dots` represents entries in which the values should be
+    ignored. These extra entries are necessary to ensure :func:`~compute_slice`
     can operate with static length.
 
     Args:
diff --git a/setup.py b/setup.py
index 9cbad0e7..6c1d2d35 100644
--- a/setup.py
+++ b/setup.py
@@ -1,38 +1,14 @@
-import os
-import shutil
-from setuptools import setup
+from pathlib import Path
+from setuptools import find_packages, setup
 
-import numpy
+this_directory = Path(__file__).parent
 
-# clean previous build
-for root, dirs, files in os.walk("./s2fft/", topdown=False):
-    for name in dirs:
-        if name == "build":
-            shutil.rmtree(name)
 
-from os import path
-
-this_directory = path.abspath(path.dirname(__file__))
-
-
-def read_requirements(file):
-    with open(file) as f:
-        return f.read().splitlines()
-
-
-def read_file(file):
-    with open(file) as f:
-        return f.read()
-
-
-long_description = read_file(".pip_readme.rst")
-required = read_requirements("requirements/requirements-core.txt")
-
-include_dirs = [
-    numpy.get_include(),
-]
+long_description = (this_directory / ".pip_readme.rst").read_text()
+requirements = (
+    (this_directory / "requirements" / "requirements-core.txt").read_text().split("\n")
+)
 
-extra_link_args = []
 
 setup(
     classifiers=[
@@ -45,15 +21,17 @@ def read_file(file):
     ],
     name="s2fft",
     version="0.0.1",
-    prefix=".",
     url="https://github.com/astro-informatics/s2fft",
     author="Authors & Contributors",
-    author_email="primary.author.e-mail@ucl.ac.uk",
     license="GNU General Public License v3 (GPLv3)",
     python_requires=">=3.8",
-    install_requires=required,
-    description="A witty project description!",
+    install_requires=requirements,
+    description=(
+        "Differentiable and accelerated spin-spherical harmonic transforms with JAX"
+    ),
     long_description_content_type="text/x-rst",
     long_description=long_description,
-    packages=["s2fft"],
+    packages=find_packages(),
+    include_package_data=True,
+    pacakge_data={"s2fft": ["default-logging-config.yaml"]}
 )
diff --git a/tox.ini b/tox.ini
new file mode 100644
index 00000000..f6823fd4
--- /dev/null
+++ b/tox.ini
@@ -0,0 +1,55 @@
+[tox]
+envlist = py38, check, docs
+
+[testenv]
+setenv =
+    py{38,39,310}: COVERAGE_FILE = .coverage.{envname}
+deps =
+    -r{toxinidir}{/}requirements{/}requirements-core.txt
+    -r{toxinidir}{/}requirements{/}requirements-tests.txt
+commands =
+    pytest --cov --cov-config={toxinidir}/.coveragerc {posargs:-vv}
+
+depends =
+    codecov: py{38,39,310}
+    report: py{38,39,310}
+    
+changedir = tests
+
+[testenv:check]
+description = Perform linting and formatting checks
+deps =
+    black
+    check-manifest
+    flake8
+skip_install = true
+commands =
+    check-manifest {toxinidir}
+    flake8 s2fft tests setup.py
+    black --check s2fft tests setup.py
+
+
+[testenv:report]
+skip_install = true
+deps = coverage
+commands =
+    coverage combine
+    coverage html
+    coverage report
+
+
+[testenv:codecov]
+deps =
+    codecov
+skip_install = true
+commands =
+    codecov {posargs}
+
+
+[testenv:docs]
+description = Build HTML docs with sphinx
+deps = 
+    -r{toxinidir}{/}requirements{/}requirements-docs.txt
+skip_install = true
+commands =
+    sphinx-build -M html {toxinidir}{/}docs {toxinidir}{/}docs{/}_build -Q


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 1bd8ca9c..434c9522 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -25,13 +25,12 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     
-    - name: Install package
+    - name: Install dependencies
       run: |
-        sudo apt install pandoc
-        python -m pip install --upgrade pip    
-        bash build_s2fft.sh
+        python -m pip install --upgrade pip
+        python -m pip install tox
     
     - name: Run tests
       run: |
-        pytest --cov-report term --cov=s2fft --cov-config=.coveragerc 
-        codecov --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702
+        tox -e py38
+        tox -e codecov -- --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702
diff --git a/requirements/requirements-tests.txt b/requirements/requirements-tests.txt
index b760a02d..e0e4ae17 100644
--- a/requirements/requirements-tests.txt
+++ b/requirements/requirements-tests.txt
@@ -1,15 +1,6 @@
-# Beautifiers
-black 
-pytest-black 
-
 # Testing
 pytest 
 pytest-cov
-codecov
-
-# Notebooks
-ipython==7.16.1
-jupyter==1.0.0
 
 # Healpix Sampling
 healpy
diff --git a/tests/test_samples.py b/tests/test_samples.py
index 3b981545..3dbe2419 100644
--- a/tests/test_samples.py
+++ b/tests/test_samples.py
@@ -108,7 +108,7 @@ def test_samples_exceptions():
 
     with pytest.raises(ValueError) as e:
         s2f.samples.ntheta_extension(L, sampling="foo")
-
+    
     with pytest.raises(ValueError) as e:
         s2f.samples.nphi_equiang(L, sampling="healpix")
 
diff --git a/tests/test_transform.py b/tests/test_transform.py
index 6aa444b0..bee26728 100644
--- a/tests/test_transform.py
+++ b/tests/test_transform.py
@@ -1,14 +1,14 @@
 import pytest
 import numpy as np
 import s2fft as s2f
-from s2fft.experimental import transforms_healpix as s2fhp
 import pyssht as ssht
 import healpy as hp
 
 
 L_to_test = [3, 4, 5]
-nside_to_test = [2, 4, 8]
 spin_to_test = [0, 1, 2]
+nside_to_test = [2, 4, 8]
+L_to_nside_ratio = [2, 3]
 
 
 @pytest.mark.parametrize("L", L_to_test)
@@ -32,12 +32,13 @@ def test_transform_inverse_direct(flm_generator, L: int, spin: int, sampling: st
 # @pytest.mark.skip(reason="Temporarily skipped for faster development")
 @pytest.mark.parametrize("nside", nside_to_test)
 def test_transform_inverse_direct_healpix(flm_generator, nside: int):
+    sampling = "healpix"
     L = 2 * nside
     flm = flm_generator(L=L, reality=True)
     flm_hp = s2f.samples.flm_2d_to_hp(flm, L)
     f_check = hp.sphtfunc.alm2map(flm_hp, nside, lmax=L - 1)
 
-    f = s2fhp.inverse_direct_healpix(flm, L, nside)
+    f = s2f.transform.inverse_direct(flm, L, 0, sampling, nside=nside)
 
     np.testing.assert_allclose(np.real(f), np.real(f_check), atol=1e-14)
 
@@ -64,13 +65,15 @@ def test_transform_inverse_sov(flm_generator, L: int, spin: int, sampling: str):
 
 # @pytest.mark.skip(reason="Temporarily skipped for faster development")
 @pytest.mark.parametrize("nside", nside_to_test)
-def test_transform_inverse_sov_healpix(flm_generator, nside: int):
-    L = 2 * nside
+@pytest.mark.parametrize("ratio", L_to_nside_ratio)
+def test_transform_inverse_sov_healpix(flm_generator, nside: int, ratio: int):
+    sampling = "healpix"
+    L = ratio * nside
     flm = flm_generator(L=L, reality=True)
     flm_hp = s2f.samples.flm_2d_to_hp(flm, L)
     f_check = hp.sphtfunc.alm2map(flm_hp, nside, lmax=L - 1)
 
-    f = s2fhp.inverse_sov_healpix(flm, L, nside)
+    f = s2f.transform.inverse_sov(flm, L, 0, sampling, nside=nside)
 
     np.testing.assert_allclose(np.real(f), np.real(f_check), atol=1e-14)
 
@@ -95,6 +98,21 @@ def test_transform_inverse_sov_fft(flm_generator, L: int, spin: int, sampling: s
     np.testing.assert_allclose(f, f_check, atol=1e-14)
 
 
+# @pytest.mark.skip(reason="Temporarily skipped for faster development")
+@pytest.mark.parametrize("nside", nside_to_test)
+@pytest.mark.parametrize("ratio", L_to_nside_ratio)
+def test_transform_inverse_sov_fft_healpix(flm_generator, nside: int, ratio: int):
+    sampling = "healpix"
+    L = ratio * nside
+    flm = flm_generator(L=L, reality=True)
+    flm_hp = s2f.samples.flm_2d_to_hp(flm, L)
+    f_check = hp.sphtfunc.alm2map(flm_hp, nside, lmax=L - 1)
+
+    f = s2f.transform.inverse_sov_fft(flm, L, 0, sampling, nside)
+
+    np.testing.assert_allclose(np.real(f), np.real(f_check), atol=1e-14)
+
+
 @pytest.mark.parametrize("L", L_to_test)
 @pytest.mark.parametrize("spin", spin_to_test)
 @pytest.mark.parametrize("sampling", ["mw", "mwss", "dh"])
@@ -113,7 +131,7 @@ def test_transform_inverse_sov_fft_vectorized(
 
 @pytest.mark.parametrize("L", L_to_test)
 @pytest.mark.parametrize("spin", spin_to_test)
-@pytest.mark.parametrize("sampling", ["dh"])
+@pytest.mark.parametrize("sampling", ["mw", "mwss", "dh"])
 def test_transform_forward_direct(flm_generator, L: int, spin: int, sampling: str):
 
     # TODO: move this and potentially do better
@@ -135,22 +153,23 @@ def test_transform_forward_direct(flm_generator, L: int, spin: int, sampling: st
 
 
 # @pytest.mark.skip(reason="Temporarily skipped for faster development")
-@pytest.mark.parametrize("L", L_to_test)
-def test_transform_forward_direct_healpix(flm_generator, L: int):
-    nside = 2 * L
+@pytest.mark.parametrize("nside", nside_to_test)
+def test_transform_forward_direct_healpix(flm_generator, nside: int):
+    L = 2 * nside
     flm = flm_generator(L=L, reality=True)
-    f = s2fhp.inverse_direct_healpix(flm, L, nside)
+    f = s2f.transform.inverse_direct(flm, L, sampling="healpix", nside=nside)
 
-    flm_direct = s2fhp.forward_direct_healpix(f, L, nside)
+    flm_direct = s2f.transform.forward_direct(f, L, sampling="healpix", nside=nside)
     flm_direct_hp = s2f.samples.flm_2d_to_hp(flm_direct, L)
-    flm_check = hp.sphtfunc.map2alm(np.real(f), lmax=L - 1)
 
-    np.testing.assert_allclose(flm_direct_hp, flm_check, atol=1e-2)
+    flm_check = hp.sphtfunc.map2alm(np.real(f), lmax=L - 1, iter=0)
+
+    np.testing.assert_allclose(flm_direct_hp, flm_check, atol=1e-14)
 
 
 @pytest.mark.parametrize("L", L_to_test)
 @pytest.mark.parametrize("spin", spin_to_test)
-@pytest.mark.parametrize("sampling", ["dh"])
+@pytest.mark.parametrize("sampling", ["mw", "mwss", "dh"])
 def test_transform_forward_sov(flm_generator, L: int, spin: int, sampling: str):
 
     # TODO: move this and potentially do better
@@ -202,7 +221,6 @@ def test_transform_forward_sov_fft(flm_generator, L: int, spin: int, sampling: s
 @pytest.mark.parametrize("L", [5])
 @pytest.mark.parametrize("spin", [0, 1, 2])
 @pytest.mark.parametrize("sampling", ["mw", "mwss", "dh"])
-# @pytest.mark.parametrize("sampling", ["mw"])
 def test_transform_forward_sov_fft_vectorized(
     flm_generator, L: int, spin: int, sampling: str
 ):
@@ -231,14 +249,49 @@ def test_transform_forward_sov_fft_vectorized(
 
 
 # @pytest.mark.skip(reason="Temporarily skipped for faster development")
-@pytest.mark.parametrize("L", L_to_test)
-def test_transform_forward_sov_healpix(flm_generator, L: int):
-    nside = 2 * L
+@pytest.mark.parametrize("nside", nside_to_test)
+@pytest.mark.parametrize("ratio", L_to_nside_ratio)
+def test_transform_forward_sov_healpix(flm_generator, nside: int, ratio: int):
+    sampling = "healpix"
+    L = ratio * nside
     flm = flm_generator(L=L, reality=True)
-    f = s2fhp.inverse_sov_healpix(flm, L, nside)
+    f = s2f.transform.inverse_sov(flm, L, 0, sampling, nside=nside)
 
-    flm_direct = s2fhp.forward_sov_healpix(f, L, nside)
+    flm_direct = s2f.transform.forward_sov(f, L, 0, sampling, nside)
     flm_direct_hp = s2f.samples.flm_2d_to_hp(flm_direct, L)
-    flm_check = hp.sphtfunc.map2alm(np.real(f), lmax=L - 1)
 
-    np.testing.assert_allclose(flm_direct_hp, flm_check, atol=1e-2)
+    flm_check = hp.sphtfunc.map2alm(np.real(f), lmax=L - 1, iter=0)
+
+    np.testing.assert_allclose(flm_direct_hp, flm_check, atol=1e-14)
+
+
+# @pytest.mark.skip(reason="Temporarily skipped for faster development")
+@pytest.mark.parametrize("nside", nside_to_test)
+@pytest.mark.parametrize("ratio", L_to_nside_ratio)
+def test_transform_forward_sov_fft_healpix(flm_generator, nside: int, ratio: int):
+    sampling = "healpix"
+    L = ratio * nside
+    flm = flm_generator(L=L, reality=True)
+    f = s2f.transform.inverse_sov(flm, L, 0, sampling, nside=nside)
+
+    flm_direct = s2f.transform.forward_sov_fft(f, L, 0, sampling, nside)
+    flm_direct_hp = s2f.samples.flm_2d_to_hp(flm_direct, L)
+
+    flm_check = hp.sphtfunc.map2alm(np.real(f), lmax=L - 1, iter=0)
+
+    np.testing.assert_allclose(flm_direct_hp, flm_check, atol=1e-14)
+
+
+@pytest.mark.parametrize("nside", nside_to_test)
+def test_healpix_nside_to_L_exceptions(flm_generator, nside: int):
+    sampling = "healpix"
+    L = 2 * nside - 1
+
+    flm = flm_generator(L=L, reality=True)
+    f = s2f.transform.inverse_sov(flm, L, 0, sampling, nside)
+
+    with pytest.raises(AssertionError) as e:
+        s2f.transform.forward_sov_fft(f, L, 0, sampling, nside)
+
+    with pytest.raises(AssertionError) as e:
+        s2f.transform.inverse_sov_fft(flm, L, 0, sampling, nside)