Installation

Install dependencies using uv:

uv sync

Usage

To train your own Flower model, have a look at flower_standalone.py. It's a self-contained model file with no dependencies beyond torch, so you can drop it directly into your project.

To reproduce results from the paper, use the flowers.train entrypoint. It uses a modular configuration system with separate configs for data, model, and training parameters.

Running Training

uv run python -m flowers.train \
  --data configs/data/<dataset>.yaml \
  --model configs/models/<model>.yaml \
  --train configs/train.yaml

For example, to train a UNet model on the viscoelastic instability dataset:

uv run python -m flowers.train \
  --data configs/data/viscoelastic_instability.yaml \
  --model configs/models/unet_convnext.yaml \
  --train configs/train.yaml

Acknowledgments

Model and Code

@misc{muser2026flowers,
      title={Flowers: A Warp Drive for Neural PDE Solvers}, 
      author={Till Muser and Alexandra Spitzer and Matti Lassas and Maarten V. de Hoop and Ivan Dokmanić},
      year={2026},
      eprint={2603.04430},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2603.04430}, 
}

If you use this repo, please cite the appropriate publications:

• If you use Flower, please cite our paper using the bibtex above.
• The training framework is built on The Well, as are many of the datasets used here.
• The scOT model implementation is adapted from POSEIDON

Data

• Primary datasets are from The Well.
• PDEBench data is taken from PDEBench
• The anisotropic helmholtz dataset is from WaveBench