prune redundant code dependncies etc

Author: CosmoMatt

.github/workflows/docs.yml

@@ -4,9 +4,7 @@ on:
   push:
     branches:
     - main
-    # Add following in for testing
-    # - feature/deploy_docs
-
+    
 jobs:
   build:
 

MANIFEST.in

@@ -7,9 +7,7 @@ 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

s2fft/default-logging-config.yaml

@@ -1,5 +1,5 @@
 # =============================================================
-# Logging setup for S2fft Software package (2022)
+# Logging setup for S2fft Software package (2023)
 # =============================================================
 
 version: 1

s2fft/precompute_transforms/construct.py

@@ -1,11 +1,9 @@
-from jax import jit, config
+from jax import config
 
 config.update("jax_enable_x64", True)
 
 import numpy as np
 import jax.numpy as jnp
-from jax import jit
-from functools import partial
 
 from s2fft.sampling import s2_samples as samples
 from s2fft.utils import quadrature, quadrature_jax

s2fft/recursions/price_mcewen.py

@@ -1,5 +1,3 @@
-from jax import local_device_count
-
 import numpy as np
 import jax.numpy as jnp
 import jax.lax as lax
@@ -244,7 +242,7 @@ def renorm_m_loop(i, args):
     # Remove redundant nans:
     # - in forward pass these are not accessed, so are irrelevant.
     # - in backward pass the adjoint computation otherwise accumulates these
-    #   nans into grads if not set to zero.
+    #   nans into grads if not explicitly set to zero.
     lrenorm = jnp.nan_to_num(lrenorm, nan=0.0, posinf=0.0, neginf=0.0)
     cpi = jnp.nan_to_num(cpi, nan=0.0, posinf=0.0, neginf=0.0)
     cp2 = jnp.nan_to_num(cp2, nan=0.0, posinf=0.0, neginf=0.0)

s2fft/sampling/so3_samples.py

@@ -1,7 +1,6 @@
 import numpy as np
 from s2fft.sampling import s2_samples as samples
 from typing import Tuple
-import s2fft
 
 
 def f_shape(

s2fft/transforms/otf_recursions.py

@@ -1,7 +1,7 @@
 import numpy as np
 import jax.numpy as jnp
 import jax.lax as lax
-from jax import jit, pmap, local_device_count, custom_vjp
+from jax import jit, pmap, local_device_count
 from functools import partial
 from typing import List
 from s2fft.sampling import s2_samples as samples
@@ -379,16 +379,7 @@ 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,
@@ -818,10 +809,7 @@ 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),