Implementation of 2408.00835 by James Halverson and Sneh Pandya.
Hierarchical decays of
This project is developed for Python3.9 interpreter on a linux machine. Using an Anaconda virtual environment is recommended. Our implementation hinges on diffrax
and numpyro
, packages built on jax
🚀.
To install dependencies, you can run
pip install -r requirements.txt
or consult online documentation for appropriate dependencies. It is highly reccomended to follow jax install directions from the source.
Our codebase is structured as follows:
src/scripts
contains the main ingredients for our methodology.
stasis_simulation_differentiable.py
is the differentiable stasis simulation, containing the differentiable Boltzmann solver and stasis finder. This is used in SVI and gradient ascent experiments.stasis_simulation_non_diff.py
is the non-differentiable stasis simulation. This is used for creating plots and features an accurate, sliding-window stasis finder algorithm. This also has parametric functionality built in via themodel
argument, which allows one to generate power-law and exponential-dependent stasis configurations.svi.py
contains code for doing stochastic variational inference with normalizing flows. This requires ayaml
file specifying experimental parameters, an example of which is given insrc/tutorials/example_svi_experiment.yaml
.
src/tutorials
contains examples which coincide with the experiments in our paper.
gradient_ascent.ipynb
gives a walk-through in how gradients in the simulation can be used to produce stasis with minimal physical biases in the initialization. It also visualizes the optimization trajectories with a gif.model_comparison.ipynb
does a numerical comparison between power-law and exponential models of stasis. It also shows that both stasis models are attractors.random_stasis.ipynb
shows that rates and abundances drawn from certain families of distributions can result in persistent epochs of stasis.
Code author: Sneh Pandya
Issues and questions: @snehjp2, pandya.sne@northeastern.edu