Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/ucl-bug/jaxdf
A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations
https://github.com/ucl-bug/jaxdf
differentiable-programming discretization jax pde
Last synced: 3 months ago
JSON representation
A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations
- Host: GitHub
- URL: https://github.com/ucl-bug/jaxdf
- Owner: ucl-bug
- License: lgpl-3.0
- Created: 2021-09-08T16:38:46.000Z (over 3 years ago)
- Default Branch: main
- Last Pushed: 2024-09-17T10:38:04.000Z (5 months ago)
- Last Synced: 2024-10-29T00:51:07.221Z (4 months ago)
- Topics: differentiable-programming, discretization, jax, pde
- Language: Python
- Homepage:
- Size: 16.9 MB
- Stars: 120
- Watchers: 7
- Forks: 7
- Open Issues: 9
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
Awesome Lists containing this project
- awesome-jax - JaxDF - Framework for differentiable simulators with arbitrary discretizations. <img src="https://img.shields.io/github/stars/ucl-bug/jaxdf?style=social" align="center"> (Libraries / New Libraries)
README
# jaxdf - JAX-based Discretization Framework
[](https://discord.gg/VtUb4fFznt)
[](https://www.gnu.org/licenses/lgpl-3.0)
[](https://codecov.io/gh/ucl-bug/jaxdf)
[](https://github.com/ucl-bug/jaxdf/actions/workflows/tests.yml)
[**Overview**](#overview)
| [**Example**](#example)
| [**Installation**](#installation)
| [**Documentation**](https://ucl-bug.github.io/jaxdf/)
| [**Support**](#support)
## Overview
Jaxdf is a package based on [JAX](https://jax.readthedocs.io/en/stable/) that provides a coding framework for creating differentiable numerical simulators with arbitrary discretizations.
The primary objective of Jaxdf is to aid in the construction of numerical models for physical systems, like wave propagation, or the numerical resolution of partial differential equations, in a manner that is easily tailored to the user's research requirements. These models are pure functions that can be seamlessly integrated into arbitrary differentiable programs written in [JAX](https://jax.readthedocs.io/en/stable/). For instance, they can be employed as layers within neural networks, or utilized in constructing a physics loss function.
## Example
The script below constructs the non-linear operator **(∇2 + sin)**, applying a Fourier spectral discretization on a square 2D domain. It then utilizes this operator to define a loss function. The gradient of this loss function is calculated using JAX's Automatic Differentiation.
```python
from jaxdf import operators as jops
from jaxdf import FourierSeries, operator
from jaxdf.geometry import Domain
from jax import numpy as jnp
from jax import jit, grad
# Defining operator
@operator
def custom_op(u, *, params=None):
grad_u = jops.gradient(u)
diag_jacobian = jops.diag_jacobian(grad_u)
laplacian = jops.sum_over_dims(diag_jacobian)
sin_u = jops.compose(u)(jnp.sin)
return laplacian + sin_u
# Defining discretizations
domain = Domain((128, 128), (1., 1.))
parameters = jnp.ones((128,128,1))
u = FourierSeries(parameters, domain)
# Define a differentiable loss function
@jit
def loss(u):
v = custom_op(u)
return jnp.mean(jnp.abs(v.on_grid)**2)
gradient = grad(loss)(u) # gradient is a FourierSeries
```
## Installation
Before proceeding with the installation of `jaxdf`, ensure that [JAX is already installed](https://github.com/google/jax#installation) on your system. If you intend to utilize `jaxdf` with NVidia GPU support, follow the instructions to install JAX accordingly.
To install `jaxdf` from PyPI, use the `pip` command:
```bash
pip install jaxdf
```
For development purposes, install `jaxdf` by either cloning the repository or downloading and extracting the compressed archive. Afterward, navigate to the root folder in a terminal, and execute the following command:
```bash
pip install --upgrade poetry
poetry install
```
This will install the dependencies and the package itself (in editable mode).
## Support
[](https://discord.gg/VtUb4fFznt)
If you encounter any issues with the code or wish to suggest new features, please feel free to open an issue. If you seek guidance, wish to discuss something, or simply want to say hi, don't hesitate to write a message in our [Discord channel](https://discord.gg/VtUb4fFznt).
## Contributing
Contributions are absolutely welcome! Most contributions start with an issue. Please don't hesitate to create issues in which you ask for features, give feedback on performances, or simply want to reach out.
To make a pull request, please look at the detailed [Contributing guide](CONTRIBUTING.md) for how to do it, but fundamentally keep in mind the following main guidelines:
- If you add a new feature or fix a bug:
- Make sure it is covered by tests
- Add a line in the changelog using `kacl-cli`
- If you changed something in the documentation, make sure that the documentation site can be correctly build using `mkdocs serve`
## Citation
[](https://arxiv.org/abs/2111.05218)
An initial version of this package was presented at the [Differentiable Programming workshop](https://diffprogramming.mit.edu/) at NeurIPS 2021.
```bibtex
@article{stanziola2021jaxdf,
author={Stanziola, Antonio and Arridge, Simon and Cox, Ben T. and Treeby, Bradley E.},
title={A research framework for writing differentiable PDE discretizations in JAX},
year={2021},
journal={Differentiable Programming workshop at Neural Information Processing Systems 2021}
}
```
#### Acknowledgements
- Some of the packaging of this repository is done by editing [this templace from @rochacbruno](https://github.com/rochacbruno/python-project-template)
- The multiple-dispatch method employed is based on `plum`, check out this amazing project: https://github.com/wesselb/plum
#### Related projects
1. [`odl`](https://github.com/odlgroup/odl) Operator Discretization Library (ODL) is a python library for fast prototyping focusing on (but not restricted to) inverse problems.
3. [`deepXDE`](https://deepxde.readthedocs.io/en/latest/): a TensorFlow and PyTorch library for scientific machine learning.
4. [`SciML`](https://sciml.ai/): SciML is a NumFOCUS sponsored open source software organization created to unify the packages for scientific machine learning.