This repository stores the collection of scripts and modules tied to the in-progress revision of the preprint, "A framework for variational inference and data assimilation of soil biogeochemical models using state space approximations and normalizing flows".

Neural moving average flow Python experiment scripts and dependent modules located in "pytorch_sbm_sde_vi" folder. Experiment scripts correspond to the various "learn... .py" files. To replicate an individual experiment, such as a fixed-$\theta$ or or joint $\theta, x$ inference (the use of joint here signifying optimization of both differential equation soil biogeochemical model $\theta$ and hidden neural network $x$ parameters), run the script associated with the experiment. "learn_NN..." scripts correspond to fixed-$\theta$ flow experiments, while "learn_theta..." scripts correspond to joint inference experiments. Synthetic data used for inference conditioning is located in the "pytorch_sbm_sde_vi/generated_data" folder, while scripts sourcing the synthetic data are in the "pytorch_sbm_sde_vi/python_notebooks" folder. Code associated with the hamiltorch SSM-NUTS benchmarking is labeled with "mcmc_..." in the filename.