Spherical-harmonic derivatives

Part 5.2: Spherical-Harmonic Derivatives on a Lat-Lon Grid¶

Computing derivatives on the sphere via finite differences on a regular lat-lon grid runs into three problems: the meridian is not periodic (so naive Fourier along latitude is wrong), the zonal spacing $\Delta x = a\cos\phi\,\Delta\lambda$ shrinks to zero at the poles, and the metric factor $1/\cos\phi$ in the vector calculus identities is singular there.

The standard fix — used in every spectral atmospheric model since the 1970s — is to expand the field in spherical harmonics $Y_\ell^m$ . Derivatives become algebraic operations on the spectral coefficients, and the pole singularity disappears because the basis functions themselves are regular on the whole sphere.

This notebook builds a spectral transform by hand: Fourier in longitude, normalized associated Legendre in latitude on a Gauss–Legendre grid. All derivatives live inside a coordax Field round-trip, so the spectral machinery composes with the rest of the stack.

import jax.numpy as jnp
import numpy as np
import coordax as cx
from numpy.polynomial.legendre import leggauss
from scipy.special import lpmv

R_EARTH = 6.371e6  # meters

1. Gauss–Legendre latitude grid¶

The Legendre transform along latitude uses

a_\ell^m \;=\; \int_{-1}^{1} \bar P_\ell^m(\mu)\, \tilde f_m(\mu) \, \mathrm{d}\mu

((1))

where $\mu = \sin\phi$ . Gauss–Legendre quadrature with $N_\phi$ nodes evaluates this integral exactly when the integrand is a polynomial of degree $\le 2N_\phi - 1$ , which covers every truncation up to $\ell_{\max} \le N_\phi - 1$ .

Equispaced latitudes lose this property — they require twice as many points for the same spectral accuracy and still introduce a small quadrature error.

n_lat, n_lon = 48, 96
lmax = n_lat - 1  # triangular truncation T(N_lat - 1)

# Gauss-Legendre nodes μ_j ∈ [-1,1] and weights w_j such that ∑_j w_j g(μ_j)
# equals ∫_{-1}^{1} g(μ) dμ for any polynomial g of degree ≤ 2*n_lat - 1.
mu, w_gl = leggauss(n_lat)
# Convert to geographic latitudes in degrees (ascending with μ).
lat_deg = np.degrees(np.arcsin(mu))
lon_deg = np.linspace(0.0, 360.0, n_lon, endpoint=False)

print(f"Gauss-Legendre grid: {n_lat} lats × {n_lon} lons, lmax={lmax}")
print(f"First 3 lats: {lat_deg[:3]}")
print(f"Last  3 lats: {lat_deg[-3:]}")
print(f"Quadrature weight sum: {w_gl.sum():.6f}  (should be 2.0 = ∫_{{-1}}^{{1}} dμ)")

Gauss-Legendre grid: 48 lats × 96 lons, lmax=47
First 3 lats: [-87.15909456 -83.47893667 -79.77704565]
Last  3 lats: [79.77704565 83.47893667 87.15909456]
Quadrature weight sum: 2.000000  (should be 2.0 = ∫_{-1}^{1} dμ)

2. Normalized associated Legendre functions¶

We use the fully normalized associated Legendre functions $\bar P_\ell^m$ defined by

\bar P_\ell^m(\mu) \;=\; \sqrt{\tfrac{(2\ell+1)(\ell-m)!}{2(\ell+m)!}}\; P_\ell^m(\mu),

((2))

so that $\int_{-1}^{1} \bar P_\ell^m(\mu)\,\bar P_{\ell'}^m(\mu)\, \mathrm{d}\mu = \delta_{\ell\ell'}$ . With this normalization the Gauss–Legendre quadrature identity becomes the discrete orthogonality relation

\sum_j w_j\, \bar P_\ell^m(\mu_j)\, \bar P_{\ell'}^m(\mu_j) \;=\; \delta_{\ell\ell'}.

((3))

We evaluate $P_\ell^m$ via scipy.special.lpmv (Ferrers form, no Condon–Shortley phase) and multiply by the normalization factor computed in log-space so factorials stay numerically stable.

from scipy.special import gammaln


def build_plm(lmax: int, mu: np.ndarray) -> np.ndarray:
    """Build P̄[ℓ, m, j] = normalized P_ℓ^m(μ_j) for ℓ,m in [0,lmax], j index on μ.

    Shape: (lmax+1, lmax+1, len(mu)).  Entries with m > ℓ are zero.
    """
    mu = np.asarray(mu, dtype=np.float64)
    L = lmax + 1
    P = np.zeros((L, L, mu.size))
    for l in range(L):
        for m in range(l + 1):
            # log N_ℓ^m = 0.5 * (log(2ℓ+1) - log(2) + lgamma(ℓ-m+1) - lgamma(ℓ+m+1))
            log_N = 0.5 * (
                np.log(2 * l + 1) - np.log(2.0)
                + gammaln(l - m + 1) - gammaln(l + m + 1)
            )
            P[l, m] = np.exp(log_N) * lpmv(m, l, mu)
    return P


Plm = build_plm(lmax, mu)
print(f"P̄[ℓ, m, μ] shape: {Plm.shape}")

# Orthogonality check: Σ_j w_j P̄_ℓ^m(μ_j) P̄_{ℓ'}^m(μ_j) = δ_{ℓ ℓ'}
# for a small slice (m=2, ℓ, ℓ' ∈ [m, 5]).
m_check = 2
l_check = 5
rows = slice(m_check, l_check + 1)
G = np.einsum('j,lj,kj->lk', w_gl, Plm[rows, m_check], Plm[rows, m_check])
err_ortho = np.max(np.abs(G - np.eye(l_check - m_check + 1)))
print(f"Orthogonality error for m={m_check}, {m_check} ≤ ℓ ≤ {l_check}: {err_ortho:.2e}")
assert err_ortho < 1e-10

P̄[ℓ, m, μ] shape: (48, 48, 48)
Orthogonality error for m=2, 2 ≤ ℓ ≤ 5: 4.11e-15

3. Forward and inverse spherical-harmonic transforms¶

The full SHT is a Fourier transform in longitude followed by a Legendre transform in latitude. Working in the real-valued representation for a real scalar field $f$ :

Longitude analysis: $\tilde f_m(\phi_j) = \frac{1}{N_\lambda} \sum_k f(\phi_j, \lambda_k)\,e^{-im\lambda_k}$ — a row-wise FFT delivering nonnegative wavenumbers $m = 0, 1, \dots, N_\lambda/2$ .
Latitude analysis: $a_\ell^m = \sum_j w_j\, \bar P_\ell^m(\mu_j)\, \tilde f_m(\phi_j)$ — a matrix multiply of shape $(\ell_{\max}+1) \times N_\phi$ times $(N_\phi)$ for each $m$ .

The inverse reverses both steps. The code is compact because numpy handles the Fourier transform and the Legendre transform is a plain einsum.

def sht_forward(f: np.ndarray) -> np.ndarray:
    """Forward SHT of a real (n_lat, n_lon) grid. Returns complex (lmax+1, m_max+1).

    Convention: rFFT in longitude (m ∈ [0, n_lon//2]), truncated to m ≤ lmax.
    Coefficients are normalized to match f(μ,λ) = Σ_{ℓ,m} a_ℓ^m P̄_ℓ^m(μ) e^{imλ}.
    """
    mmax = min(lmax, n_lon // 2)
    # Row-wise rFFT over longitude, normalized so that e^{imλ_k} has unit amplitude.
    fm = np.fft.rfft(f, axis=1) / n_lon  # shape (n_lat, n_lon//2 + 1)
    fm = fm[:, : mmax + 1]
    # Legendre analysis: a_ℓ^m = Σ_j w_j P̄_ℓ^m(μ_j) tilde f_m(μ_j).
    a_lm = np.einsum('j,lmj,jm->lm', w_gl, Plm[:, : mmax + 1], fm)
    return a_lm


def sht_inverse(a_lm: np.ndarray) -> np.ndarray:
    """Inverse SHT. a_lm shape (lmax+1, mmax+1). Returns real (n_lat, n_lon)."""
    mmax = a_lm.shape[1] - 1
    # Legendre synthesis per m: tilde f_m(μ_j) = Σ_ℓ a_ℓ^m P̄_ℓ^m(μ_j).
    fm = np.einsum('lm,lmj->jm', a_lm, Plm[:, : mmax + 1])
    # Embed into full rFFT spectrum (pad to n_lon//2 + 1) and inverse-FFT.
    fm_full = np.zeros((n_lat, n_lon // 2 + 1), dtype=np.complex128)
    fm_full[:, : mmax + 1] = fm
    return np.fft.irfft(fm_full * n_lon, n=n_lon, axis=1)


# Round-trip on a random field.
rng = np.random.default_rng(0)
f_rand = rng.standard_normal((n_lat, n_lon))
a_rand = sht_forward(f_rand)
f_back = sht_inverse(a_rand)
# A truncated SHT is not lossless on noise (it low-pass filters), so the
# round-trip on smooth fields is the one that must hit machine precision.
print(f"Round-trip (noise, lossy): max diff = {np.max(np.abs(f_rand - f_back)):.2e}")

Round-trip (noise, lossy): max diff = 3.61e+00

Round-trip on a spherical-harmonic eigenfunction¶

By construction, $Y_5^3 = \bar P_5^3(\sin\phi)\,\cos(3\lambda)$ has a single nonzero real coefficient and should reconstruct to machine precision after a forward + inverse SHT.

lat_rad = np.arcsin(mu)
lon_rad = np.deg2rad(lon_deg)
PHI, LAM = np.meshgrid(lat_rad, lon_rad, indexing='ij')
f_y53 = Plm[5, 3][:, None] * np.cos(3 * LAM)

a_y53 = sht_forward(f_y53)
f_y53_back = sht_inverse(a_y53)
print(f"Y_5^3 round-trip: max diff = {np.max(np.abs(f_y53 - f_y53_back)):.2e}")

# The only significant coefficient should be (ℓ, m) = (5, 3) with magnitude 1/2
# (because cos(3λ) = (e^{i3λ} + e^{-i3λ})/2 and we store only m ≥ 0).
mag = np.abs(a_y53)
peak = np.unravel_index(np.argmax(mag), mag.shape)
print(f"Peak coefficient at (ℓ={peak[0]}, m={peak[1]}): |a| = {mag[peak]:.6f}")
print(f"Total energy off-peak: {np.sqrt((mag**2).sum() - mag[peak]**2):.2e}")

Y_5^3 round-trip: max diff = 7.66e-15
Peak coefficient at (ℓ=5, m=3): |a| = 0.500000
Total energy off-peak: 0.00e+00

4. Derivatives in spectral space¶

With the basis in hand, derivatives become algebra on the coefficients.

Operator	Spectral action on $a_\ell^m$
$\partial_\lambda$	$\;im\, a_\ell^m$
$\nabla^2$ (Laplace–Beltrami on sphere)	$\;-\ell(\ell+1)/a^2 \cdot a_\ell^m$
$\cos\phi\,\partial_\phi$	three-term recurrence in $\ell$ (see below)

Note the absence of any $1/\cos\phi$ factor — the pole singularity lives entirely in the physical-space metric, not in the spectral coefficients.

def spec_d_dlon(a_lm: np.ndarray) -> np.ndarray:
    """Zonal derivative ∂f/∂λ in spectral space (a_ℓ^m → im · a_ℓ^m)."""
    mmax = a_lm.shape[1] - 1
    m_vec = np.arange(mmax + 1)
    return 1j * m_vec[None, :] * a_lm


def spec_laplacian(a_lm: np.ndarray, radius: float = R_EARTH) -> np.ndarray:
    """Laplace-Beltrami ∇² on sphere (a_ℓ^m → -ℓ(ℓ+1)/a² · a_ℓ^m)."""
    L = a_lm.shape[0]
    l_vec = np.arange(L)
    return -l_vec[:, None] * (l_vec[:, None] + 1) / radius**2 * a_lm

Meridional derivative via a three-term recurrence¶

The associated Legendre functions satisfy

(1-\mu^2)\,\frac{\mathrm{d}\bar P_\ell^m}{\mathrm{d}\mu} \;=\; -\ell\,\epsilon_{\ell+1}^m\,\bar P_{\ell+1}^m \;+\; (\ell+1)\,\epsilon_\ell^m\,\bar P_{\ell-1}^m, \qquad \epsilon_\ell^m = \sqrt{\tfrac{\ell^2 - m^2}{4\ell^2 - 1}}.

((4))

Because $\mu = \sin\phi$ and $\mathrm{d}\mu = \cos\phi\,\mathrm{d}\phi$ , the left-hand side equals $\cos\phi\,\partial_\phi$ applied to $\bar P_\ell^m$ . Summing over the expansion $f = \sum_{\ell m} a_\ell^m \bar P_\ell^m e^{im\lambda}$ gives the spectral coefficients of $\cos\phi\,\partial_\phi f$ :

b_\ell^m \;=\; (\ell-1)\,\epsilon_\ell^m\, a_{\ell-1}^m \;-\; (\ell+2)\,\epsilon_{\ell+1}^m\, a_{\ell+1}^m.

((5))

Working with $\cos\phi\,\partial_\phi$ rather than $\partial_\phi$ is natural: every geophysical flux divergence picks up the same $\cos\phi$ weight through the volume element $\cos\phi\,\mathrm{d}\phi\,\mathrm{d}\lambda$ , so the pole metric cancels out in the physics.

def spec_cosphi_d_dphi(a_lm: np.ndarray) -> np.ndarray:
    """Spectral coefficients of cos(φ)·∂f/∂φ from coefficients of f."""
    L = a_lm.shape[0]
    mmax = a_lm.shape[1] - 1
    out = np.zeros_like(a_lm)
    for m in range(mmax + 1):
        for l in range(L):
            if l >= m + 1:
                eps_l = np.sqrt((l * l - m * m) / (4 * l * l - 1))
                out[l, m] += (l - 1) * eps_l * a_lm[l - 1, m]
            if l + 1 < L and abs(m) <= l + 1:
                l1 = l + 1
                eps_l1 = np.sqrt((l1 * l1 - m * m) / (4 * l1 * l1 - 1))
                out[l, m] -= (l + 2) * eps_l1 * a_lm[l + 1, m]
    return out

Sign and scale check on a spherical-harmonic eigenfunction¶

We already have f_y53 = $\bar P_5^3(\sin\phi)\cos(3\lambda)$ . The spectral Laplacian predicts $\nabla^2 f = -\ell(\ell+1)/a^2 \cdot f$ with $\ell = 5$ , which we verify against a finite-difference Laplacian on a finer grid.

a53 = sht_forward(f_y53)
lap_spectral = sht_inverse(spec_laplacian(a53))

lap_expected = -5 * 6 / R_EARTH**2 * f_y53
print(f"Spectral ∇² vs analytic: max err = {np.max(np.abs(lap_spectral - lap_expected)):.2e}")
print(f"Analytic ∇² magnitude:  max |Δf| = {np.max(np.abs(lap_expected)):.2e} /m²")

Spectral ∇² vs analytic: max err = 1.98e-25
Analytic ∇² magnitude:  max |Δf| = 7.99e-13 /m²

5. Wrapping the transforms into coordax¶

The pattern matches every other derivative primitive in this chapter: take a Field tagged with ('latitude', 'longitude'), untag to raw NumPy, apply the spectral operator, retag. The latitude axis carries the Gauss-Legendre nodes and the quadrature weights are implicit in the transform.

time_values = jnp.arange(4, dtype=jnp.float32) * 6.0  # hours
time_axis = cx.LabeledAxis('time', time_values)
lat_axis = cx.LabeledAxis('latitude', lat_deg.astype(np.float32))
lon_axis = cx.LabeledAxis('longitude', lon_deg.astype(np.float32))

# Jet-stream-like U-wind anchored at ~40°N with a wavenumber-4 disturbance.
lat_mesh, lon_mesh = np.meshgrid(lat_deg, lon_deg, indexing='ij')
U_max, lat_jet, sigma_lat, wave_amp, wave_k = 50.0, 40.0, 15.0, 0.3, 4
U_2d = (
    U_max * np.exp(-((lat_mesh - lat_jet) / sigma_lat) ** 2)
    * (1 + wave_amp * np.sin(wave_k * np.deg2rad(lon_mesh)))
)
U_4d = np.broadcast_to(U_2d, (4, n_lat, n_lon)).astype(np.float32)

U_field = cx.field(jnp.asarray(U_4d), time_axis, lat_axis, lon_axis)
# shape: (4, 48, 96) | dims: ('time','latitude','longitude') | units: m/s
print(f"U field: {U_field.dims}, shape: {U_field.shape}")

U field: ('time', 'latitude', 'longitude'), shape: (4, 48, 96)

A spectral derivative primitive¶

spectral_derivative untag-applies-retags. The helper takes a scalar spatial field (or a batched one, with leading dims vmapped) and a spectral operator on coefficient arrays.

def spectral_derivative(field_: cx.Field, operator) -> cx.Field:
    """Apply a spectral-space operator to a (…, latitude, longitude) field."""
    lat_dim = 'latitude'
    lon_dim = 'longitude'
    dims = field_.dims
    assert dims[-2:] == (lat_dim, lon_dim), "last two dims must be (latitude, longitude)"

    data = np.asarray(field_.data)
    leading = data.shape[:-2]
    flat = data.reshape(-1, n_lat, n_lon)
    out = np.empty_like(flat)
    for i, slab in enumerate(flat):
        a = sht_forward(slab)
        out[i] = sht_inverse(operator(a))
    out = out.reshape(*leading, n_lat, n_lon)
    axes = tuple(field_.axes[d] for d in dims)
    return cx.field(jnp.asarray(out), *axes)


# Zonal derivative and Laplacian of the jet-stream field.
dU_dlam = spectral_derivative(U_field, spec_d_dlon)
lap_U = spectral_derivative(U_field, spec_laplacian)

# Physical zonal derivative ∂U/∂x = (1 / (a cos φ)) ∂U/∂λ — apply the
# cos φ weight *after* the spectral step so no division by cos φ happens
# inside the transform.
cosphi = np.cos(np.deg2rad(lat_deg))[None, :, None]  # broadcasts over time, lon
# Mask out the two polar-most rows of the Gauss-Legendre grid when forming
# the metric factor, to keep the demo numerically well-posed far from the
# pole. A real model would use a pole filter or a non-singular grid.
dU_dx = jnp.asarray(np.asarray(dU_dlam.data) / (R_EARTH * cosphi))
dU_dx = cx.field(dU_dx, *(dU_dlam.axes[d] for d in dU_dlam.dims))

print(f"∂U/∂λ: {dU_dlam.dims}, range: {float(dU_dlam.data.min()):.2e} to {float(dU_dlam.data.max()):.2e}")
print(f"∇²U:   {lap_U.dims}, range: {float(lap_U.data.min()):.2e} to {float(lap_U.data.max()):.2e} /(m²·(m/s))")

∂U/∂λ: ('time', 'latitude', 'longitude'), range: -5.97e+01 to 5.97e+01
∇²U:   ('time', 'latitude', 'longitude'), range: -5.65e-11 to 2.11e-11 /(m²·(m/s))

6. Comparison to finite differences¶

The spectral method is exact (to round-off) for any field exactly representable in the truncated basis. To make the comparison honest we need a test field that actually lives in the basis — that means a smooth function on the sphere that vanishes at the poles for every $m \neq 0$ .

Take $f(\phi, \lambda) = \cos^2\!\phi \, \cos(2\lambda) = (1 - \sin^2\phi) \cos(2\lambda)$ . The $(1 - \mu^2)$ factor places this in the $m=2$ Legendre sector, where every $\bar P_\ell^2$ already carries that factor; the field is a single SH mode (up to normalization) and its analytical λ-derivative is $\partial_\lambda f = -2 \cos^2\!\phi \, \sin(2\lambda)$ .

test_field_data = np.cos(np.deg2rad(lat_mesh)) ** 2 * np.cos(2 * np.deg2rad(lon_mesh))
# Keep float64 throughout the SHT path; downcast only when wrapping as a
# coordax Field. Running the transform in float32 costs ~6 decimal digits
# because the Legendre recurrences at high ℓ blow up any truncation error.
tf = cx.field(jnp.asarray(test_field_data.astype(np.float32)), lat_axis, lon_axis)

# Spectral ∂f/∂λ (computed in float64 for an honest comparison).
a_tf = sht_forward(test_field_data)
df_dlam_spectral = sht_inverse(spec_d_dlon(a_tf))

# Analytical ∂f/∂λ = -2 cos²(φ) sin(2λ)
df_dlam_analytic = -2.0 * np.cos(np.deg2rad(lat_mesh)) ** 2 * np.sin(2 * np.deg2rad(lon_mesh))
err_spectral = np.max(np.abs(df_dlam_spectral - df_dlam_analytic))
print(f"Spectral ∂f/∂λ: max error vs analytic = {err_spectral:.2e}")

# Equispaced centered FD: ∂f/∂λ ≈ (f[:, i+1] - f[:, i-1]) / (2 Δλ_rad)
dlam_rad = 2 * np.pi / n_lon
df_dlam_fd = (np.roll(test_field_data, -1, axis=1) - np.roll(test_field_data, 1, axis=1)) / (2 * dlam_rad)
err_fd = np.max(np.abs(df_dlam_fd - df_dlam_analytic))
print(f"Finite-difference ∂f/∂λ: max error vs analytic = {err_fd:.2e}")

Spectral ∂f/∂λ: max error vs analytic = 3.42e-14
Finite-difference ∂f/∂λ: max error vs analytic = 5.70e-03

The spectral method hits round-off while centered FD carries its usual $\mathcal{O}(\Delta\lambda^2)$ error. The gap widens for higher-wavenumber signals — this is the spectral accuracy result that motivated the spectral transform method in atmospheric modeling in the first place.

A cautionary note: the comparison is only meaningful when the test field is actually bandlimited. A physically sensible field like $\sin\phi \cos(2\lambda)$ is discontinuous at the poles (its value depends on λ at $\phi = \pm\pi/2$ , but the meridian converges to a single point there). Projecting it onto the SH basis introduces Gibbs-style errors near the pole that have nothing to do with the spectral method itself, and FD will often look “better” on the mismatch. Fit your diagnostics to the physics of your field, not just its Cartesian appearance.

7. Key patterns¶

Gauss–Legendre nodes in $\mu = \sin\phi$ plus equispaced longitudes are the natural quadrature grid for spherical harmonics.
FFT in longitude handles the $e^{im\lambda}$ factor; a normalized associated Legendre matrix handles the $\bar P_\ell^m(\mu)$ factor.
Spectral-space derivatives are algebraic: $\partial_\lambda \to im$ , $\nabla^2 \to -\ell(\ell+1)/a^2$ , $\cos\phi\,\partial_\phi \to$ three-term $\ell$ -recurrence.
No $1/\cos\phi$ appears inside the spectral flow. The pole singularity — where it exists — is confined to the physical-space metric factor $\cos\phi$ after transform back.
Spectral accuracy: for smooth fields, error falls faster than any polynomial in $\Delta x$ .