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

@InProceedings{pmlr-v306-muser26a,
  title     = {Flowers: A Warp Drive for Neural {PDE} Solvers},
  author    = {Muser, Till and Spitzer, Alexandra and Lassas, Matti and De Hoop, Maarten V. and Dokmani\'{c}, Ivan},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  pages     = {91390--91418},
  year      = {2026},
  volume    = {306},
  series    = {Proceedings of Machine Learning Research},
  publisher = {PMLR},
  url       = {https://proceedings.mlr.press/v306/muser26a.html}
}

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