DISCLAIMER: Package is in continued development.

StochFlow is a Python library that implements the stochastic interpolant models and algorithms described in the paper Stochastic Interpolants via Diffusive and Flow-based Processes by Chen et al [preprint]. The library provides a set of tools for designing generative models using stochastic processes associated with time-evolving probability density functions as well as the forward and backward Fokker-Planck equations.

• Implementation of the stochastic interpolant models and algorithms described in the paper.
• Numerical methods for solving the forward and backward Fokker-Planck equations, including the Crank-Nicolson method and the Strang splitting method.
• Tools for generating samples from a given probability distribution using the forward and backward generative models.
• Support for both flow-based and diffusion-based methods for generative modelling.
• Examples and tutorials demonstrating the use of the library.

To install StochFlow, simply run the following command:

pip3 install stochflow
pip3 install $(grep -ivE "torch|torchvision|scikit-learn" requirements.txt)

It may be necessary to manually install PyTorch and scikit-learn.

PyTorch ➞ Select the stable build for your operating system, your preferred package manager, and compute platform depending on your GPU capabilities. Then, run the recommended commands to install PyTorch locally.

scikit-learn ➞ Select your operating system/packager, and run the provided command to install the latest official release.

To use StochFlow, simply import the relevant modules and functions:

import stochflow as sf

# Define the initial and final probability density functions
rho_0 = sf.GaussianMixtureModel(...)
rho_1 = sf.GaussianMixtureModel(...)

# Solve the forward Fokker-Planck equation to generate samples from rho_1
samples_forward = sf.forward_generative_model(rho_0, rho_1, ...)

# Solve the backward Fokker-Planck equation to generate samples from rho_0
samples_backward = sf.backward_generative_model(rho_0, rho_1, ...)

The library provides a range of options and parameters for customizing the generative models and numerical methods used to solve the Fokker-Planck equations. See the documentation and examples for more details.

The library includes several examples and tutorials demonstrating the use of the stochastic interpolant models and algorithms. These examples cover a range of applications, including image generation, time series modeling, and data augmentation.

Contributions to StochFlow are welcome! If you find a bug or have a feature request, please open an issue on the GitHub repository. If you would like to contribute code, please fork the repository and submit a pull request.

StochFlow is released under the MIT License. See the LICENSE file for more details.