Skip to content

Spherical Harmonic Solvers

SphericalPoissonSolver

Bases: Module

Spectral Poisson solver on the sphere: ∇²φ = f.

In SHT-coefficient space the mode-by-mode inversion is

φ̂(l, m) = −f̂(l, m) · [l(l+1)/R²]⁻¹    (l ≥ 1)

The l=0 mode is always set to zero (∇² annihilates constants on the sphere, so it is undefined).

Attributes:

Name Type Description
grid SphericalGrid1D or SphericalGrid2D

Underlying spherical grid.

Examples:

>>> import jax.numpy as jnp
>>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
>>> solver = SphericalPoissonSolver(grid=grid)
>>> PHI, THETA = grid.X
>>> # Laplacian of cos(θ) is −2 cos(θ)/R², so Poisson RHS is that:
>>> R = grid.Ly / jnp.pi
>>> f = -2.0 * jnp.cos(THETA) / R**2
>>> phi = solver.solve(f)  # ≈ cos(θ) up to an additive constant
Source code in spectraldiffx/_src/spherical/solvers.py 
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
class SphericalPoissonSolver(eqx.Module):
    """Spectral Poisson solver on the sphere:  ∇²φ = f.

    In SHT-coefficient space the mode-by-mode inversion is

        φ̂(l, m) = −f̂(l, m) · [l(l+1)/R²]⁻¹    (l ≥ 1)

    The l=0 mode is always set to zero (∇² annihilates constants on the
    sphere, so it is undefined).

    Attributes
    ----------
    grid : SphericalGrid1D or SphericalGrid2D
        Underlying spherical grid.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
    >>> solver = SphericalPoissonSolver(grid=grid)
    >>> PHI, THETA = grid.X
    >>> # Laplacian of cos(θ) is −2 cos(θ)/R², so Poisson RHS is that:
    >>> R = grid.Ly / jnp.pi
    >>> f = -2.0 * jnp.cos(THETA) / R**2
    >>> phi = solver.solve(f)  # ≈ cos(θ) up to an additive constant
    """

    grid: SphericalGrid1D | SphericalGrid2D

    def solve(
        self,
        f: Num[Array, "..."],
        zero_mean: bool = True,
        spectral: bool = False,
    ) -> Float[Array, "..."]:
        """Solve ∇²φ = f on the sphere.

        Parameters
        ----------
        f : Num[Array, ...]
            Source field.  Shape ``(N,)`` for 1D or ``(Nlat, Nlon)`` for 2D.
        zero_mean : bool
            Must be ``True`` (default): the l=0 mode of φ is undefined and is
            set to zero. ``False`` raises ``ValueError`` (gh-92).
        spectral : bool
            If ``True``, ``f`` is already a DLT/SHT coefficient array.

        Returns
        -------
        Float[Array, ...]
            Solution in physical space (same shape as ``f``).
        """
        if not zero_mean:
            raise ValueError(
                "zero_mean=False: the l=0 mode of a Poisson solution on the "
                "sphere is undefined (gh-92)."
            )
        R = _sphere_radius(self.grid)
        f_hat = f if spectral else self.grid.transform(f)
        l = self.grid.l
        eigenval = l * (l + 1) / (R**2)
        if not isinstance(self.grid, SphericalGrid1D):
            eigenval = eigenval[:, None]
        # Guard the l=0 division, then set the undefined mode to zero.
        denom = jnp.where(eigenval == 0.0, 1.0, eigenval)
        phi_hat = jnp.where(eigenval == 0.0, 0.0, -f_hat / denom)

        return self.grid.transform(phi_hat, inverse=True)

Functions

solve(f, zero_mean=True, spectral=False)

Solve ∇²φ = f on the sphere.

Parameters:

Name Type Description Default
f Num[Array, ...]

Source field. Shape (N,) for 1D or (Nlat, Nlon) for 2D.

 required
zero_mean bool

Must be True (default): the l=0 mode of φ is undefined and is set to zero. False raises ValueError (gh-92).

 True
spectral bool

If True, f is already a DLT/SHT coefficient array.

 False

Returns:

Type Description
Float[Array, ...]

Solution in physical space (same shape as f).
Source code in spectraldiffx/_src/spherical/solvers.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def solve(
    self,
    f: Num[Array, "..."],
    zero_mean: bool = True,
    spectral: bool = False,
) -> Float[Array, "..."]:
    """Solve ∇²φ = f on the sphere.

    Parameters
    ----------
    f : Num[Array, ...]
        Source field.  Shape ``(N,)`` for 1D or ``(Nlat, Nlon)`` for 2D.
    zero_mean : bool
        Must be ``True`` (default): the l=0 mode of φ is undefined and is
        set to zero. ``False`` raises ``ValueError`` (gh-92).
    spectral : bool
        If ``True``, ``f`` is already a DLT/SHT coefficient array.

    Returns
    -------
    Float[Array, ...]
        Solution in physical space (same shape as ``f``).
    """
    if not zero_mean:
        raise ValueError(
            "zero_mean=False: the l=0 mode of a Poisson solution on the "
            "sphere is undefined (gh-92)."
        )
    R = _sphere_radius(self.grid)
    f_hat = f if spectral else self.grid.transform(f)
    l = self.grid.l
    eigenval = l * (l + 1) / (R**2)
    if not isinstance(self.grid, SphericalGrid1D):
        eigenval = eigenval[:, None]
    # Guard the l=0 division, then set the undefined mode to zero.
    denom = jnp.where(eigenval == 0.0, 1.0, eigenval)
    phi_hat = jnp.where(eigenval == 0.0, 0.0, -f_hat / denom)

    return self.grid.transform(phi_hat, inverse=True)

SphericalHelmholtzSolver

Bases: Module

Spectral Helmholtz solver on the sphere: (∇² − α) φ = f.

In SHT-coefficient space:

φ̂(l, m) = −f̂(l, m) / [l(l+1)/R² + α]

Non-singular for α > 0, where the l=0 mode (the mean) is solved like every other mode. For α = 0 this reduces to Poisson and the undefined l=0 mode is set to zero (gh-92).

Attributes:

Name Type Description
grid SphericalGrid1D or SphericalGrid2D

Underlying spherical grid.

Examples:

>>> import jax.numpy as jnp
>>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
>>> solver = SphericalHelmholtzSolver(grid=grid)
>>> PHI, THETA = grid.X
>>> R = grid.Ly / jnp.pi
>>> alpha = 4.0
>>> # For φ = cos θ: (∇² − α) φ = (−2/R² − α) cos θ
>>> f = (-2.0 / R**2 - alpha) * jnp.cos(THETA)
>>> phi = solver.solve(f, alpha=alpha)  # ≈ cos(θ)
Source code in spectraldiffx/_src/spherical/solvers.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
class SphericalHelmholtzSolver(eqx.Module):
    """Spectral Helmholtz solver on the sphere:  (∇² − α) φ = f.

    In SHT-coefficient space:

        φ̂(l, m) = −f̂(l, m) / [l(l+1)/R² + α]

    Non-singular for α > 0, where the l=0 mode (the mean) is solved like
    every other mode. For α = 0 this reduces to Poisson and the undefined
    l=0 mode is set to zero (gh-92).

    Attributes
    ----------
    grid : SphericalGrid1D or SphericalGrid2D
        Underlying spherical grid.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
    >>> solver = SphericalHelmholtzSolver(grid=grid)
    >>> PHI, THETA = grid.X
    >>> R = grid.Ly / jnp.pi
    >>> alpha = 4.0
    >>> # For φ = cos θ: (∇² − α) φ = (−2/R² − α) cos θ
    >>> f = (-2.0 / R**2 - alpha) * jnp.cos(THETA)
    >>> phi = solver.solve(f, alpha=alpha)  # ≈ cos(θ)
    """

    grid: SphericalGrid1D | SphericalGrid2D

    def solve(
        self,
        f: Num[Array, "..."],
        alpha: float = 0.0,
        zero_mean: bool | None = None,
        spectral: bool = False,
    ) -> Float[Array, "..."]:
        """Solve (∇² − α) φ = f on the sphere.

        Parameters
        ----------
        f : Num[Array, ...]
            Source field (1D ``(N,)`` or 2D ``(Nlat, Nlon)``).
        alpha : float
            Helmholtz parameter (≥ 0).  α=0 falls back to Poisson.
        zero_mean : bool or None
            ``None`` (default): zero the l=0 mode only when α = 0, where it
            is undefined. ``True``: always zero it. ``False``: keep it; an
            error when α = 0 (gh-92).
        spectral : bool
            If ``True``, ``f`` is already a DLT/SHT coefficient array.

        Returns
        -------
        Float[Array, ...]
            Solution in physical space.
        """
        if alpha < 0:
            raise ValueError(f"alpha must be >= 0, got {alpha}")
        if zero_mean is False and alpha == 0:
            raise ValueError(
                "zero_mean=False with alpha=0: the l=0 mode of the solution is "
                "undefined (gh-92). Use zero_mean=None (default) or True."
            )
        R = _sphere_radius(self.grid)
        f_hat = f if spectral else self.grid.transform(f)
        l = self.grid.l
        if not isinstance(self.grid, SphericalGrid1D):
            l = l[:, None]
        denom = l * (l + 1) / (R**2) + alpha
        denom_safe = jnp.where(denom == 0.0, 1.0, denom)
        # An undefined (zero-denominator) mode is set to 0, never -f_hat.
        phi_hat = jnp.where(denom == 0.0, 0.0, -f_hat / denom_safe)
        if zero_mean:
            phi_hat = jnp.where(l == 0.0, 0.0, phi_hat)

        return self.grid.transform(phi_hat, inverse=True)

Functions

solve(f, alpha=0.0, zero_mean=None, spectral=False)

Solve (∇² − α) φ = f on the sphere.

Parameters:

Name Type Description Default
f Num[Array, ...]

Source field (1D (N,) or 2D (Nlat, Nlon)).

 required
alpha float

Helmholtz parameter (≥ 0). α=0 falls back to Poisson.

 0.0
zero_mean bool or None

None (default): zero the l=0 mode only when α = 0, where it is undefined. True: always zero it. False: keep it; an error when α = 0 (gh-92).

 None
spectral bool

If True, f is already a DLT/SHT coefficient array.

 False

Returns:

Type Description
Float[Array, ...]

Solution in physical space.
Source code in spectraldiffx/_src/spherical/solvers.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
def solve(
    self,
    f: Num[Array, "..."],
    alpha: float = 0.0,
    zero_mean: bool | None = None,
    spectral: bool = False,
) -> Float[Array, "..."]:
    """Solve (∇² − α) φ = f on the sphere.

    Parameters
    ----------
    f : Num[Array, ...]
        Source field (1D ``(N,)`` or 2D ``(Nlat, Nlon)``).
    alpha : float
        Helmholtz parameter (≥ 0).  α=0 falls back to Poisson.
    zero_mean : bool or None
        ``None`` (default): zero the l=0 mode only when α = 0, where it
        is undefined. ``True``: always zero it. ``False``: keep it; an
        error when α = 0 (gh-92).
    spectral : bool
        If ``True``, ``f`` is already a DLT/SHT coefficient array.

    Returns
    -------
    Float[Array, ...]
        Solution in physical space.
    """
    if alpha < 0:
        raise ValueError(f"alpha must be >= 0, got {alpha}")
    if zero_mean is False and alpha == 0:
        raise ValueError(
            "zero_mean=False with alpha=0: the l=0 mode of the solution is "
            "undefined (gh-92). Use zero_mean=None (default) or True."
        )
    R = _sphere_radius(self.grid)
    f_hat = f if spectral else self.grid.transform(f)
    l = self.grid.l
    if not isinstance(self.grid, SphericalGrid1D):
        l = l[:, None]
    denom = l * (l + 1) / (R**2) + alpha
    denom_safe = jnp.where(denom == 0.0, 1.0, denom)
    # An undefined (zero-denominator) mode is set to 0, never -f_hat.
    phi_hat = jnp.where(denom == 0.0, 0.0, -f_hat / denom_safe)
    if zero_mean:
        phi_hat = jnp.where(l == 0.0, 0.0, phi_hat)

    return self.grid.transform(phi_hat, inverse=True)

Geophysical inversions

SphericalVorticityInversionSolver

Bases: Module

Vorticity-inversion solver on the sphere.

Given scalar vorticity ζ = (∇×V)·r̂, solves for the streamfunction ψ via Poisson

∇²ψ = ζ

and returns the rotational (non-divergent) velocity field

V = ẑ × ∇ψ       ⇔       V_θ = −(1/sin θ) ∂ψ/∂φ ·(1/R),
                           V_φ = (1/R) ∂ψ/∂θ.

This is the canonical vorticity–streamfunction inversion used in barotropic and quasigeostrophic spherical models.

Attributes:

Name Type Description
grid SphericalGrid2D

Underlying 2D lat-lon grid.

Examples:

>>> import jax.numpy as jnp
>>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
>>> solver = SphericalVorticityInversionSolver(grid=grid)
>>> # zonal vorticity ζ = −2 cos θ / R² corresponds to ψ = cos θ
>>> PHI, THETA = grid.X
>>> R = grid.Ly / jnp.pi
>>> zeta = -2.0 * jnp.cos(THETA) / R**2
>>> psi, (v_theta, v_phi) = solver.solve(zeta)
Source code in spectraldiffx/_src/spherical/solvers.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
class SphericalVorticityInversionSolver(eqx.Module):
    """Vorticity-inversion solver on the sphere.

    Given scalar vorticity ζ = (∇×V)·r̂, solves for the streamfunction ψ
    via Poisson

        ∇²ψ = ζ

    and returns the rotational (non-divergent) velocity field

        V = ẑ × ∇ψ       ⇔       V_θ = −(1/sin θ) ∂ψ/∂φ ·(1/R),
                                   V_φ = (1/R) ∂ψ/∂θ.

    This is the canonical vorticity–streamfunction inversion used in
    barotropic and quasigeostrophic spherical models.

    Attributes
    ----------
    grid : SphericalGrid2D
        Underlying 2D lat-lon grid.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
    >>> solver = SphericalVorticityInversionSolver(grid=grid)
    >>> # zonal vorticity ζ = −2 cos θ / R² corresponds to ψ = cos θ
    >>> PHI, THETA = grid.X
    >>> R = grid.Ly / jnp.pi
    >>> zeta = -2.0 * jnp.cos(THETA) / R**2
    >>> psi, (v_theta, v_phi) = solver.solve(zeta)
    """

    grid: SphericalGrid2D

    def solve(
        self,
        zeta: Num[Array, "Nlat Nlon"],
        spectral: bool = False,
    ) -> tuple[
        Float[Array, "Nlat Nlon"],
        tuple[Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"]],
    ]:
        """Solve ∇²ψ = ζ and recover the rotational velocity V = ẑ × ∇ψ.

        Parameters
        ----------
        zeta : Num[Array, "Nlat Nlon"]
            Vorticity in physical space (or spectral if ``spectral=True``).
        spectral : bool
            If ``True``, treat ``zeta`` as SHT coefficients.

        Returns
        -------
        (psi, (v_theta, v_phi))
            Streamfunction ψ in physical space, and the tangent velocity
            field decomposed into (colatitude, longitude) components.
        """
        poisson = SphericalPoissonSolver(grid=self.grid)
        psi = poisson.solve(zeta, zero_mean=True, spectral=spectral)
        # V = ẑ × ∇ψ   ⇒   V_θ = −∇_φ ψ,  V_φ = +∇_θ ψ.
        deriv = SphericalDerivative2D(grid=self.grid)
        grad_theta_psi, grad_phi_psi = deriv.gradient(psi)
        v_theta = -grad_phi_psi
        v_phi = grad_theta_psi
        return psi, (v_theta, v_phi)

Functions

solve(zeta, spectral=False)

Solve ∇²ψ = ζ and recover the rotational velocity V = ẑ × ∇ψ.

Parameters:

Name Type Description Default
zeta Num[Array, 'Nlat Nlon']

Vorticity in physical space (or spectral if spectral=True).

 required
spectral bool

If True, treat zeta as SHT coefficients.

 False

Returns:

Type Description
(psi, (v_theta, v_phi))

Streamfunction ψ in physical space, and the tangent velocity field decomposed into (colatitude, longitude) components.
Source code in spectraldiffx/_src/spherical/solvers.py 
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
def solve(
    self,
    zeta: Num[Array, "Nlat Nlon"],
    spectral: bool = False,
) -> tuple[
    Float[Array, "Nlat Nlon"],
    tuple[Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"]],
]:
    """Solve ∇²ψ = ζ and recover the rotational velocity V = ẑ × ∇ψ.

    Parameters
    ----------
    zeta : Num[Array, "Nlat Nlon"]
        Vorticity in physical space (or spectral if ``spectral=True``).
    spectral : bool
        If ``True``, treat ``zeta`` as SHT coefficients.

    Returns
    -------
    (psi, (v_theta, v_phi))
        Streamfunction ψ in physical space, and the tangent velocity
        field decomposed into (colatitude, longitude) components.
    """
    poisson = SphericalPoissonSolver(grid=self.grid)
    psi = poisson.solve(zeta, zero_mean=True, spectral=spectral)
    # V = ẑ × ∇ψ   ⇒   V_θ = −∇_φ ψ,  V_φ = +∇_θ ψ.
    deriv = SphericalDerivative2D(grid=self.grid)
    grad_theta_psi, grad_phi_psi = deriv.gradient(psi)
    v_theta = -grad_phi_psi
    v_phi = grad_theta_psi
    return psi, (v_theta, v_phi)

SphericalDivergenceInversionSolver

Bases: Module

Divergence-inversion solver on the sphere.

Given horizontal divergence δ = ∇·V, solves for the velocity potential χ via Poisson

∇²χ = δ

and returns the irrotational (curl-free) velocity field

V = ∇χ       ⇔       V_θ = (1/R) ∂χ/∂θ,  V_φ = (1/(R sin θ)) ∂χ/∂φ.

Attributes:

Name Type Description
grid SphericalGrid2D

Underlying 2D lat-lon grid.

Examples:

>>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
>>> solver = SphericalDivergenceInversionSolver(grid=grid)
>>> delta = ...  # horizontal divergence field
>>> chi, (v_theta, v_phi) = solver.solve(delta)
Source code in spectraldiffx/_src/spherical/solvers.py 
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
class SphericalDivergenceInversionSolver(eqx.Module):
    """Divergence-inversion solver on the sphere.

    Given horizontal divergence δ = ∇·V, solves for the velocity
    potential χ via Poisson

        ∇²χ = δ

    and returns the irrotational (curl-free) velocity field

        V = ∇χ       ⇔       V_θ = (1/R) ∂χ/∂θ,  V_φ = (1/(R sin θ)) ∂χ/∂φ.

    Attributes
    ----------
    grid : SphericalGrid2D
        Underlying 2D lat-lon grid.

    Examples
    --------
    >>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
    >>> solver = SphericalDivergenceInversionSolver(grid=grid)
    >>> delta = ...  # horizontal divergence field  # doctest: +SKIP
    >>> chi, (v_theta, v_phi) = solver.solve(delta)  # doctest: +SKIP
    """

    grid: SphericalGrid2D

    def solve(
        self,
        delta: Num[Array, "Nlat Nlon"],
        spectral: bool = False,
    ) -> tuple[
        Float[Array, "Nlat Nlon"],
        tuple[Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"]],
    ]:
        """Solve ∇²χ = δ and recover the irrotational velocity V = ∇χ."""
        poisson = SphericalPoissonSolver(grid=self.grid)
        chi = poisson.solve(delta, zero_mean=True, spectral=spectral)
        deriv = SphericalDerivative2D(grid=self.grid)
        grad_theta_chi, grad_phi_chi = deriv.gradient(chi)
        return chi, (grad_theta_chi, grad_phi_chi)

Functions

solve(delta, spectral=False)

Solve ∇²χ = δ and recover the irrotational velocity V = ∇χ.

Source code in spectraldiffx/_src/spherical/solvers.py 
310
311
312
313
314
315
316
317
318
319
320
321
322
323
def solve(
    self,
    delta: Num[Array, "Nlat Nlon"],
    spectral: bool = False,
) -> tuple[
    Float[Array, "Nlat Nlon"],
    tuple[Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"]],
]:
    """Solve ∇²χ = δ and recover the irrotational velocity V = ∇χ."""
    poisson = SphericalPoissonSolver(grid=self.grid)
    chi = poisson.solve(delta, zero_mean=True, spectral=spectral)
    deriv = SphericalDerivative2D(grid=self.grid)
    grad_theta_chi, grad_phi_chi = deriv.gradient(chi)
    return chi, (grad_theta_chi, grad_phi_chi)

SphericalHelmholtzDecomposition

Bases: Module

Helmholtz decomposition of a horizontal vector field on the sphere.

Given a tangent field V = (V_θ, V_φ) on the sphere, decomposes

V = ∇χ + ẑ × ∇ψ

into a curl-free part (velocity potential χ) and a divergence-free part (streamfunction ψ). The scalar potentials are obtained by inverting the horizontal Laplacian applied to the divergence and vorticity of V:

∇²χ = ∇·V = δ
∇²ψ = (∇×V)·r̂ = ζ

This is the spherical analogue of the classical Helmholtz decomposition for 2D incompressible/irrotational flow, and is the foundation of vorticity–divergence spectral GFD models.

Accuracy note

The intermediate divergence and curl rely on :class:SphericalDerivative2D, whose colatitude derivative uses a 1D Legendre transform column-by-column. That is exact only for zonal (m = 0) modes; for m ≠ 0 modes there is a small truncation error proportional to the resolved smoothness of V. The Laplace–Beltrami inversion itself is spectrally exact. Pipelines that require machine-precision round-trip reconstruction of arbitrary fields should first project V onto its vorticity / divergence in SHT space (future work).

Attributes:

Name Type Description
grid SphericalGrid2D

Underlying 2D lat-lon grid.

Examples:

>>> import jax.numpy as jnp
>>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
>>> decomp = SphericalHelmholtzDecomposition(grid=grid)
>>> v_theta = jnp.zeros((grid.Ny, grid.Nx))
>>> v_phi = jnp.sin(grid.y)[:, None] * jnp.ones((grid.Ny, grid.Nx))
>>> psi, chi, v_rot, v_div = decomp.decompose(v_theta, v_phi)
Source code in spectraldiffx/_src/spherical/solvers.py 
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
class SphericalHelmholtzDecomposition(eqx.Module):
    """Helmholtz decomposition of a horizontal vector field on the sphere.

    Given a tangent field V = (V_θ, V_φ) on the sphere, decomposes

        V = ∇χ + ẑ × ∇ψ

    into a curl-free part (velocity potential χ) and a divergence-free
    part (streamfunction ψ).  The scalar potentials are obtained by
    inverting the horizontal Laplacian applied to the divergence and
    vorticity of V:

        ∇²χ = ∇·V = δ
        ∇²ψ = (∇×V)·r̂ = ζ

    This is the spherical analogue of the classical Helmholtz
    decomposition for 2D incompressible/irrotational flow, and is the
    foundation of vorticity–divergence spectral GFD models.

    Accuracy note
    -------------
    The intermediate divergence and curl rely on
    :class:`SphericalDerivative2D`, whose colatitude derivative uses a
    1D Legendre transform column-by-column.  That is exact only for
    zonal (m = 0) modes; for m ≠ 0 modes there is a small truncation
    error proportional to the resolved smoothness of V.  The
    Laplace–Beltrami inversion itself is spectrally exact.  Pipelines
    that require machine-precision round-trip reconstruction of
    arbitrary fields should first project V onto its vorticity /
    divergence in SHT space (future work).

    Attributes
    ----------
    grid : SphericalGrid2D
        Underlying 2D lat-lon grid.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> grid = SphericalGrid2D.from_N_L(Nx=32, Ny=16)
    >>> decomp = SphericalHelmholtzDecomposition(grid=grid)
    >>> v_theta = jnp.zeros((grid.Ny, grid.Nx))
    >>> v_phi = jnp.sin(grid.y)[:, None] * jnp.ones((grid.Ny, grid.Nx))
    >>> psi, chi, v_rot, v_div = decomp.decompose(v_theta, v_phi)
    """

    grid: SphericalGrid2D

    def decompose(
        self,
        v_theta: Num[Array, "Nlat Nlon"],
        v_phi: Num[Array, "Nlat Nlon"],
    ) -> tuple[
        Float[Array, "Nlat Nlon"],
        Float[Array, "Nlat Nlon"],
        tuple[Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"]],
        tuple[Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"]],
    ]:
        """Compute (ψ, χ, V_rot, V_div) for a tangent vector field.

        Parameters
        ----------
        v_theta : Num[Array, "Nlat Nlon"]
            Colatitude component of V (physical space).
        v_phi : Num[Array, "Nlat Nlon"]
            Longitude component of V (physical space).

        Returns
        -------
        psi : Float[Array, "Nlat Nlon"]
            Streamfunction (divergence-free potential).
        chi : Float[Array, "Nlat Nlon"]
            Velocity potential (curl-free potential).
        v_rot : (Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"])
            Divergence-free velocity components (V_θ, V_φ) = ẑ × ∇ψ.
        v_div : (Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"])
            Curl-free velocity components (V_θ, V_φ) = ∇χ.
        """
        deriv = SphericalDerivative2D(grid=self.grid)
        zeta = deriv.curl(v_theta, v_phi)
        delta = deriv.divergence(v_theta, v_phi)

        rot = SphericalVorticityInversionSolver(grid=self.grid)
        div = SphericalDivergenceInversionSolver(grid=self.grid)
        psi, v_rot = rot.solve(zeta)
        chi, v_div = div.solve(delta)
        return psi, chi, v_rot, v_div

Functions

decompose(v_theta, v_phi)

Compute (ψ, χ, V_rot, V_div) for a tangent vector field.

Parameters:

Name Type Description Default
v_theta Num[Array, 'Nlat Nlon']

Colatitude component of V (physical space).

 required
v_phi Num[Array, 'Nlat Nlon']

Longitude component of V (physical space).

 required

Returns:

Name Type Description
psi Float[Array, 'Nlat Nlon']

Streamfunction (divergence-free potential).
chi Float[Array, 'Nlat Nlon']

Velocity potential (curl-free potential).
v_rot (Float[Array, 'Nlat Nlon'], Float[Array, 'Nlat Nlon'])

Divergence-free velocity components (V_θ, V_φ) = ẑ × ∇ψ.
v_div (Float[Array, 'Nlat Nlon'], Float[Array, 'Nlat Nlon'])

Curl-free velocity components (V_θ, V_φ) = ∇χ.
Source code in spectraldiffx/_src/spherical/solvers.py 
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def decompose(
    self,
    v_theta: Num[Array, "Nlat Nlon"],
    v_phi: Num[Array, "Nlat Nlon"],
) -> tuple[
    Float[Array, "Nlat Nlon"],
    Float[Array, "Nlat Nlon"],
    tuple[Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"]],
    tuple[Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"]],
]:
    """Compute (ψ, χ, V_rot, V_div) for a tangent vector field.

    Parameters
    ----------
    v_theta : Num[Array, "Nlat Nlon"]
        Colatitude component of V (physical space).
    v_phi : Num[Array, "Nlat Nlon"]
        Longitude component of V (physical space).

    Returns
    -------
    psi : Float[Array, "Nlat Nlon"]
        Streamfunction (divergence-free potential).
    chi : Float[Array, "Nlat Nlon"]
        Velocity potential (curl-free potential).
    v_rot : (Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"])
        Divergence-free velocity components (V_θ, V_φ) = ẑ × ∇ψ.
    v_div : (Float[Array, "Nlat Nlon"], Float[Array, "Nlat Nlon"])
        Curl-free velocity components (V_θ, V_φ) = ∇χ.
    """
    deriv = SphericalDerivative2D(grid=self.grid)
    zeta = deriv.curl(v_theta, v_phi)
    delta = deriv.divergence(v_theta, v_phi)

    rot = SphericalVorticityInversionSolver(grid=self.grid)
    div = SphericalDivergenceInversionSolver(grid=self.grid)
    psi, v_rot = rot.solve(zeta)
    chi, v_div = div.solve(delta)
    return psi, chi, v_rot, v_div