# Code for On Flow-based Generative Models for Probabilistic Forecasting

This repository contains the code for the paper "On Flow-based Generative Models for Probabilistic Forecasting".

## Installation
Create a new conda environment, install the dependencies and build the linear SDE library:
```
conda create -n neurips_2025 python=3.10
pip install -r requirements.txt
pip install ./DiffusionCRF
```
Install GPU support for JAX by following the instructions [here](https://docs.jax.dev/en/latest/installation.html#installation).

## Usage
Create the dynamical system data by running:
```
python -m Data.dynamical_system.data_gen
```
and generate the noisy double pendulum data by running:
```
python -m Data.pendulum.double_pendulum
```
Data generation may be slower in GPU mode compared to CPU mode.
Generate the scripts for the experiments by running:
```
python -m script_generator.generate_scripts
```
Run the experiments that are written to the `script_generator/python_commands.txt`.  For example the following command will train the NeuralSDE model on the noisy double pendulum data using the Wiener velocity model as the underlying SDE.
```
python main.py --config_file=Config/noisy_double_pendulum.yaml --freq=1 --sde_type=tracking --model_name=my_neural_sde_rnn_bwd --group=final_models --global_key_seed=0 --train
```





