Data Generation API

Neojax provides some data generation functions. It currently only supports generating data for Burgers equation.

Setup

Generating data requires the installation of the optional diffrax dependency.

pip3 install "neojax-operators[gen]"

or pip3 install ".[gen]" from the project root.

1D Burgers Equation

Burgers equation is a classical example of a 1D non-linear PDE used in Neural Operator publications:

\[ \begin{aligned} \frac{\partial}{\partial t}u(x, t) + \frac{1}{2}\frac{\partial}{\partial x}(u(x,t))^2 &= \nu \frac{\partial^2}{\partial x^2}u(x,t), \; \quad x \in (0,2\pi), t \in (0, \infty) \\ u(x, 0) &= u_0(x), \quad \quad \quad \quad x \in (0, 2\pi) \end{aligned} \]

with a fixed viscosity \(\nu = 10^{-1}\). The initial condition is drawn from \(\mu = \mathcal{N}(0, C)\), i.e., \(u_0 \sim \mu\), where

\[ C = 625(-\frac{d^2}{dx^2} + 25I)^{-2}. \]

neojax.data.generation.burgers_1d.generate_burgers_1d(n_samples: int = 1000, n_grid_points: int = 1024, nu: float = 0.1, t1: float = 1.0, batch_size: int = 100, seed: int = 42, use_bundle: bool = False) -> neojax.data.datasets.raw_dataset.RawDataset | neojax.data.datasets.bundle_dataset.BundleDataset

Generates 1D Burgers equation dataset.

Solves equation using a pseudo-spectral split step method as in the reference.

Generation is batched to limit memory usage.

Parameters:

Name	Type	Description	Default
n_samples	int	Number of samples. Default is 1000.	1000
n_grid_points	int	Number of spatial grid points. Default is 1024.	1024
nu	float	Viscosity. Default is 0.1.	0.1
t1	float	Maximum time. Default is 1.0.	1.0
batch_size	int	Batch size for generation to save memory. Default is 100.	100
seed	int	Jax random seed. Default is 42.	42
use_bundle	bool	Whether to return a BundleDataset instead of RawDataset. Default is False.	False

Returns:

Type	Description
RawDataset | BundleDataset	An instance of either RawDataset or BundleDataset. If RawDataset, assigns attrs inputs and labels, else assigns fields of BundleDataset as concatenated inputs and labels along first axis (and assigns corrds).

Cite

Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs

@article{kovachki2023neural,
    title={Neural operator: Learning maps between
    function spaces with applications to pdes},
    author={Kovachki, Nikola and Li, Zongyi and Liu,
    Burigede and Azizzadenesheli, Kamyar and Bhattacharya,
    Kaushik and Stuart, Andrew and Anandkumar, Anima},
    journal={Journal of Machine Learning Research},
    volume={24},
    number={89},
    pages={1--97},
    year={2023}
}