torchsde

Project Overview: Torchsde

Torchsde is a powerful library implemented in Python for solving stochastic differential equations (SDEs) with enhanced computational efficiency and support for GPU accelerations. Stochastic differential equations are widely used in various fields like finance, physics, and machine learning to model systems affected by random fluctuations. This library provides differentiable SDE solvers, making them compatible with PyTorch, a leading open-source machine learning library.

Installation

To get started with torchsde, it can be easily installed using pip, which is a package manager for Python packages. Python version 3.8 or higher and PyTorch version 1.6.0 or higher are required.

pip install torchsde

Documentation

Detailed documentation is available here, offering comprehensive guidance on using the library effectively.

Practical Examples

Torchsde provides several examples to demonstrate its capabilities:

Quick Example

In a quick example, users can easily define an SDE within a custom class inheriting from torch.nn.Module. This involves specifying the drift and diffusion functions, f and g, respectively. The example illustrates how to initialize the SDE system and solve it over a specified time interval.

import torch
import torchsde

batch_size, state_size, brownian_size = 32, 3, 2
t_size = 20

class SDE(torch.nn.Module):
    noise_type = 'general'
    sde_type = 'ito'

    def __init__(self):
        super().__init__()
        self.mu = torch.nn.Linear(state_size, state_size)
        self.sigma = torch.nn.Linear(state_size, state_size * brownian_size)

    def f(self, t, y):
        return self.mu(y)

    def g(self, t, y):
        return self.sigma(y).view(batch_size, state_size, brownian_size)

sde = SDE()
y0 = torch.full((batch_size, state_size), 0.1)
ts = torch.linspace(0, 1, t_size)
ys = torchsde.sdeint(sde, y0, ts)

Notebook

A helpful Jupyter Notebook example (examples/demo.ipynb) is provided, guiding users on solving SDEs and addressing topics such as managing randomness and selecting appropriate noise types.

Learning Latent SDEs

For more advanced usage, the latent_sde.py example demonstrates learning a latent stochastic differential equation, comparing it to a variational autoencoder. This involves fitting an SDE to data, and users can execute it with a single command, specifying a directory for output figures.

python -m examples.latent_sde --train-dir <TRAIN_DIR>

Neural SDEs as GANs

The sde_gan.py example explores using SDEs within a Generative Adversarial Network (GAN) framework, where an SDE acts as the generator, while a neural controlled differential equation (CDE) functions as the discriminator.

python -m examples.sde_gan

Citation and References

Torchsde has been featured in notable academic conferences. Researchers using it in their work are encouraged to cite the related papers for academic credit.

• Scalable Gradients for Stochastic Differential Equations
• Neural SDEs as Infinite-Dimensional GANs

These resources reflect the cutting-edge research that underpins the torchsde project, making it an invaluable tool for scientific and machine learning applications involving stochastic processes.