Improved docstrings in sht

10cf65f6 · apaaris · Boris Bonev · 742499ae · 10cf65f6
Commit 10cf65f6 authored Jun 30, 2025 by apaaris Committed by Boris Bonev Jul 21, 2025
Show whitespace changes
Inline Side-by-side

Showing with 142 additions and 10 deletions

torch_harmonics/sht.py torch_harmonics/sht.py +142 -10

No files found.
--- a/torch_harmonics/sht.py
+++ b/torch_harmonics/sht.py
@@ -43,6 +43,30 @@ class RealSHT(nn.Module):
    Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.
    The SHT is applied to the last two dimensions of the input

+    Parameters
+    -----------
+    nlat: int
+        Number of latitude points
+    nlon: int
+        Number of longitude points
+    lmax: int
+        Maximum spherical harmonic degree
+    mmax: int
+        Maximum spherical harmonic order
+    grid: str
+        Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+    norm: str
+        Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+    csphase: bool
+        Whether to apply the Condon-Shortley phase factor, by default True
+
+    Returns
+    -------
+    x: torch.Tensor
+        Tensor of shape (..., lmax, mmax)
+
+    References
+    ----------
    [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
    """
@@ -51,10 +75,27 @@ class RealSHT(nn.Module):
        r"""
        Initializes the SHT Layer, precomputing the necessary quadrature weights

-        Parameters:
-        nlat: input grid resolution in the latitudinal direction
-        nlon: input grid resolution in the longitudinal direction
-        grid: grid in the latitude direction (for now only tensor product grids are supported)
+        Parameters
+        -----------
+        nlat: int
+            Number of latitude points
+        nlon: int
+            Number of longitude points
+        lmax: int
+            Maximum spherical harmonic degree
+        mmax: int
+            Maximum spherical harmonic order
+        grid: str
+            Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+        norm: str
+            Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+        csphase: bool
+            Whether to apply the Condon-Shortley phase factor, by default True
+
+        Returns
+        -------
+        x: torch.Tensor
+            Tensor of shape (..., lmax, mmax)
        """

        super().__init__()
@@ -138,9 +179,35 @@ class InverseRealSHT(nn.Module):
    r"""
    Defines a module for computing the inverse (real-valued) SHT.
    Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.
-    nlat, nlon: Output dimensions
-    lmax, mmax: Input dimensions (spherical coefficients). For convenience, these are inferred from the output dimensions

+    Parameters
+    -----------
+    nlat: int
+        Number of latitude points
+    nlon: int
+        Number of longitude points
+    lmax: int
+        Maximum spherical harmonic degree
+    mmax: int
+        Maximum spherical harmonic order
+    grid: str
+        Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+    norm: str
+        Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+    csphase: bool
+        Whether to apply the Condon-Shortley phase factor, by default True
+
+    Raises
+    ------
+    ValueError: If the grid type is unknown
+
+    Returns
+    -------
+    x: torch.Tensor
+        Tensor of shape (..., lmax, mmax)
+
+    References
+    ----------
    [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
    """
@@ -218,6 +285,30 @@ class RealVectorSHT(nn.Module):
    Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.
    The SHT is applied to the last three dimensions of the input.

+    Parameters
+    -----------
+    nlat: int
+        Number of latitude points
+    nlon: int
+        Number of longitude points
+    lmax: int
+        Maximum spherical harmonic degree
+    mmax: int
+        Maximum spherical harmonic order
+    grid: str
+        Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+    norm: str
+        Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+    csphase: bool
+        Whether to apply the Condon-Shortley phase factor, by default True
+
+    Returns
+    -------
+    x: torch.Tensor
+        Tensor of shape (..., lmax, mmax)
+
+    References
+    ----------
    [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
    """
@@ -226,10 +317,27 @@ class RealVectorSHT(nn.Module):
        r"""
        Initializes the vector SHT Layer, precomputing the necessary quadrature weights

-        Parameters:
-        nlat: input grid resolution in the latitudinal direction
-        nlon: input grid resolution in the longitudinal direction
-        grid: type of grid the data lives on
+        Parameters
+        -----------
+        nlat: int
+            Number of latitude points
+        nlon: int
+            Number of longitude points
+        lmax: int
+            Maximum spherical harmonic degree
+        mmax: int
+            Maximum spherical harmonic order
+        grid: str
+            Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+        norm: str
+            Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+        csphase: bool
+            Whether to apply the Condon-Shortley phase factor, by default True
+
+        Returns
+        -------
+        x: torch.Tensor
+            Tensor of shape (..., lmax, mmax)
        """

        super().__init__()
@@ -326,6 +434,30 @@ class InverseRealVectorSHT(nn.Module):
    Defines a module for computing the inverse (real-valued) vector SHT.
    Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.

+    Parameters
+    -----------
+    nlat: int
+        Number of latitude points
+    nlon: int
+        Number of longitude points
+    lmax: int
+        Maximum spherical harmonic degree
+    mmax: int
+        Maximum spherical harmonic order
+    grid: str
+        Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+    norm: str
+        Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+    csphase: bool
+        Whether to apply the Condon-Shortley phase factor, by default True
+
+    Returns
+    -------
+    x: torch.Tensor
+        Tensor of shape (..., lmax, mmax)
+
+    References
+    ----------
    [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
    """