# A research framework for writing differentiable PDE discretizations in JAX

arXiv: https://arxiv.org/abs/2111.05218 · HTML: https://mlsanity.com/paper/2111.05218/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 24 references that are in the map.
- First posted on arXiv 9 Nov 2021.

## Builds on (most cited first, 14 of 14)

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [DeepXDE: A deep learning library for solving differential equations](https://mlsanity.com/paper/1907.04502/) (2019)
- [The frontier of simulation-based inference](https://mlsanity.com/paper/1911.01429/) (2019)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Universal Differential Equations for Scientific Machine Learning](https://mlsanity.com/paper/2001.04385/) (2020)
- [PDE-Net: Learning PDEs from Data](https://mlsanity.com/paper/1710.09668/) (2017)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Solving ill-posed inverse problems using iterative deep neural networks](https://mlsanity.com/paper/1704.04058/) (2017)
- [Learning to Control PDEs with Differentiable Physics](https://mlsanity.com/paper/2001.07457/) (2020)
- [gradSim: Differentiable simulation for system identification and visuomotor control](https://mlsanity.com/paper/2104.02646/) (2021)
- [A Differentiable Programming System to Bridge Machine Learning and Scientific Computing](https://mlsanity.com/paper/1907.07587/) (2019)
- [NeuralSim: Augmenting Differentiable Simulators with Neural Networks](https://mlsanity.com/paper/2011.04217/) (2020)
- [DiffEqFlux.jl - A Julia Library for Neural Differential Equations](https://mlsanity.com/paper/1902.02376/) (2019)
- [The Koopman Expectation: An Operator Theoretic Method for Efficient Analysis and Optimization of Uncertain Hybrid Dynamical Systems](https://mlsanity.com/paper/2008.08737/) (2020)

## Built on by (most cited first, 2 of 2)

- [On conditional diffusion models for PDE simulations](https://mlsanity.com/paper/2410.16415/) (2024)
- [j-Wave: An open-source differentiable wave simulator](https://mlsanity.com/paper/2207.01499/) (2022)

## Abstract (arXiv)

Differentiable simulators are an emerging concept with applications in several fields, from reinforcement learning to optimal control. Their distinguishing feature is the ability to calculate analytic gradients with respect to the input parameters. Like neural networks, which are constructed by composing several building blocks called layers, a simulation often requires computing the output of an operator that can itself be decomposed into elementary units chained together. While each layer of a neural network represents a specific discrete operation, the same operator can have multiple representations, depending on the discretization employed and the research question that needs to be addressed. Here, we propose a simple design pattern to construct a library of differentiable operators and discretizations, by representing operators as mappings between families of continuous functions, parametrized by finite vectors. We demonstrate the approach on an acoustic optimization problem, where the Helmholtz equation is discretized using Fourier spectral methods, and differentiability is demonstrated using gradient descent to optimize the speed of sound of an acoustic lens. The proposed framework is open-sourced and available at urlhttps://github.com/ucl-bug/jaxdf