Skip to main content

Ctrl+K

Differentiable Earth System Modeling with JAX

Differentiable Earth System Modeling with JAX

Part 0 — Why differentiable modeling?

0.1 The adjoint for free
0.2 Science questions, reframed as gradients

Part 1 — JAX fundamentals

1.1 Coding in JAX
1.2 First climate gradient
1.3 Add time: the energy balance model

Part 2 — Gradient-based science on a small model

2.1 Sensitivity maps
2.2 Calibration as optimization
2.3 Variational data assimilation
2.4 Hybrid physics–ML modeling

Part 3 — The hard parts

3.1 Gradients through chaos
3.2 Memory and the cost of long simulations
3.3 Non-smooth physics
3.4 Differentiating through solvers
3.5 Do you trust your gradient?

Appendices

Setup
JAX gotchas cheat sheet
References

Repository
Open issue

Index

By Differentiable Earth System Science Research Group (LEAP)

© Copyright 2026.