Merge pull request #193 from astro-informatics/feature/gl_sampling

Feature/gl sampling
astro-informatics · Mar 5, 2024 · b3d033c · b3d033c
2 parents 8e111f4 + 43b2198
commit b3d033c
Show file tree

Hide file tree

Showing 28 changed files with 251 additions and 81 deletions.
diff --git a/README.md b/README.md
@@ -57,8 +57,9 @@ isolattitude sampling scheme. A number of sampling schemes are currently
 supported.
 
 The equiangular sampling schemes of [McEwen & Wiaux
-(2012)](https://arxiv.org/abs/1110.6298) and [Driscoll & Healy
-(1995)](https://www.sciencedirect.com/science/article/pii/S0196885884710086)
+(2012)](https://arxiv.org/abs/1110.6298), [Driscoll & Healy
+(1995)](https://www.sciencedirect.com/science/article/pii/S0196885884710086) 
+and [Gauss-Legendre (1986)](https://link.springer.com/article/10.1007/BF02519350)
 are supported, which exhibit associated sampling theorems and so
 harmonic transforms can be computed to machine precision. Note that the
 McEwen & Wiaux sampling theorem reduces the Nyquist rate on the sphere

diff --git a/docs/index.rst b/docs/index.rst
@@ -21,7 +21,8 @@ Algorithms |:zap:|
 ``S2FFT`` leverages new algorithmic structures that can he highly parallelised and
 distributed, and so map very well onto the architecture of hardware accelerators (i.e.
 GPUs and TPUs).  In particular, these algorithms are based on new Wigner-d recursions
-that are stable to high angular resolution :math:`L`.  The diagram below illustrates the recursions (for further details see Price & McEwen 2023).
+that are stable to high angular resolution :math:`L`.  The diagram below illustrates the 
+recursions (for further details see Price & McEwen 2023).
 
 .. image:: ./assets/figures/Wigner_recursion_github_docs.png
 
@@ -46,7 +47,8 @@ Sampling |:earth_africa:|
 
 The structure of the algorithms implemented in ``S2FFT`` can support any isolattitude sampling scheme.  A number of sampling schemes are currently supported.
 
-The equiangular sampling schemes of `McEwen & Wiaux (2012) <https://arxiv.org/abs/1110.6298>`_ and `Driscoll & Healy (1995) <https://www.sciencedirect.com/science/article/pii/S0196885884710086>`_ are supported, which exhibit associated sampling theorems and so harmonic transforms can be computed to machine precision.  Note that the McEwen & Wiaux sampling theorem reduces the Nyquist rate on the sphere by a factor of two compared to the Driscoll & Healy approach, halving the number of spherical samples required. 
+The equiangular sampling schemes of `McEwen & Wiaux (2012) <https://arxiv.org/abs/1110.6298>`_, 
+`Driscoll & Healy (1995) <https://www.sciencedirect.com/science/article/pii/S0196885884710086>`_, and `Gauss-Legendre (1986) <https://link.springer.com/article/10.1007/BF02519350>`_ are supported, which exhibit associated sampling theorems and so harmonic transforms can be computed to machine precision.  Note that the McEwen & Wiaux sampling theorem reduces the Nyquist rate on the sphere by a factor of two compared to the Driscoll & Healy approach, halving the number of spherical samples required. 
 
 The popular `HEALPix <https://healpix.jpl.nasa.gov>`_ sampling scheme (`Gorski et al. 2005 <https://arxiv.org/abs/astro-ph/0409513>`_) is also supported.  The HEALPix sampling does not exhibit a sampling theorem and so the corresponding harmonic transforms do not achieve machine precision but exhibit some error.  However, the HEALPix sampling provides pixels of equal areas, which has many practical advantages.
 

diff --git a/requirements/requirements-core.txt b/requirements/requirements-core.txt
@@ -3,4 +3,7 @@ numpy>=1.20
 colorlog
 pyyaml
 jax>=0.3.13
-jaxlib
+jaxlib
+
+# Remove when subpackage functionality is fixed.
+torch 
diff --git a/s2fft/_version.py b/s2fft/_version.py
@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '1.0.1.dev12'
+__version_tuple__ = version_tuple = (1, 0, 1, 'dev12')
diff --git a/s2fft/base_transforms/spherical.py b/s2fft/base_transforms/spherical.py
@@ -27,7 +27,7 @@ def inverse(
         spin (int, optional): Harmonic spin. Defaults to 0.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}.  Defaults to "mw".
 
         nside (int, optional): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".  Defaults to None.

diff --git a/s2fft/base_transforms/wigner.py b/s2fft/base_transforms/wigner.py
@@ -33,7 +33,7 @@ def inverse(
         L_lower (int, optional): Harmonic lower bound. Defaults to 0.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "gl"}.  Defaults to "mw".
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.  Defaults to

diff --git a/s2fft/precompute_transforms/construct.py b/s2fft/precompute_transforms/construct.py
@@ -34,7 +34,7 @@ def spin_spherical_kernel(
             Defaults to False.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
 
         nside (int): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".
@@ -106,7 +106,7 @@ def spin_spherical_kernel_jax(
             Defaults to False.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
 
         nside (int): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".
@@ -185,7 +185,7 @@ def wigner_kernel(
             Defaults to False.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
 
         nside (int): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".
@@ -258,7 +258,7 @@ def wigner_kernel_jax(
             Defaults to False.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
 
         nside (int): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".

diff --git a/s2fft/precompute_transforms/spherical.py b/s2fft/precompute_transforms/spherical.py
@@ -32,7 +32,7 @@ def inverse(
         kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -88,7 +88,7 @@ def inverse_transform(
         L (int): Harmonic band-limit.
 
         sampling (str): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.
+            {"mw", "mwss", "dh", "gl", "healpix"}.
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -152,7 +152,7 @@ def inverse_transform_jax(
         L (int): Harmonic band-limit.
 
         sampling (str): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.
+            {"mw", "mwss", "dh", "gl", "healpix"}.
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -277,7 +277,7 @@ def forward(
         kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -333,7 +333,7 @@ def forward_transform(
         L (int): Harmonic band-limit.
 
         sampling (str): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.
+            {"mw", "mwss", "dh", "gl", "healpix"}.
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -401,7 +401,7 @@ def forward_transform_jax(
         L (int): Harmonic band-limit.
 
         sampling (str): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.
+            {"mw", "mwss", "dh", "gl", "healpix"}.
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.

diff --git a/s2fft/precompute_transforms/wigner.py b/s2fft/precompute_transforms/wigner.py
@@ -32,7 +32,7 @@ def inverse(
         kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -89,7 +89,7 @@ def inverse_transform(
         N (int): Directional band-limit.
 
         sampling (str): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.
+            {"mw", "mwss", "dh", "gl", "healpix"}.
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -150,7 +150,7 @@ def inverse_transform_jax(
         N (int): Directional band-limit.
 
         sampling (str): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.
+            {"mw", "mwss", "dh", "gl", "healpix"}.
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -306,7 +306,7 @@ def forward(
         kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -361,7 +361,7 @@ def forward_transform(
         N (int): Directional band-limit.
 
         sampling (str): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.
+            {"mw", "mwss", "dh", "gl", "healpix"}.
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.
@@ -442,7 +442,7 @@ def forward_transform_jax(
         N (int): Directional band-limit.
 
         sampling (str): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.
+            {"mw", "mwss", "dh", "gl", "healpix"}.
 
         reality (bool, optional): Whether the signal on the sphere is real.  If so,
             conjugate symmetry is exploited to reduce computational costs.

diff --git a/s2fft/sampling/s2_samples.py b/s2fft/sampling/s2_samples.py
@@ -10,7 +10,7 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int:
             Defaults to None.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}.  Defaults to "mw".
 
         nside (int, optional): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".  Defaults to None.
@@ -31,7 +31,7 @@ def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int:
             f"Sampling scheme sampling={sampling} with L={L} not supported"
         )
 
-    if sampling.lower() == "mw":
+    if sampling.lower() in ["mw", "gl"]:
         return L
 
     elif sampling.lower() == "mwss":
@@ -93,7 +93,7 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int:
         L (int): Harmonic band-limit.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "gl"}.  Defaults to "mw".
 
     Raises:
         ValueError: HEALPix sampling scheme.
@@ -104,7 +104,7 @@ def nphi_equiang(L: int, sampling: str = "mw") -> int:
         int: Number of :math:`\phi` samples.
     """
 
-    if sampling.lower() == "mw":
+    if sampling.lower() in ["mw", "gl"]:
         return 2 * L - 1
 
     elif sampling.lower() == "mwss":
@@ -129,7 +129,7 @@ def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> Tuple[int, int
         L (int): Harmonic band-limit.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}.  Defaults to "mw".
 
         nside (int, optional): HEALPix Nside resolution parameter.
 
@@ -144,7 +144,7 @@ def ftm_shape(L: int, sampling: str = "mw", nside: int = None) -> Tuple[int, int
     if sampling.lower() in ["mwss", "healpix"]:
         return ntheta(L, sampling, nside), 2 * L
 
-    elif sampling.lower() in ["mw", "dh"]:
+    elif sampling.lower() in ["mw", "dh", "gl"]:
         return ntheta(L, sampling, nside), 2 * L - 1
 
     else:
@@ -203,14 +203,17 @@ def thetas(L: int = None, sampling: str = "mw", nside: int = None) -> np.ndarray
             Defaults to None.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}.  Defaults to "mw".
 
         nside (int, optional): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".  Defaults to None.
 
     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]))
+
     t = np.arange(0, ntheta(L=L, sampling=sampling, nside=nside)).astype(np.float64)
 
     return t2theta(t, L, sampling, nside)
@@ -228,7 +231,7 @@ def t2theta(
             Defaults to None.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}.  Defaults to "mw".
 
         nside (int, optional): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".  Defaults to None.
@@ -346,7 +349,7 @@ def phis_equiang(L: int, sampling: str = "mw") -> np.ndarray:
         L (int, optional): Harmonic band-limit.
 
         sampling (str, optional): Sampling scheme.  Supported equiangular sampling
-            schemes include {"mw", "mwss", "dh"}.  Defaults to "mw".
+            schemes include {"mw", "mwss", "dh", "gl"}.  Defaults to "mw".
 
     Returns:
         np.ndarray: Array of :math:`\phi` samples for given sampling scheme.
@@ -365,7 +368,7 @@ def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray:
         p (int): :math:`\phi` index.
 
         sampling (str, optional): Sampling scheme.  Supported equiangular sampling
-            schemes include {"mw", "mwss", "dh"}.  Defaults to "mw".
+            schemes include {"mw", "mwss", "dh", "gl"}.  Defaults to "mw".
 
     Raises:
         ValueError: HEALPix sampling not support (only equiangular schemes supported).
@@ -376,7 +379,7 @@ def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray:
         np.ndarray: :math:`\phi` sample(s) for given sampling scheme.
     """
 
-    if sampling.lower() == "mw":
+    if sampling.lower() in ["mw", "gl"]:
         return 2 * p * np.pi / (2 * L - 1)
 
     elif sampling.lower() == "mwss":
@@ -431,7 +434,7 @@ def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> Tuple[int
         L (int, optional): Harmonic band-limit.
 
         sampling (str, optional): Sampling scheme.  Supported sampling schemes include
-            {"mw", "mwss", "dh", "healpix"}.  Defaults to "mw".
+            {"mw", "mwss", "dh", "gl", "healpix"}.  Defaults to "mw".
 
         nside (int, optional): HEALPix Nside resolution parameter.  Only required
             if sampling="healpix".  Defaults to None.