1.2 First climate gradient

1.2 First climate gradient#

The simplest climate model is an algebraic statement of energy balance: sunlight absorbed equals infrared radiation emitted. For a planet with albedo (reflectivity) \(\alpha\) receiving solar irradiance \(S_0\) and emitting like a grey body with emissivity \(\varepsilon\),

\[\frac{S_0}{4}(1-\alpha) \;=\; \varepsilon \sigma T^4 \quad\Longrightarrow\quad T_\mathrm{eq} = \left[\frac{S_0 (1-\alpha)}{4 \varepsilon \sigma}\right]^{1/4}.\]

Using jax.grad take the derivative of this model with respect to albedo.

Outline

  • Compute the derivative of a model output with jax.grad and interpret it physically

  • Verify a gradient against finite differences

  • Use jax.jvp and jax.vjp, and connect them to the tangent linear and adjoint models

  • Assess the cost: reverse mode prices all input sensitivities at a few model runs

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

jax.config.update("jax_enable_x64", True)   # climate work needs double precision
# Running on Colab? First run:  %pip install -q optax equinox diffrax

The model, and its gradient#

S0 = 1361.0       # solar constant, W m^-2
sigma = 5.67e-8   # Stefan-Boltzmann constant, W m^-2 K^-4

def equilibrium_temperature(albedo, emissivity=0.61):
    """Global-mean surface temperature (K) in radiative equilibrium."""
    return (S0 * (1.0 - albedo) / (4.0 * emissivity * sigma)) ** 0.25

T = equilibrium_temperature(0.30)
dT_dalbedo = jax.grad(equilibrium_temperature)(0.30)
print(f"T_eq = {T:.1f} K")
print(f"dT/d(albedo) = {dT_dalbedo:.1f} K per unit albedo")

T_eq = 288.1 K
dT/d(albedo) = -102.9 K per unit albedo

With Earth-like numbers (\(\alpha = 0.30\), \(\varepsilon = 0.61\)) the model gives a realistic 288 K. The gradient says that increasing planetary albedo by 0.01 — roughly the difference made by a large volcanic eruption — cools the planet by about 1 K. That is a sensitivity, obtained without running any perturbation experiment.

jax.grad differentiates by the chain rule, applied mechanically to each elementary operation in the function — the same mathematics you would use by hand, executed by the computer on the code as written. It is not a numerical approximation. For this function the derivative also exists in closed form, \(\partial T_\mathrm{eq}/\partial\alpha = -T_\mathrm{eq}\,/\,[4(1-\alpha)]\), so we can check that autodiff is exact:

analytic = -T / (4.0 * (1.0 - 0.30))
print(f"autodiff: {dT_dalbedo:.10f}")
print(f"analytic: {analytic:.10f}")

autodiff: -102.8816186994
analytic: -102.8816186994

Always check against finite differences#

Real models have no closed-form derivative, so the standard independent check is a finite difference: perturb the input by a small step \(h\) and difference the outputs, \([f(x+h) - f(x-h)]\,/\,2h\). This estimate carries its own errors — too large an \(h\) misses curvature, too small an \(h\) amplifies floating-point roundoff — which is exactly why autodiff is preferable for the science, and why the finite difference is best used as a check rather than as the method:

def finite_difference(f, x, h):
    return (f(x + h) - f(x - h)) / (2.0 * h)

print(" h          FD estimate      |FD - autodiff|")
for h in [1e-2, 1e-4, 1e-6, 1e-8, 1e-10]:
    fd = finite_difference(equilibrium_temperature, 0.30, h)
    print(f"{h:8.0e}   {fd:14.8f}   {abs(fd - dT_dalbedo):.2e}")

 h          FD estimate      |FD - autodiff|
   1e-02    -102.88621211   4.59e-03
   1e-04    -102.88161916   4.59e-07
   1e-06    -102.88161869   1.19e-08
   1e-08    -102.88162002   1.32e-06
   1e-10    -102.88147223   1.46e-04

The error shrinks as \(h\) decreases, then grows again as roundoff takes over — with the best step giving 8–10 correct digits. Autodiff has no step size and no such trade-off. We will run a check like this on every gradient in this book; Chapter 3.5 turns it into a routine.

Two modes: jvp and vjp#

grad is a convenience wrapper around two more fundamental operations.

Forward mode (jax.jvp, “Jacobian-vector product”) pushes an input perturbation through the model and returns the resulting output perturbation — the tangent linear model. Reverse mode (jax.vjp, “vector-Jacobian product”) propagates a sensitivity backwards from output to inputs — this is the adjoint model. Chapter 0.1 describes the correspondence; here is what the calls look like:

# forward mode: effect on T of an albedo perturbation of +0.01
_, dT = jax.jvp(equilibrium_temperature, (0.30,), (0.01,))
print(f"tangent linear: dT = {dT:.3f} K for d(albedo) = +0.01")

# reverse mode: sensitivity of T to the input, from one backward pass
T_value, backward = jax.vjp(equilibrium_temperature, 0.30)
(sens,) = backward(1.0)          # propagate dT = 1 back to the input
print(f"adjoint: dT/d(albedo) = {sens:.1f} K")

tangent linear: dT = -1.029 K for d(albedo) = +0.01
adjoint: dT/d(albedo) = -102.9 K

The cost#

The two modes differ in what they charge for. Forward mode costs one model evaluation per input direction; reverse mode costs a few model evaluations total, regardless of how many inputs there are, because one backward pass yields the sensitivity to everything at once.

The demonstration below uses a function of 1,000 inputs. Computing all 1,000 sensitivities takes forward mode 1,000 differentiation passes and reverse mode one. (The measured gap is smaller than 1,000× because forward mode evaluates its passes in one batched sweep, but the gap widens without bound as inputs are added, while reverse mode’s cost stays flat.)

import time

key = jax.random.key(0)
M = jax.random.normal(key, (1000, 1000)) / 1000.0

def scalar_diagnostic(v):                     # many inputs -> one output
    return jnp.sum(jnp.tanh(M @ v) ** 2)

v = jnp.ones(1000)
rev = jax.jit(jax.grad(scalar_diagnostic))    # reverse mode: one backward pass
fwd = jax.jit(jax.jacfwd(scalar_diagnostic))  # forward mode: one pass per input

rev(v).block_until_ready(); fwd(v).block_until_ready()   # compile both
t = time.perf_counter(); rev(v).block_until_ready()
t_rev = time.perf_counter() - t
t = time.perf_counter(); fwd(v).block_until_ready()
t_fwd = time.perf_counter() - t
print(f"reverse mode: {t_rev*1000:7.1f} ms")
print(f"forward mode: {t_fwd*1000:7.1f} ms   ({t_fwd/t_rev:.0f}x slower for 1000 inputs)")
print(f"same result: {bool(jnp.allclose(rev(v), fwd(v)))}")

reverse mode:     0.6 ms
forward mode:     8.8 ms   (16x slower for 1000 inputs)
same result: True

The rule of thumb:

You want

Use

Cost

Effect of one perturbation on many outputs

forward (jvp)

~1 model run

Sensitivity of one diagnostic to many inputs

reverse (grad, vjp)

~a few model runs

Full Jacobian, few inputs

jax.jacfwd

~1 run per input

Full Jacobian, few outputs

jax.jacrev

~1 run per output

Reverse mode’s constant price is what makes adjoint sensitivity maps (Chapter 2.1), variational data assimilation (2.3), and neural-network training (2.4) affordable. Its hidden cost — memory — is the subject of Chapter 3.2.

Exercises.

  1. Compute \(\partial T_\mathrm{eq}/\partial \varepsilon\) and check it against a finite difference. Is the sign what you expect physically?

  2. The second derivative is one line: jax.grad(jax.grad(equilibrium_temperature)). Evaluate it and interpret the sign.

  3. Use jax.jacfwd to get \(\partial T/\partial(\alpha, \varepsilon)\) jointly, by rewriting the function to take a single array of both parameters.