Repository files navigation

This repository provides an efficient implementation injax of a score matching and physics informed neural network-based algorithm for solving the stationary Fokker-Planck equation in high dimension.

The implementation is built on Google'sjax package for accelerated linear algebra and DeepMind'shaiku package for neural networks. Both can be installed by following the guidelines at the linked repositories.

Routines common to all implemented simulations can be found inpy/common, including implementations of the various neural networks used, systems studied, and loss functions used.

Simulation code to launch learning experiments can be found inpy/launchers.

Code for generating datasets can be found inpy/dataset_gen.

Code for visualizing the output of simulations and for producing the publication figures can be found innotebooks.

Slurmsbatch scripts used to launch the experiments in the paper can be found underslurm_scripts.

Experiment tracking is implemented inWeights and Biases. You will need to input a project title in the corresponding simulation launcher in the call towandb.init.

If you found this repository useful, please cite:

[1] N. M. Boffi and Eric Vanden-Eijnden. “Deep learning probability flows and entropy production rates in active matter", arXiv: 2309.12991.

@misc{boffi2023deep,      title={Deep learning probability flows and entropy production rates in active matter},       author={Nicholas M. Boffi and Eric Vanden-Eijnden},      year={2023},      eprint={2309.12991},      archivePrefix={arXiv},      primaryClass={cond-mat.stat-mech}}