Spherical Derivative Operators¶

`SphericalDerivative1D` ¶

Bases: Module

1D spectral derivative operator on Gauss-Legendre latitude grid.

Supports colatitude (theta) gradient and Laplacian using Legendre polynomial spectral coefficients.

Mathematical Formulation:¶

Expand u in normalised Legendre polynomials: u(theta) = sum_l c_l * P_l_norm(cos(theta))

Colatitude gradient (using recurrence P_l'(cos(theta)) = -P_l^1(cos(theta))): du/d_theta = -sum_l c_l * sqrt(l*(l+1)) * P_l^1_norm(cos(theta))

Laplacian on the unit sphere (eigenvalue relation): nabla^2_sphere u = -sum_l c_l * l*(l+1) / R^2 * P_l_norm(cos(theta))

Attributes:¶

grid : SphericalGrid1D The 1D Gauss-Legendre grid.

Source code in spectraldiffx/_src/spherical/operators.py

class SphericalDerivative1D(eqx.Module):
    """
    1D spectral derivative operator on Gauss-Legendre latitude grid.

    Supports colatitude (theta) gradient and Laplacian using Legendre polynomial
    spectral coefficients.

    Mathematical Formulation:
    -------------------------
    Expand u in normalised Legendre polynomials:
        u(theta) = sum_l c_l * P_l_norm(cos(theta))

    Colatitude gradient (using recurrence P_l'(cos(theta)) = -P_l^1(cos(theta))):
        du/d_theta = -sum_l c_l * sqrt(l*(l+1)) * P_l^1_norm(cos(theta))

    Laplacian on the unit sphere (eigenvalue relation):
        nabla^2_sphere u = -sum_l c_l * l*(l+1) / R^2 * P_l_norm(cos(theta))

    Attributes:
    -----------
    grid : SphericalGrid1D
        The 1D Gauss-Legendre grid.
    """

    grid: SphericalGrid1D
    # Precomputed m=1 ALP matrix for gradient reconstruction
    _P1_matrix: Float[Array, "N N"]
    # Precomputed gradient coefficients: sqrt(l*(l+1))
    _grad_coeff: Float[Array, "N"]

    def __init__(self, grid: SphericalGrid1D):
        self.grid = grid
        N = grid.N
        nodes_np = np.array(grid._nodes)
        P1_np = _gradient_alp_matrix(N, nodes_np)
        self._P1_matrix = jnp.asarray(P1_np)
        l = np.arange(N, dtype=np.float64)
        self._grad_coeff = jnp.asarray(np.sqrt(l * (l + 1)))

    def to_spectral(self, u: Float[Array, "N"]) -> Float[Array, "N"]:
        """
        Forward Discrete Legendre Transform.

        c_l = sum_j w_j * P_l_norm(cos(theta_j)) * u(theta_j)

        Parameters:
        -----------
        u : Float[Array, "N"]
            Physical field at GL nodes.

        Returns:
        --------
        c : Float[Array, "N"]
            Legendre spectral coefficients.
        """
        return self.grid.transform(u, inverse=False)

    def from_spectral(self, c: Float[Array, "N"]) -> Float[Array, "N"]:
        """
        Inverse Discrete Legendre Transform.

        u(theta_j) ≈ sum_l c_l * P_l_norm(cos(theta_j))

        Parameters:
        -----------
        c : Float[Array, "N"]
            Legendre spectral coefficients.

        Returns:
        --------
        u : Float[Array, "N"]
            Physical field at GL nodes.
        """
        return self.grid.transform(c, inverse=True)

    def gradient(
        self, u: Float[Array, "N"], spectral: bool = False
    ) -> Float[Array, "N"]:
        """
        Colatitude gradient: du/d_theta.

        Algorithm:
            1. Forward DLT: c_l = P_norm_matrix @ (w * u)
            2. Multiply:    c_grad_l = -sqrt(l*(l+1)) * c_l
            3. Reconstruct: du/d_theta = P1_norm_matrix.T @ c_grad

        Parameters:
        -----------
        u : Float[Array, "N"]
            Physical field or spectral coefficients (if spectral=True).
        spectral : bool
            If True, u contains Legendre coefficients.

        Returns:
        --------
        du_dtheta : Float[Array, "N"]
            Colatitude derivative du/d_theta at GL nodes.
        """
        c = u if spectral else self.to_spectral(u)
        # scipy.special.lpmv uses the Condon-Shortley phase convention (-1)^m.
        # For m=1, lpmv(1, l, mu) = -sqrt(1-mu^2) * dP_l/dmu, so the stored
        # P1 matrix already carries the negative sign for nodes with mu > 0.
        # Consequently d P_l_norm / d_theta = +sqrt(l*(l+1)) * P_l^1_norm
        # (no additional minus sign needed here).
        c_grad = self._grad_coeff * c
        return self._P1_matrix.T @ c_grad

    def laplacian(
        self, u: Float[Array, "N"], spectral: bool = False
    ) -> Float[Array, "N"]:
        """
        Spherical Laplacian: (1/sin(theta)) * d/d_theta [sin(theta) * du/d_theta]
        (zonal, m=0 case): nabla^2_sphere u = -l*(l+1)/R^2 * u in spectral space.

        Parameters:
        -----------
        u : Float[Array, "N"]
            Physical field or spectral coefficients.
        spectral : bool
            If True, u contains Legendre coefficients.

        Returns:
        --------
        lap_u : Float[Array, "N"]
            Laplacian at GL nodes.
        """
        R = self.grid.L / jnp.pi  # sphere radius
        c = u if spectral else self.to_spectral(u)
        l = self.grid.l
        c_lap = -(l * (l + 1)) / (R**2) * c
        return self.from_spectral(c_lap)

    def __call__(
        self, u: Float[Array, "N"], order: int = 1, spectral: bool = False
    ) -> Float[Array, "N"]:
        """
        Apply derivative operator.

        Parameters:
        -----------
        u : Float[Array, "N"]
            Physical field or spectral coefficients.
        order : int
            1 = gradient (du/d_theta), 2 = Laplacian.
        spectral : bool
            If True, u contains Legendre coefficients.

        Returns:
        --------
        Array [N]
        """
        if order == 1:
            return self.gradient(u, spectral=spectral)
        elif order == 2:
            return self.laplacian(u, spectral=spectral)
        raise ValueError(f"order must be 1 or 2, got {order}")

Functions¶

`to_spectral(u)` ¶

Forward Discrete Legendre Transform.

c_l = sum_j w_j * P_l_norm(cos(theta_j)) * u(theta_j)

Parameters:¶

u : Float[Array, "N"] Physical field at GL nodes.

Returns:¶

c : Float[Array, "N"] Legendre spectral coefficients.

Source code in spectraldiffx/_src/spherical/operators.py

def to_spectral(self, u: Float[Array, "N"]) -> Float[Array, "N"]:
    """
    Forward Discrete Legendre Transform.

    c_l = sum_j w_j * P_l_norm(cos(theta_j)) * u(theta_j)

    Parameters:
    -----------
    u : Float[Array, "N"]
        Physical field at GL nodes.

    Returns:
    --------
    c : Float[Array, "N"]
        Legendre spectral coefficients.
    """
    return self.grid.transform(u, inverse=False)

`from_spectral(c)` ¶

Inverse Discrete Legendre Transform.

u(theta_j) ≈ sum_l c_l * P_l_norm(cos(theta_j))

Parameters:¶

c : Float[Array, "N"] Legendre spectral coefficients.

Returns:¶

u : Float[Array, "N"] Physical field at GL nodes.

Source code in spectraldiffx/_src/spherical/operators.py

def from_spectral(self, c: Float[Array, "N"]) -> Float[Array, "N"]:
    """
    Inverse Discrete Legendre Transform.

    u(theta_j) ≈ sum_l c_l * P_l_norm(cos(theta_j))

    Parameters:
    -----------
    c : Float[Array, "N"]
        Legendre spectral coefficients.

    Returns:
    --------
    u : Float[Array, "N"]
        Physical field at GL nodes.
    """
    return self.grid.transform(c, inverse=True)

`gradient(u, spectral=False)` ¶

Colatitude gradient: du/d_theta.

Algorithm

Forward DLT: c_l = P_norm_matrix @ (w * u)
Multiply: c_grad_l = -sqrt(l*(l+1)) * c_l
Reconstruct: du/d_theta = P1_norm_matrix.T @ c_grad

Parameters:¶

u : Float[Array, "N"] Physical field or spectral coefficients (if spectral=True). spectral : bool If True, u contains Legendre coefficients.

Returns:¶

du_dtheta : Float[Array, "N"] Colatitude derivative du/d_theta at GL nodes.

Source code in spectraldiffx/_src/spherical/operators.py

def gradient(
    self, u: Float[Array, "N"], spectral: bool = False
) -> Float[Array, "N"]:
    """
    Colatitude gradient: du/d_theta.

    Algorithm:
        1. Forward DLT: c_l = P_norm_matrix @ (w * u)
        2. Multiply:    c_grad_l = -sqrt(l*(l+1)) * c_l
        3. Reconstruct: du/d_theta = P1_norm_matrix.T @ c_grad

    Parameters:
    -----------
    u : Float[Array, "N"]
        Physical field or spectral coefficients (if spectral=True).
    spectral : bool
        If True, u contains Legendre coefficients.

    Returns:
    --------
    du_dtheta : Float[Array, "N"]
        Colatitude derivative du/d_theta at GL nodes.
    """
    c = u if spectral else self.to_spectral(u)
    # scipy.special.lpmv uses the Condon-Shortley phase convention (-1)^m.
    # For m=1, lpmv(1, l, mu) = -sqrt(1-mu^2) * dP_l/dmu, so the stored
    # P1 matrix already carries the negative sign for nodes with mu > 0.
    # Consequently d P_l_norm / d_theta = +sqrt(l*(l+1)) * P_l^1_norm
    # (no additional minus sign needed here).
    c_grad = self._grad_coeff * c
    return self._P1_matrix.T @ c_grad

`laplacian(u, spectral=False)` ¶

Spherical Laplacian: (1/sin(theta)) * d/d_theta [sin(theta) * du/d_theta] (zonal, m=0 case): nabla^2_sphere u = -l*(l+1)/R^2 * u in spectral space.

Parameters:¶

u : Float[Array, "N"] Physical field or spectral coefficients. spectral : bool If True, u contains Legendre coefficients.

Returns:¶

lap_u : Float[Array, "N"] Laplacian at GL nodes.

Source code in spectraldiffx/_src/spherical/operators.py

def laplacian(
    self, u: Float[Array, "N"], spectral: bool = False
) -> Float[Array, "N"]:
    """
    Spherical Laplacian: (1/sin(theta)) * d/d_theta [sin(theta) * du/d_theta]
    (zonal, m=0 case): nabla^2_sphere u = -l*(l+1)/R^2 * u in spectral space.

    Parameters:
    -----------
    u : Float[Array, "N"]
        Physical field or spectral coefficients.
    spectral : bool
        If True, u contains Legendre coefficients.

    Returns:
    --------
    lap_u : Float[Array, "N"]
        Laplacian at GL nodes.
    """
    R = self.grid.L / jnp.pi  # sphere radius
    c = u if spectral else self.to_spectral(u)
    l = self.grid.l
    c_lap = -(l * (l + 1)) / (R**2) * c
    return self.from_spectral(c_lap)

`call(u, order=1, spectral=False)` ¶

Apply derivative operator.

Parameters:¶

u : Float[Array, "N"] Physical field or spectral coefficients. order : int 1 = gradient (du/d_theta), 2 = Laplacian. spectral : bool If True, u contains Legendre coefficients.

Returns:¶

Array [N]

Source code in spectraldiffx/_src/spherical/operators.py

def __call__(
    self, u: Float[Array, "N"], order: int = 1, spectral: bool = False
) -> Float[Array, "N"]:
    """
    Apply derivative operator.

    Parameters:
    -----------
    u : Float[Array, "N"]
        Physical field or spectral coefficients.
    order : int
        1 = gradient (du/d_theta), 2 = Laplacian.
    spectral : bool
        If True, u contains Legendre coefficients.

    Returns:
    --------
    Array [N]
    """
    if order == 1:
        return self.gradient(u, spectral=spectral)
    elif order == 2:
        return self.laplacian(u, spectral=spectral)
    raise ValueError(f"order must be 1 or 2, got {order}")

`SphericalDerivative2D` ¶

Bases: Module

2D pseudo-spectral derivative operators on the sphere.

Computes physical-space differential operators (gradient, divergence, curl, Laplacian) for a field u(theta, phi) on the sphere using the full Spherical Harmonic Transform.

Mathematical Formulation:¶

For a field u(theta, phi) on a sphere of radius R = Ly / pi:

Gradient (covariant components):
    grad_theta u = (1/R) * du/d_theta
    grad_phi   u = 1 / (R * sin(theta)) * du/d_phi

Divergence of vector (V_theta, V_phi) (in physical coordinates):
    div V = 1/(R*sin(theta)) * [d(V_theta*sin(theta))/d_theta + d(V_phi)/d_phi]

Scalar curl (vertical component of curl):
    curl V = 1/(R*sin(theta)) * [d(V_theta)/d_phi - d(V_phi*sin(theta))/d_theta]

Laplacian (via eigenvalue in spectral space):
    nabla^2 u = sum_{l,m} -l*(l+1)/R^2 * u_hat(l,m) * Y_l^m

Attributes:¶

grid : SphericalGrid2D The full lat-lon grid. deriv_theta : SphericalDerivative1D 1D operator reused for latitude derivatives.

Source code in spectraldiffx/_src/spherical/operators.py

class SphericalDerivative2D(eqx.Module):
    """
    2D pseudo-spectral derivative operators on the sphere.

    Computes physical-space differential operators (gradient, divergence, curl,
    Laplacian) for a field u(theta, phi) on the sphere using the full Spherical
    Harmonic Transform.

    Mathematical Formulation:
    -------------------------
    For a field u(theta, phi) on a sphere of radius R = Ly / pi:

        Gradient (covariant components):
            grad_theta u = (1/R) * du/d_theta
            grad_phi   u = 1 / (R * sin(theta)) * du/d_phi

        Divergence of vector (V_theta, V_phi) (in physical coordinates):
            div V = 1/(R*sin(theta)) * [d(V_theta*sin(theta))/d_theta + d(V_phi)/d_phi]

        Scalar curl (vertical component of curl):
            curl V = 1/(R*sin(theta)) * [d(V_theta)/d_phi - d(V_phi*sin(theta))/d_theta]

        Laplacian (via eigenvalue in spectral space):
            nabla^2 u = sum_{l,m} -l*(l+1)/R^2 * u_hat(l,m) * Y_l^m

    Attributes:
    -----------
    grid : SphericalGrid2D
        The full lat-lon grid.
    deriv_theta : SphericalDerivative1D
        1D operator reused for latitude derivatives.
    """

    grid: SphericalGrid2D
    deriv_theta: SphericalDerivative1D

    def __init__(self, grid: SphericalGrid2D):
        self.grid = grid
        # Build a compatible 1D grid from the latitude parameters
        grid_1d = SphericalGrid1D(
            N=grid.Ny,
            L=grid.Ly,
            dx=grid.dy,
            dealias=grid.dealias,
        )
        self.deriv_theta = SphericalDerivative1D(grid_1d)

    def gradient(
        self, u: Float[Array, "Ny Nx"], spectral: bool = False
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """
        Compute the gradient of u on the sphere.

            grad_theta u = (1/R) * du/d_theta
            grad_phi   u = 1 / (R * sin(theta)) * du/d_phi

        Parameters:
        -----------
        u : Float[Array, "Ny Nx"]
            Physical field.
        spectral : bool
            If True, u is already in spectral (SHT) space.

        Returns:
        --------
        (grad_theta, grad_phi) : tuple of Float[Array, "Ny Nx"]
            Covariant gradient components.
        """
        R = self.grid.Ly / jnp.pi

        # --- theta-component: vmap 1D gradient over longitude columns ---
        u_phys = self.grid.transform(u, inverse=True).real if spectral else u
        du_dtheta = jax.vmap(self.deriv_theta.gradient, in_axes=1, out_axes=1)(
            u_phys
        )  # (Ny, Nx)

        # --- phi-component via FFT ---
        u_hat_fft = jnp.fft.fft(u_phys, axis=-1)  # (Ny, Nx)
        m_phys = 2 * jnp.pi * jnp.fft.fftfreq(self.grid.Nx, self.grid.dx)  # (Nx,)
        du_dphi = jnp.fft.ifft(
            1j * m_phys[None, :] * u_hat_fft, axis=-1
        ).real  # (Ny, Nx)

        # Apply metric factors
        sin_theta = jnp.sin(self.grid.y)[:, None]  # (Ny, 1)
        grad_theta = du_dtheta / R
        grad_phi = du_dphi / (R * sin_theta)
        return grad_theta, grad_phi

    def divergence(
        self,
        v_theta: Float[Array, "Ny Nx"],
        v_phi: Float[Array, "Ny Nx"],
        spectral: bool = False,
    ) -> Float[Array, "Ny Nx"]:
        """
        Divergence of a vector field (V_theta, V_phi) on the sphere.

            div V = 1/(R*sin(theta)) * [d(V_theta*sin(theta))/d_theta + dV_phi/d_phi]

        Parameters:
        -----------
        v_theta : Float[Array, "Ny Nx"]
            Colatitude component (physical space).
        v_phi : Float[Array, "Ny Nx"]
            Longitude component (physical space).
        spectral : bool
            Unused (reserved for API consistency; inputs expected in physical space).

        Returns:
        --------
        div : Float[Array, "Ny Nx"]
            Scalar divergence field.
        """
        R = self.grid.Ly / jnp.pi
        sin_theta = jnp.sin(self.grid.y)[:, None]  # (Ny, 1)

        # d(V_theta * sin(theta))/d_theta
        vs = v_theta * sin_theta  # (Ny, Nx)
        d_vs_dtheta = jax.vmap(self.deriv_theta.gradient, in_axes=1, out_axes=1)(
            vs
        )  # (Ny, Nx)

        # dV_phi/d_phi via FFT
        m_phys = 2 * jnp.pi * jnp.fft.fftfreq(self.grid.Nx, self.grid.dx)  # (Nx,)
        vp_hat = jnp.fft.fft(v_phi, axis=-1)  # (Ny, Nx)
        d_vp_dphi = jnp.fft.ifft(
            1j * m_phys[None, :] * vp_hat, axis=-1
        ).real  # (Ny, Nx)

        return (d_vs_dtheta + d_vp_dphi) / (R * sin_theta)

    def curl(
        self,
        v_theta: Num[Array, "Nlat Nlon"],
        v_phi: Num[Array, "Nlat Nlon"],
        spectral: bool = False,
    ) -> Float[Array, "Nlat Nlon"]:
        """Scalar (radial) vorticity (∇×V)·r̂ of a tangent field V on the sphere.

            (∇×V)·r̂ = (1 / (R·sin θ)) · [∂(V_φ·sin θ)/∂θ − ∂V_θ/∂φ]

        This is the standard GFD sign convention: for solid-body rotation
        V_φ = Ω R sin θ the formula returns the planetary vorticity
        2 Ω cos θ, positive in the northern hemisphere.

        Parameters
        ----------
        v_theta : Num[Array, "Nlat Nlon"]
            Colatitude component.
        v_phi : Num[Array, "Nlat Nlon"]
            Longitude component.
        spectral : bool
            Unused (reserved for API consistency).

        Returns
        -------
        Float[Array, "Nlat Nlon"]
            Scalar vorticity.
        """
        R = self.grid.Ly / jnp.pi
        sin_theta = jnp.sin(self.grid.y)[:, None]  # (Nlat, 1)

        # ∂V_θ/∂φ via FFT along the longitude axis.
        m_phys = 2 * jnp.pi * jnp.fft.fftfreq(self.grid.Nx, self.grid.dx)  # (Nlon,)
        vt_hat = jnp.fft.fft(v_theta, axis=-1)
        d_vt_dphi = jnp.fft.ifft(1j * m_phys[None, :] * vt_hat, axis=-1).real

        # ∂(V_φ·sin θ)/∂θ via the GL-latitude derivative.
        vs = v_phi * sin_theta
        d_vs_dtheta = jax.vmap(self.deriv_theta.gradient, in_axes=1, out_axes=1)(vs)

        # Standard curl convention (positive in NH for solid-body rotation):
        return (d_vs_dtheta - d_vt_dphi) / (R * sin_theta)

    def laplacian(
        self,
        u: Num[Array, "Nlat Nlon"],
        spectral: bool = False,
    ) -> Num[Array, "Nlat Nlon"]:
        """Scalar Laplace–Beltrami operator on the sphere.

        In spectral space:

            ∇² û(l, m) = −l(l+1)/R² · û(l, m)

        Parameters
        ----------
        u : Num[Array, "Nlat Nlon"]
            Physical field, or SHT coefficients if ``spectral=True``.
        spectral : bool
            If ``True``, ``u`` is already in spectral space.

        Returns
        -------
        Num[Array, "Nlat Nlon"]
            Real (float) array in physical space when ``spectral=False``;
            complex SHT coefficients when ``spectral=True``.
        """
        return self.iterated_laplacian(u, n=1, spectral=spectral)

    def iterated_laplacian(
        self,
        u: Num[Array, "Nlat Nlon"],
        n: int = 1,
        spectral: bool = False,
    ) -> Num[Array, "Nlat Nlon"]:
        """Iterated Laplacian ∇^(2n) — useful for hyperdiffusion.

        In spectral space:

            ∇^(2n) û(l, m) = [−l(l+1)/R²]ⁿ · û(l, m)

        Parameters
        ----------
        u : Num[Array, "Nlat Nlon"]
            Physical field or SHT coefficients.
        n : int
            Number of Laplacian applications (≥ 1).  ``n=2`` is the
            biharmonic ∇⁴, etc.
        spectral : bool
            If ``True``, ``u`` is already spectral and the result stays
            spectral; otherwise the function transforms to and from
            physical space.

        Returns
        -------
        Num[Array, "Nlat Nlon"]
            ∇^(2n) u.  Real-valued (float) when ``spectral=False``; complex
            SHT coefficients when ``spectral=True``.

        Examples
        --------
        >>> lap_u = deriv.iterated_laplacian(u, n=1)
        >>> bih_u = deriv.iterated_laplacian(u, n=2)  # biharmonic ∇⁴ u
        """
        if n < 1:
            raise ValueError(f"n must be >= 1, got {n}")
        R = self.grid.Ly / jnp.pi
        u_hat = u if spectral else self.grid.transform(u)
        l = self.grid.l  # (Nl,)
        lap_eig = -(l * (l + 1)) / (R**2)  # (Nl,)
        eig_n = lap_eig**n  # (Nl,)
        out_hat = eig_n[:, None] * u_hat  # (Nl, Nm)
        if spectral:
            return out_hat
        return self.grid.transform(out_hat, inverse=True)

    def biharmonic(
        self,
        u: Num[Array, "Nlat Nlon"],
        spectral: bool = False,
    ) -> Num[Array, "Nlat Nlon"]:
        """Biharmonic operator ∇⁴ u (alias for ``iterated_laplacian(u, n=2)``).

        ∇⁴ û(l, m) = [l(l+1)/R²]² · û(l, m)

        Returns a real (float) array when ``spectral=False`` and complex
        SHT coefficients when ``spectral=True``.
        """
        return self.iterated_laplacian(u, n=2, spectral=spectral)

    def advection_scalar(
        self,
        v_theta: Num[Array, "Nlat Nlon"],
        v_phi: Num[Array, "Nlat Nlon"],
        q: Num[Array, "Nlat Nlon"],
    ) -> Float[Array, "Nlat Nlon"]:
        """Pseudo-spectral scalar advection (V·∇) q.

        Evaluates the gradient of q spectrally, then multiplies by V
        in physical space:

            (V·∇) q = V_θ · ∇_θ q + V_φ · ∇_φ q
        """
        grad_theta_q, grad_phi_q = self.gradient(q)
        return v_theta * grad_theta_q + v_phi * grad_phi_q

Functions¶

`gradient(u, spectral=False)` ¶

Compute the gradient of u on the sphere.

grad_theta u = (1/R) * du/d_theta
grad_phi   u = 1 / (R * sin(theta)) * du/d_phi

Parameters:¶

u : Float[Array, "Ny Nx"] Physical field. spectral : bool If True, u is already in spectral (SHT) space.

Returns:¶

(grad_theta, grad_phi) : tuple of Float[Array, "Ny Nx"] Covariant gradient components.

Source code in spectraldiffx/_src/spherical/operators.py

def gradient(
    self, u: Float[Array, "Ny Nx"], spectral: bool = False
) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
    """
    Compute the gradient of u on the sphere.

        grad_theta u = (1/R) * du/d_theta
        grad_phi   u = 1 / (R * sin(theta)) * du/d_phi

    Parameters:
    -----------
    u : Float[Array, "Ny Nx"]
        Physical field.
    spectral : bool
        If True, u is already in spectral (SHT) space.

    Returns:
    --------
    (grad_theta, grad_phi) : tuple of Float[Array, "Ny Nx"]
        Covariant gradient components.
    """
    R = self.grid.Ly / jnp.pi

    # --- theta-component: vmap 1D gradient over longitude columns ---
    u_phys = self.grid.transform(u, inverse=True).real if spectral else u
    du_dtheta = jax.vmap(self.deriv_theta.gradient, in_axes=1, out_axes=1)(
        u_phys
    )  # (Ny, Nx)

    # --- phi-component via FFT ---
    u_hat_fft = jnp.fft.fft(u_phys, axis=-1)  # (Ny, Nx)
    m_phys = 2 * jnp.pi * jnp.fft.fftfreq(self.grid.Nx, self.grid.dx)  # (Nx,)
    du_dphi = jnp.fft.ifft(
        1j * m_phys[None, :] * u_hat_fft, axis=-1
    ).real  # (Ny, Nx)

    # Apply metric factors
    sin_theta = jnp.sin(self.grid.y)[:, None]  # (Ny, 1)
    grad_theta = du_dtheta / R
    grad_phi = du_dphi / (R * sin_theta)
    return grad_theta, grad_phi

`divergence(v_theta, v_phi, spectral=False)` ¶

Divergence of a vector field (V_theta, V_phi) on the sphere.

div V = 1/(R*sin(theta)) * [d(V_theta*sin(theta))/d_theta + dV_phi/d_phi]

Parameters:¶

v_theta : Float[Array, "Ny Nx"] Colatitude component (physical space). v_phi : Float[Array, "Ny Nx"] Longitude component (physical space). spectral : bool Unused (reserved for API consistency; inputs expected in physical space).

Returns:¶

div : Float[Array, "Ny Nx"] Scalar divergence field.

Source code in spectraldiffx/_src/spherical/operators.py

def divergence(
    self,
    v_theta: Float[Array, "Ny Nx"],
    v_phi: Float[Array, "Ny Nx"],
    spectral: bool = False,
) -> Float[Array, "Ny Nx"]:
    """
    Divergence of a vector field (V_theta, V_phi) on the sphere.

        div V = 1/(R*sin(theta)) * [d(V_theta*sin(theta))/d_theta + dV_phi/d_phi]

    Parameters:
    -----------
    v_theta : Float[Array, "Ny Nx"]
        Colatitude component (physical space).
    v_phi : Float[Array, "Ny Nx"]
        Longitude component (physical space).
    spectral : bool
        Unused (reserved for API consistency; inputs expected in physical space).

    Returns:
    --------
    div : Float[Array, "Ny Nx"]
        Scalar divergence field.
    """
    R = self.grid.Ly / jnp.pi
    sin_theta = jnp.sin(self.grid.y)[:, None]  # (Ny, 1)

    # d(V_theta * sin(theta))/d_theta
    vs = v_theta * sin_theta  # (Ny, Nx)
    d_vs_dtheta = jax.vmap(self.deriv_theta.gradient, in_axes=1, out_axes=1)(
        vs
    )  # (Ny, Nx)

    # dV_phi/d_phi via FFT
    m_phys = 2 * jnp.pi * jnp.fft.fftfreq(self.grid.Nx, self.grid.dx)  # (Nx,)
    vp_hat = jnp.fft.fft(v_phi, axis=-1)  # (Ny, Nx)
    d_vp_dphi = jnp.fft.ifft(
        1j * m_phys[None, :] * vp_hat, axis=-1
    ).real  # (Ny, Nx)

    return (d_vs_dtheta + d_vp_dphi) / (R * sin_theta)

`curl(v_theta, v_phi, spectral=False)` ¶

Scalar (radial) vorticity (∇×V)·r̂ of a tangent field V on the sphere.

(∇×V)·r̂ = (1 / (R·sin θ)) · [∂(V_φ·sin θ)/∂θ − ∂V_θ/∂φ]

This is the standard GFD sign convention: for solid-body rotation V_φ = Ω R sin θ the formula returns the planetary vorticity 2 Ω cos θ, positive in the northern hemisphere.

Parameters¶

v_theta : Num[Array, "Nlat Nlon"] Colatitude component. v_phi : Num[Array, "Nlat Nlon"] Longitude component. spectral : bool Unused (reserved for API consistency).

Returns¶

Float[Array, "Nlat Nlon"] Scalar vorticity.

Source code in spectraldiffx/_src/spherical/operators.py

def curl(
    self,
    v_theta: Num[Array, "Nlat Nlon"],
    v_phi: Num[Array, "Nlat Nlon"],
    spectral: bool = False,
) -> Float[Array, "Nlat Nlon"]:
    """Scalar (radial) vorticity (∇×V)·r̂ of a tangent field V on the sphere.

        (∇×V)·r̂ = (1 / (R·sin θ)) · [∂(V_φ·sin θ)/∂θ − ∂V_θ/∂φ]

    This is the standard GFD sign convention: for solid-body rotation
    V_φ = Ω R sin θ the formula returns the planetary vorticity
    2 Ω cos θ, positive in the northern hemisphere.

    Parameters
    ----------
    v_theta : Num[Array, "Nlat Nlon"]
        Colatitude component.
    v_phi : Num[Array, "Nlat Nlon"]
        Longitude component.
    spectral : bool
        Unused (reserved for API consistency).

    Returns
    -------
    Float[Array, "Nlat Nlon"]
        Scalar vorticity.
    """
    R = self.grid.Ly / jnp.pi
    sin_theta = jnp.sin(self.grid.y)[:, None]  # (Nlat, 1)

    # ∂V_θ/∂φ via FFT along the longitude axis.
    m_phys = 2 * jnp.pi * jnp.fft.fftfreq(self.grid.Nx, self.grid.dx)  # (Nlon,)
    vt_hat = jnp.fft.fft(v_theta, axis=-1)
    d_vt_dphi = jnp.fft.ifft(1j * m_phys[None, :] * vt_hat, axis=-1).real

    # ∂(V_φ·sin θ)/∂θ via the GL-latitude derivative.
    vs = v_phi * sin_theta
    d_vs_dtheta = jax.vmap(self.deriv_theta.gradient, in_axes=1, out_axes=1)(vs)

    # Standard curl convention (positive in NH for solid-body rotation):
    return (d_vs_dtheta - d_vt_dphi) / (R * sin_theta)

`laplacian(u, spectral=False)` ¶

Scalar Laplace–Beltrami operator on the sphere.

In spectral space:

∇² û(l, m) = −l(l+1)/R² · û(l, m)

Parameters¶

u : Num[Array, "Nlat Nlon"] Physical field, or SHT coefficients if spectral=True. spectral : bool If True, u is already in spectral space.

Returns¶

Num[Array, "Nlat Nlon"] Real (float) array in physical space when spectral=False; complex SHT coefficients when spectral=True.

Source code in spectraldiffx/_src/spherical/operators.py

def laplacian(
    self,
    u: Num[Array, "Nlat Nlon"],
    spectral: bool = False,
) -> Num[Array, "Nlat Nlon"]:
    """Scalar Laplace–Beltrami operator on the sphere.

    In spectral space:

        ∇² û(l, m) = −l(l+1)/R² · û(l, m)

    Parameters
    ----------
    u : Num[Array, "Nlat Nlon"]
        Physical field, or SHT coefficients if ``spectral=True``.
    spectral : bool
        If ``True``, ``u`` is already in spectral space.

    Returns
    -------
    Num[Array, "Nlat Nlon"]
        Real (float) array in physical space when ``spectral=False``;
        complex SHT coefficients when ``spectral=True``.
    """
    return self.iterated_laplacian(u, n=1, spectral=spectral)

`iterated_laplacian(u, n=1, spectral=False)` ¶

Iterated Laplacian ∇^(2n) — useful for hyperdiffusion.

In spectral space:

∇^(2n) û(l, m) = [−l(l+1)/R²]ⁿ · û(l, m)

Parameters¶

u : Num[Array, "Nlat Nlon"] Physical field or SHT coefficients. n : int Number of Laplacian applications (≥ 1). n=2 is the biharmonic ∇⁴, etc. spectral : bool If True, u is already spectral and the result stays spectral; otherwise the function transforms to and from physical space.

Returns¶

Num[Array, "Nlat Nlon"] ∇^(2n) u. Real-valued (float) when spectral=False; complex SHT coefficients when spectral=True.

Examples¶

lap_u = deriv.iterated_laplacian(u, n=1) bih_u = deriv.iterated_laplacian(u, n=2) # biharmonic ∇⁴ u

Source code in spectraldiffx/_src/spherical/operators.py

def iterated_laplacian(
    self,
    u: Num[Array, "Nlat Nlon"],
    n: int = 1,
    spectral: bool = False,
) -> Num[Array, "Nlat Nlon"]:
    """Iterated Laplacian ∇^(2n) — useful for hyperdiffusion.

    In spectral space:

        ∇^(2n) û(l, m) = [−l(l+1)/R²]ⁿ · û(l, m)

    Parameters
    ----------
    u : Num[Array, "Nlat Nlon"]
        Physical field or SHT coefficients.
    n : int
        Number of Laplacian applications (≥ 1).  ``n=2`` is the
        biharmonic ∇⁴, etc.
    spectral : bool
        If ``True``, ``u`` is already spectral and the result stays
        spectral; otherwise the function transforms to and from
        physical space.

    Returns
    -------
    Num[Array, "Nlat Nlon"]
        ∇^(2n) u.  Real-valued (float) when ``spectral=False``; complex
        SHT coefficients when ``spectral=True``.

    Examples
    --------
    >>> lap_u = deriv.iterated_laplacian(u, n=1)
    >>> bih_u = deriv.iterated_laplacian(u, n=2)  # biharmonic ∇⁴ u
    """
    if n < 1:
        raise ValueError(f"n must be >= 1, got {n}")
    R = self.grid.Ly / jnp.pi
    u_hat = u if spectral else self.grid.transform(u)
    l = self.grid.l  # (Nl,)
    lap_eig = -(l * (l + 1)) / (R**2)  # (Nl,)
    eig_n = lap_eig**n  # (Nl,)
    out_hat = eig_n[:, None] * u_hat  # (Nl, Nm)
    if spectral:
        return out_hat
    return self.grid.transform(out_hat, inverse=True)

`biharmonic(u, spectral=False)` ¶

Biharmonic operator ∇⁴ u (alias for iterated_laplacian(u, n=2)).

∇⁴ û(l, m) = [l(l+1)/R²]² · û(l, m)

Returns a real (float) array when spectral=False and complex SHT coefficients when spectral=True.

Source code in spectraldiffx/_src/spherical/operators.py

def biharmonic(
    self,
    u: Num[Array, "Nlat Nlon"],
    spectral: bool = False,
) -> Num[Array, "Nlat Nlon"]:
    """Biharmonic operator ∇⁴ u (alias for ``iterated_laplacian(u, n=2)``).

    ∇⁴ û(l, m) = [l(l+1)/R²]² · û(l, m)

    Returns a real (float) array when ``spectral=False`` and complex
    SHT coefficients when ``spectral=True``.
    """
    return self.iterated_laplacian(u, n=2, spectral=spectral)

`advection_scalar(v_theta, v_phi, q)` ¶

Pseudo-spectral scalar advection (V·∇) q.

Evaluates the gradient of q spectrally, then multiplies by V in physical space:

(V·∇) q = V_θ · ∇_θ q + V_φ · ∇_φ q

Source code in spectraldiffx/_src/spherical/operators.py

def advection_scalar(
    self,
    v_theta: Num[Array, "Nlat Nlon"],
    v_phi: Num[Array, "Nlat Nlon"],
    q: Num[Array, "Nlat Nlon"],
) -> Float[Array, "Nlat Nlon"]:
    """Pseudo-spectral scalar advection (V·∇) q.

    Evaluates the gradient of q spectrally, then multiplies by V
    in physical space:

        (V·∇) q = V_θ · ∇_θ q + V_φ · ∇_φ q
    """
    grad_theta_q, grad_phi_q = self.gradient(q)
    return v_theta * grad_theta_q + v_phi * grad_phi_q

Spherical Derivative Operators¶

SphericalDerivative1D ¶

Mathematical Formulation:¶

Attributes:¶

Functions¶

to_spectral(u) ¶

Parameters:¶

Returns:¶

from_spectral(c) ¶

Parameters:¶

Returns:¶

gradient(u, spectral=False) ¶

Parameters:¶

Returns:¶

laplacian(u, spectral=False) ¶

Parameters:¶

Returns:¶

__call__(u, order=1, spectral=False) ¶

Parameters:¶

Returns:¶

SphericalDerivative2D ¶

Mathematical Formulation:¶

Attributes:¶

Functions¶

gradient(u, spectral=False) ¶

Parameters:¶

Returns:¶

divergence(v_theta, v_phi, spectral=False) ¶

Parameters:¶

Returns:¶

curl(v_theta, v_phi, spectral=False) ¶

Parameters¶

Returns¶

laplacian(u, spectral=False) ¶

Parameters¶

Returns¶

iterated_laplacian(u, n=1, spectral=False) ¶

Parameters¶

Returns¶

Examples¶

biharmonic(u, spectral=False) ¶

advection_scalar(v_theta, v_phi, q) ¶

`SphericalDerivative1D` ¶

`to_spectral(u)` ¶

`from_spectral(c)` ¶

`gradient(u, spectral=False)` ¶

`laplacian(u, spectral=False)` ¶

`call(u, order=1, spectral=False)` ¶

`SphericalDerivative2D` ¶

`gradient(u, spectral=False)` ¶

`divergence(v_theta, v_phi, spectral=False)` ¶

`curl(v_theta, v_phi, spectral=False)` ¶

`laplacian(u, spectral=False)` ¶

`iterated_laplacian(u, n=1, spectral=False)` ¶

`biharmonic(u, spectral=False)` ¶

`advection_scalar(v_theta, v_phi, q)` ¶