Replace ssht.sample_positions & MWSS -> mwss (#221)

astro-informatics · Sep 29, 2023 · dcd9ca8 · dcd9ca8
1 parent cbf148b
commit dcd9ca8
Show file tree

Hide file tree

Showing 20 changed files with 185 additions and 81 deletions.
diff --git a/examples/arbitrary/_slepian_wavelet_covariance.py b/examples/arbitrary/_slepian_wavelet_covariance.py
@@ -5,7 +5,7 @@
 
 import sleplet
 
-SAMPLING_SCHEME = "MWSS"
+SAMPLING_SCHEME = "mwss"
 
 
 def compute_slepian_wavelet_covariance(
@@ -29,7 +29,7 @@ def compute_slepian_wavelet_covariance(
         covariance.shape[0],
     )
     np.testing.assert_equal(
-        ssht.sample_shape(slepian_wavelets.L, Method=SAMPLING_SCHEME),
+        ssht.sample_shape(slepian_wavelets.L, Method=SAMPLING_SCHEME.upper()),
         covariance.shape[1:],
     )
     return covariance
diff --git a/examples/misc/_denoising_axisym.py b/examples/misc/_denoising_axisym.py
@@ -5,7 +5,7 @@
 
 import sleplet
 
-SAMPLING_SCHEME = "MWSS"
+SAMPLING_SCHEME = "mwss"
 
 
 def denoising_axisym(  # noqa: PLR0913
@@ -60,5 +60,5 @@ def denoising_axisym(  # noqa: PLR0913
         else flm
     )
 
-    f = ssht.inverse(flm, signal.L, Method=SAMPLING_SCHEME)
+    f = ssht.inverse(flm, signal.L, Method=SAMPLING_SCHEME.upper())
     return f, noised_signal.snr, denoised_snr
diff --git a/examples/misc/translation_normalisation.py b/examples/misc/translation_normalisation.py
@@ -1,5 +1,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
+import s2fft
 import seaborn as sns
 
 import pyssht as ssht
@@ -10,13 +11,13 @@
 
 ALPHA_DEFAULT = 0.75
 L = 128
-SAMPLING_SCHEME = "MWSS"
+SAMPLING_SCHEME = "mwss"
 
 
 def compute_translation_normalisation_theta() -> None:
     """Analysis of the translation norm for referee."""
     hg = sleplet.functions.HarmonicGaussian(L)
-    thetas, _ = ssht.sample_positions(L, Method=SAMPLING_SCHEME)
+    thetas = s2fft.samples.thetas(L, sampling=SAMPLING_SCHEME)
     norm = np.zeros(len(thetas))
     for i, theta in enumerate(thetas):
         print(f"compute norm {i+1}/{len(thetas)}")

diff --git a/examples/polar_cap/slepian_error.py b/examples/polar_cap/slepian_error.py
@@ -10,15 +10,15 @@
 sns.set(context="paper")
 
 L = 16
-SAMPLING_SCHEME = "MWSS"
+SAMPLING_SCHEME = "mwss"
 THETA_MAX = 40
 
 
 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)
+    field = ssht.inverse(earth.coefficients, L, Method=SAMPLING_SCHEME.upper())
     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

diff --git a/examples/wavelets/axisymmetric_wavelet_covariance.py b/examples/wavelets/axisymmetric_wavelet_covariance.py
@@ -9,7 +9,7 @@
 J_MIN = 2
 L = 128
 RANDOM_SEED = 30
-SAMPLING_SCHEME = "MWSS"
+SAMPLING_SCHEME = "mwss"
 
 
 def _compute_wavelet_covariance(
@@ -83,7 +83,7 @@ def axisymmetric_wavelet_covariance(
 
         # compute covariance from data
         for j, coefficient in enumerate(wlm):
-            f_wav_j = ssht.inverse(coefficient, L, Method=SAMPLING_SCHEME)
+            f_wav_j = ssht.inverse(coefficient, L, Method=SAMPLING_SCHEME.upper())
             covar_data[i, j] = (
                 f_wav_j.var() if _is_ergodic(j_min, j=j) else f_wav_j[0, 0]
             )

diff --git a/src/sleplet/_integration_methods.py b/src/sleplet/_integration_methods.py
@@ -3,18 +3,26 @@
 
 import numpy as np
 import numpy.typing as npt
-
-import pyssht as ssht
+import s2fft
 
 import sleplet._vars
 
 
 def calc_integration_weight(L: int) -> npt.NDArray[np.float_]:
     """Computes the spherical Jacobian for the integration."""
-    thetas, phis = ssht.sample_positions(
-        L,
-        Grid=True,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+    thetas = np.tile(
+        s2fft.samples.thetas(L, sampling=sleplet._vars.SAMPLING_SCHEME),
+        (
+            s2fft.samples.nphi_equiang(
+                L,
+                sampling=sleplet._vars.SAMPLING_SCHEME,
+            ),
+            1,
+        ),
+    ).T
+    phis = np.tile(
+        s2fft.samples.phis_equiang(L, sampling=sleplet._vars.SAMPLING_SCHEME),
+        (s2fft.samples.ntheta(L, sampling=sleplet._vars.SAMPLING_SCHEME), 1),
     )
     delta_theta = np.ediff1d(thetas[:, 0]).mean()
     delta_phi = np.ediff1d(phis[0]).mean()

diff --git a/src/sleplet/_mask_methods.py b/src/sleplet/_mask_methods.py
@@ -4,6 +4,7 @@
 import numpy as np
 import numpy.typing as npt
 import platformdirs
+import s2fft
 
 import pyssht as ssht
 
@@ -32,10 +33,19 @@ def create_mask_region(
                                    phi_min or phi_max is provided
     * arbitrary - just checks the shape of the input mask.
     """
-    thetas, phis = ssht.sample_positions(
-        L,
-        Grid=True,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+    thetas = np.tile(
+        s2fft.samples.thetas(L, sampling=sleplet._vars.SAMPLING_SCHEME),
+        (
+            s2fft.samples.nphi_equiang(
+                L,
+                sampling=sleplet._vars.SAMPLING_SCHEME,
+            ),
+            1,
+        ),
+    ).T
+    phis = np.tile(
+        s2fft.samples.phis_equiang(L, sampling=sleplet._vars.SAMPLING_SCHEME),
+        (s2fft.samples.ntheta(L, sampling=sleplet._vars.SAMPLING_SCHEME), 1),
     )
 
     match region.region_type:
@@ -88,7 +98,7 @@ def ensure_masked_flm_bandlimited(
         L,
         Reality=reality,
         Spin=spin,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+        Method=sleplet._vars.SAMPLING_SCHEME.upper(),
     )
     mask = create_mask_region(L, region)
     field = np.where(mask, field, 0)
@@ -97,7 +107,7 @@ def ensure_masked_flm_bandlimited(
         L,
         Reality=reality,
         Spin=spin,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+        Method=sleplet._vars.SAMPLING_SCHEME.upper(),
     )
 
 
@@ -159,13 +169,18 @@ def _create_africa_mask(
         rot_flm,
         L,
         Reality=True,
-        Method=sleplet._vars.SAMPLING_SCHEME,
-    )
-    thetas, _ = ssht.sample_positions(
-        L,
-        Grid=True,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+        Method=sleplet._vars.SAMPLING_SCHEME.upper(),
     )
+    thetas = np.tile(
+        s2fft.samples.thetas(L, sampling=sleplet._vars.SAMPLING_SCHEME),
+        (
+            s2fft.samples.nphi_equiang(
+                L,
+                sampling=sleplet._vars.SAMPLING_SCHEME,
+            ),
+            1,
+        ),
+    ).T
     return (thetas <= _AFRICA_RANGE) & (earth_f >= 0)
 
 
@@ -179,13 +194,18 @@ def _create_south_america_mask(
         rot_flm,
         L,
         Reality=True,
-        Method=sleplet._vars.SAMPLING_SCHEME,
-    )
-    thetas, _ = ssht.sample_positions(
-        L,
-        Grid=True,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+        Method=sleplet._vars.SAMPLING_SCHEME.upper(),
     )
+    thetas = np.tile(
+        s2fft.samples.thetas(L, sampling=sleplet._vars.SAMPLING_SCHEME),
+        (
+            s2fft.samples.nphi_equiang(
+                L,
+                sampling=sleplet._vars.SAMPLING_SCHEME,
+            ),
+            1,
+        ),
+    ).T
     return (thetas <= _SOUTH_AMERICA_RANGE) & (earth_f >= 0)
 
 

diff --git a/src/sleplet/_vars.py b/src/sleplet/_vars.py
@@ -4,7 +4,7 @@
 PHI_MAX_DEFAULT = 2 * np.pi
 PHI_MIN_DEFAULT = 0.0
 RANDOM_SEED = 30
-SAMPLING_SCHEME = "MWSS"
+SAMPLING_SCHEME = "mwss"
 SPHERE_UNSEEN = -1.56e30
 THETA_0 = 0.0
 THETA_MAX_DEFAULT = np.pi

diff --git a/src/sleplet/functions/africa.py b/src/sleplet/functions/africa.py
@@ -65,7 +65,7 @@ def _grid_fun(
             rot_flm,
             self.L,
             Reality=self.reality,
-            Method=sleplet._vars.SAMPLING_SCHEME,
+            Method=sleplet._vars.SAMPLING_SCHEME.upper(),
         )
         mask_name = f"{self.name}_L{self.L}.npy"
         mask_location = sleplet._data.setup_pooch.find_on_pooch_then_local(

diff --git a/src/sleplet/functions/slepian_dirac_delta.py b/src/sleplet/functions/slepian_dirac_delta.py
@@ -4,6 +4,7 @@
 import numpy as np
 import numpy.typing as npt
 import pydantic.v1 as pydantic
+import s2fft
 
 import pyssht as ssht
 
@@ -52,15 +53,36 @@ def _setup_args(self) -> None:
 
     def _compute_angles(self) -> None:
         """Computes alpha/beta if not provided."""
-        thetas, phis = ssht.sample_positions(
-            self.L,
-            Grid=True,
-            Method=sleplet._vars.SAMPLING_SCHEME,
+        thetas = np.tile(
+            s2fft.samples.thetas(
+                self.L,
+                sampling=sleplet._vars.SAMPLING_SCHEME,
+            ),
+            (
+                s2fft.samples.nphi_equiang(
+                    self.L,
+                    sampling=sleplet._vars.SAMPLING_SCHEME,
+                ),
+                1,
+            ),
+        ).T
+        phis = np.tile(
+            s2fft.samples.phis_equiang(
+                self.L,
+                sampling=sleplet._vars.SAMPLING_SCHEME,
+            ),
+            (
+                s2fft.samples.ntheta(
+                    self.L,
+                    sampling=sleplet._vars.SAMPLING_SCHEME,
+                ),
+                1,
+            ),
         )
         sp = ssht.inverse(
             self.slepian.eigenvectors[0],
             self.L,
-            Method=sleplet._vars.SAMPLING_SCHEME,
+            Method=sleplet._vars.SAMPLING_SCHEME.upper(),
         )
         idx = tuple(np.argwhere(sp == sp.max())[0])
         self._alpha = phis[idx]

diff --git a/src/sleplet/functions/south_america.py b/src/sleplet/functions/south_america.py
@@ -65,7 +65,7 @@ def _grid_fun(
             rot_flm,
             self.L,
             Reality=self.reality,
-            Method=sleplet._vars.SAMPLING_SCHEME,
+            Method=sleplet._vars.SAMPLING_SCHEME.upper(),
         )
         mask_name = f"{self.name}_L{self.L}.npy"
         mask_location = sleplet._data.setup_pooch.find_on_pooch_then_local(

diff --git a/src/sleplet/harmonic_methods.py b/src/sleplet/harmonic_methods.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 import numpy.typing as npt
+import s2fft
 
 import pyssht as ssht
 
@@ -63,7 +64,7 @@ def invert_flm_boosted(
         resolution,
         Reality=reality,
         Spin=spin,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+        Method=sleplet._vars.SAMPLING_SCHEME.upper(),
     )
 
 
@@ -81,18 +82,27 @@ def _ensure_f_bandlimited(
     If the function created is created in pixel space rather than harmonic
     space then need to transform it into harmonic space first before using it.
     """
-    thetas, phis = ssht.sample_positions(
-        L,
-        Grid=True,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+    thetas = np.tile(
+        s2fft.samples.thetas(L, sampling=sleplet._vars.SAMPLING_SCHEME),
+        (
+            s2fft.samples.nphi_equiang(
+                L,
+                sampling=sleplet._vars.SAMPLING_SCHEME,
+            ),
+            1,
+        ),
+    ).T
+    phis = np.tile(
+        s2fft.samples.phis_equiang(L, sampling=sleplet._vars.SAMPLING_SCHEME),
+        (s2fft.samples.ntheta(L, sampling=sleplet._vars.SAMPLING_SCHEME), 1),
     )
     f = grid_fun(thetas, phis)
     return ssht.forward(
         f,
         L,
         Reality=reality,
         Spin=spin,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+        Method=sleplet._vars.SAMPLING_SCHEME.upper(),
     )
 
 

diff --git a/src/sleplet/noise.py b/src/sleplet/noise.py
@@ -104,7 +104,7 @@ def _create_slepian_noise(
     flm = ssht.forward(
         sleplet.slepian_methods.slepian_inverse(slepian_signal, L, slepian),
         L,
-        Method=sleplet._vars.SAMPLING_SCHEME,
+        Method=sleplet._vars.SAMPLING_SCHEME.upper(),
     )
     nlm = _create_noise(L, flm, snr_in)
     return sleplet.slepian_methods.slepian_forward(L, slepian, flm=nlm)
@@ -141,12 +141,12 @@ def harmonic_hard_thresholding(
     _logger.info("begin harmonic hard thresholding")
     for j, coefficient in enumerate(wav_coeffs[1:]):
         _logger.info(f"start Psi^{j + 1}/{len(wav_coeffs)-1}")
-        f = ssht.inverse(coefficient, L, Method=sleplet._vars.SAMPLING_SCHEME)
+        f = ssht.inverse(coefficient, L, Method=sleplet._vars.SAMPLING_SCHEME.upper())
         f_thresholded = _perform_hard_thresholding(f, sigma_j[j], n_sigma)
         wav_coeffs[j + 1] = ssht.forward(
             f_thresholded,
             L,
-            Method=sleplet._vars.SAMPLING_SCHEME,
+            Method=sleplet._vars.SAMPLING_SCHEME.upper(),
         )
     return wav_coeffs