Replace ssht.inverse in tests

documentation/DOCUMENTATION.md

@@ -53,13 +53,15 @@ done
 
 ```python
 import numpy as np
-import pyssht as ssht
+import s2fft
 import sleplet
 
 for ell in range(2, 0, -1):
     f = sleplet.functions.HarmonicGaussian(128, l_sigma=10**ell, m_sigma=10)
     flm = f.translate(alpha=0.75 * np.pi, beta=0.125 * np.pi)
-    f_sphere = ssht.inverse(flm, f.L, Method="MWSS")
+    f_sphere = s2fft.inverse(
+        s2fft.samples.flm_1d_to_2d(flm, f.L), f.L, method="jax", sampling="mwss"
+    )
     sleplet.plotting.PlotSphere(
         f_sphere,
         f.L,
@@ -75,12 +77,14 @@ sphere earth -L 128
 ```
 
 ```python
-import pyssht as ssht
+import s2fft
 import sleplet
 
 f = sleplet.functions.Earth(128)
 flm = sleplet.harmonic_methods.rotate_earth_to_south_america(f.coefficients, f.L)
-f_sphere = ssht.inverse(flm, f.L, Method="MWSS")
+f_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(flm, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(f_sphere, f.L, "fig_2").execute()
 ```
 
@@ -93,15 +97,17 @@ done
 ```
 
 ```python
-import pyssht as ssht
+import s2fft
 import sleplet
 
 for ell in range(2, 0, -1):
     f = sleplet.functions.HarmonicGaussian(128, l_sigma=10**ell, m_sigma=10)
     g = sleplet.functions.Earth(128)
     flm = f.convolve(f.coefficients, g.coefficients.conj())
     flm_rot = sleplet.harmonic_methods.rotate_earth_to_south_america(flm, f.L)
-    f_sphere = ssht.inverse(flm_rot, f.L, Method="MWSS")
+    f_sphere = s2fft.inverse(
+        s2fft.samples.flm_1d_to_2d(flm_rot, f.L), f.L, method="jax", sampling="mwss"
+    )
     sleplet.plotting.PlotSphere(f_sphere, f.L, f"fig_3_ell_{ell}").execute()
 ```
 
@@ -126,13 +132,15 @@ sphere slepian_south_america -L 128 -s 2 -u
 ```
 
 ```python
-import pyssht as ssht
+import s2fft
 import sleplet
 
 # a
 f = sleplet.functions.Earth(128, smoothing=2)
 flm = sleplet.harmonic_methods.rotate_earth_to_south_america(f.coefficients, f.L)
-f_sphere = ssht.inverse(flm, f.L, Method="MWSS")
+f_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(flm, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(f_sphere, f.L, "fig_3_a", normalise=False).execute()
 # b
 region = sleplet.slepian.Region(mask_name="south_america")
@@ -282,13 +290,15 @@ sphere slepian_africa -L 128 -s 2 -u
 ```
 
 ```python
-import pyssht as ssht
+import s2fft
 import sleplet
 
 # a
 f = sleplet.functions.Earth(128, smoothing=2)
 flm = sleplet.harmonic_methods.rotate_earth_to_africa(f.coefficients, f.L)
-f_sphere = ssht.inverse(flm, f.L, Method="MWSS")
+f_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(flm, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(f_sphere, f.L, "fig_9_a", normalise=False).execute()
 # b
 region = sleplet.slepian.Region(mask_name="africa")
@@ -639,13 +649,18 @@ done
 ```
 
 ```python
-import pyssht as ssht
+import s2fft
 import sleplet
 
 for ell in range(5):
     for m in range(ell + 1):
         f = sleplet.functions.SphericalHarmonic(128, ell=ell, m=m)
-        f_sphere = ssht.inverse(f.coefficients, f.L, Method="MWSS")
+        f_sphere = s2fft.inverse(
+            s2fft.samples.flm_1d_to_2d(f.coefficients, f.L),
+            f.L,
+            method="jax",
+            sampling="mwss",
+        )
         sleplet.plotting.PlotSphere(
             f_sphere,
             f.L,
@@ -670,17 +685,21 @@ sphere elongated_gaussian -e -1 -1 -L 128 -m rotate -a 0.25 -b 0.25 -g 0.25
 
 ```python
 import numpy as np
-import pyssht as ssht
+import s2fft
 import sleplet
 
 # a
 f = sleplet.functions.ElongatedGaussian(128, p_sigma=0.1, t_sigma=0.1)
-f_sphere = ssht.inverse(f.coefficients, f.L, Method="MWSS")
+f_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(f.coefficients, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(f_sphere, f.L, "fig_2_2_a", annotations=[]).execute()
 # b-d
 for a, b, g in [(0, 0, 0.25), (0, 0.25, 0.25), (0.25, 0.25, 0.25)]:
     glm_rot = f.rotate(alpha=a * np.pi, beta=b * np.pi, gamma=g * np.pi)
-    g_sphere = ssht.inverse(glm_rot, f.L, Method="MWSS")
+    g_sphere = s2fft.inverse(
+        s2fft.samples.flm_1d_to_2d(glm_rot, f.L), f.L, method="jax", sampling="mwss"
+    )
     sleplet.plotting.PlotSphere(
         g_sphere,
         f.L,
@@ -713,12 +732,17 @@ done
 ```
 
 ```python
-import pyssht as ssht
+import s2fft
 import sleplet
 
 for j in [None, *list(range(4))]:
     f = sleplet.functions.AxisymmetricWavelets(128, B=3, j_min=2, j=j)
-    f_sphere = ssht.inverse(f.coefficients, f.L, Method="MWSS")
+    f_sphere = s2fft.inverse(
+        s2fft.samples.flm_1d_to_2d(f.coefficients, f.L),
+        f.L,
+        method="jax",
+        sampling="mwss",
+    )
     sleplet.plotting.PlotSphere(
         f_sphere,
         f.L,
@@ -758,16 +782,20 @@ sphere gaussian -a 0.75 -b 0.125 -L 128 -m translate -o
 
 ```python
 import numpy as np
-import pyssht as ssht
+import s2fft
 import sleplet
 
 # a
 f = sleplet.functions.Gaussian(128)
-f_sphere = ssht.inverse(f.coefficients, f.L, Method="MWSS")
+f_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(f.coefficients, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(f_sphere, f.L, "fig_3_1_a", annotations=[]).execute()
 # b
 glm_trans = f.translate(alpha=0.75 * np.pi, beta=0.125 * np.pi)
-g_sphere = ssht.inverse(glm_trans, f.L, Method="MWSS")
+g_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(glm_trans, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(g_sphere, f.L, "fig_3_1_b", annotations=[]).execute()
 ```
 
@@ -782,16 +810,20 @@ sphere squashed_gaussian -a 0.75 -b 0.125 -L 128 -m translate -o
 
 ```python
 import numpy as np
-import pyssht as ssht
+import s2fft
 import sleplet
 
 # a
 f = sleplet.functions.SquashedGaussian(128)
-f_sphere = ssht.inverse(f.coefficients, f.L, Method="MWSS")
+f_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(f.coefficients, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(f_sphere, f.L, "fig_3_2_a", annotations=[]).execute()
 # b
 glm_trans = f.translate(alpha=0.75 * np.pi, beta=0.125 * np.pi)
-g_sphere = ssht.inverse(glm_trans, f.L, Method="MWSS")
+g_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(glm_trans, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(g_sphere, f.L, "fig_3_2_b", annotations=[]).execute()
 ```
 
@@ -806,16 +838,20 @@ sphere elongated_gaussian -a 0.75 -b 0.125 -L 128 -m translate -o
 
 ```python
 import numpy as np
-import pyssht as ssht
+import s2fft
 import sleplet
 
 # a
 f = sleplet.functions.ElongatedGaussian(128)
-f_sphere = ssht.inverse(f.coefficients, f.L, Method="MWSS")
+f_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(f.coefficients, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(f_sphere, f.L, "fig_3_3_a", annotations=[]).execute()
 # b
 glm_trans = f.translate(alpha=0.75 * np.pi, beta=0.125 * np.pi)
-g_sphere = ssht.inverse(glm_trans, f.L, Method="MWSS")
+g_sphere = s2fft.inverse(
+    s2fft.samples.flm_1d_to_2d(glm_trans, f.L), f.L, method="jax", sampling="mwss"
+)
 sleplet.plotting.PlotSphere(g_sphere, f.L, "fig_3_3_b", annotations=[]).execute()
 ```
 
@@ -830,12 +866,17 @@ done
 ```
 
 ```python
-import pyssht as ssht
+import s2fft
 import sleplet
 
 for ell in range(2, 0, -1):
     f = sleplet.functions.HarmonicGaussian(128, l_sigma=10**ell, m_sigma=10)
-    f_sphere = ssht.inverse(f.coefficients, f.L, Method="MWSS")
+    f_sphere = s2fft.inverse(
+        s2fft.samples.flm_1d_to_2d(f.coefficients, f.L),
+        f.L,
+        method="jax",
+        sampling="mwss",
+    )
     sleplet.plotting.PlotSphere(
         f_sphere,
         f.L,

examples/misc/_denoising_axisym.py

@@ -1,10 +1,11 @@
 import numpy as np
 import numpy.typing as npt
 
-import pyssht as ssht
+import s2fft
 
 import sleplet
 
+EXECUTION_MODE = "jax"
 SAMPLING_SCHEME = "mwss"
 
 
@@ -60,5 +61,10 @@ def denoising_axisym(  # noqa: PLR0913
         else flm
     )
 
-    f = ssht.inverse(flm, signal.L, Method=SAMPLING_SCHEME.upper())
+    f = s2fft.inverse(
+        s2fft.samples.flm_1d_to_2d(flm, signal.L),
+        signal.L,
+        method=EXECUTION_MODE,
+        sampling=SAMPLING_SCHEME,
+    )
     return f, noised_signal.snr, denoised_snr

examples/polar_cap/slepian_error.py

@@ -3,12 +3,13 @@
 import numpy.typing as npt
 import seaborn as sns
 
-import pyssht as ssht
+import s2fft
 
 import sleplet
 
 sns.set(context="paper")
 
+EXECUTION_MODE = "jax"
 L = 16
 SAMPLING_SCHEME = "mwss"
 THETA_MAX = 40
@@ -18,7 +19,12 @@ def main() -> None:
     """Creates a plot of Slepian coefficients against rank."""
     region = sleplet.slepian.Region(theta_max=np.deg2rad(THETA_MAX))
     earth = sleplet.functions.Earth(L, region=region)
-    field = ssht.inverse(earth.coefficients, L, Method=SAMPLING_SCHEME.upper())
+    field = s2fft.inverse(
+        s2fft.samples.flm_1d_to_2d(earth.coefficients, L),
+        L,
+        method=EXECUTION_MODE,
+        sampling=SAMPLING_SCHEME,
+    )
     integrate_region = _helper_region(L, region, field, earth.coefficients)
     integrate_sphere = _helper_sphere(L, region, field, earth.coefficients)
     N = sleplet.slepian.SlepianPolarCap(L, np.deg2rad(THETA_MAX)).N

examples/wavelets/axisymmetric_wavelet_covariance.py

@@ -1,11 +1,12 @@
 import numpy as np
 import numpy.typing as npt
 
-import pyssht as ssht
+import s2fft
 
 import sleplet
 
 B = 3
+EXECUTION_MODE = "jax"
 J_MIN = 2
 L = 128
 RANDOM_SEED = 30
@@ -83,7 +84,12 @@ def axisymmetric_wavelet_covariance(
 
         # compute covariance from data
         for j, coefficient in enumerate(wlm):
-            f_wav_j = ssht.inverse(coefficient, L, Method=SAMPLING_SCHEME.upper())
+            f_wav_j = s2fft.inverse(
+                s2fft.samples.flm_1d_to_2d(coefficient, L),
+                L,
+                method=EXECUTION_MODE,
+                sampling=SAMPLING_SCHEME,
+            )
             covar_data[i, j] = (
                 f_wav_j.var() if _is_ergodic(j_min, j=j) else f_wav_j[0, 0]
             )