https://github.com/pmocz/advectiondiffusion-jax

Solve the advection diffusion equations looped into an optimization problem with JAX/autodiff
https://github.com/pmocz/advectiondiffusion-jax

automatic-differentiation jax scientific-computing tutorial

Last synced: 11 months ago
JSON representation

Solve the advection diffusion equations looped into an optimization problem with JAX/autodiff

• Host: GitHub
• URL: https://github.com/pmocz/advectiondiffusion-jax
• Owner: pmocz
• License: gpl-3.0
• Created: 2024-06-04T15:15:58.000Z (over 1 year ago)
• Default Branch: main
• Last Pushed: 2024-06-07T18:55:01.000Z (over 1 year ago)
• Last Synced: 2024-06-07T21:47:46.915Z (over 1 year ago)
• Topics: automatic-differentiation, jax, scientific-computing, tutorial
• Language: Python
• Homepage:
• Size: 367 KB
• Stars: 0
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

          
# advectiondiffusion-jax

Solve the advection diffusion equations looped into an optimization problem with JAX/autodiff

## Create Your Own Automatically Differentiable Simulation (With Python/JAX)

### Philip Mocz (2024) Flatiron Institute, [@PMocz](https://twitter.com/PMocz)

### [📝 Read the Algorithm Write-up on Medium](https://philip-mocz.medium.com/create-your-own-automatically-differentiable-simulation-with-python-jax-46951e120fbb)

Simulate the advection-diffusion equations with finite difference and calculate gradients automatically with JAX,

in order to find which simulation input parameters lead to 'worst-case' pollution scenario.

Run the code in the command line with:

```console

python advection_diffusion.py

```

Simulation result:

![Simulation](./simulation.png)

'Worst-case' pollution scenario:

![Pollution](./pollution.png)