Source-linked AI summary
Dandelion: A Spherical Flower for Neural Simulation of Planetary Dynamics
Till Muser, Giovanni Abati, Ivan Dokmanić
TL;DR
Euclidean neural architectures are geometrically mismatched with spherical planetary dynamics, motivating native spherical PDE solvers. Dandelion uses tangent-plane geodesic warps and spherical-harmonic pooling, and it performs best or second-best across the benchmark while remaining subject to resolution and rollout-scope limits.
Problem
Euclidean grids distort spherical geometry and weaken structural rotational robustness for planetary dynamics.
Method
Dandelion adapts Flower with tangent-plane displacements transported along great circles and spherical-harmonic resolution changes, leaving warps as the sole spatial-mixing operator.
Results
Dandelion is best or second-best on every dataset and strictly best on a majority, while spherical adaptation cuts Flower2D rollout error by 1.5–4.5× on four datasets.
Takeaways & Limitations
The released suite provides six new spherical PDE problems plus planetswe, enabling mid-size models to be trained from scratch in hours.
Takeaways & Limitations
The study is limited to resolutions up to 256 × 512, at most two autoregressive training steps, and a synthetic benchmark rather than operational weather data.
Abstract
from arXiv · showhide
Many dynamical processes unfold on the sphere but the default scientific machine learning architectures are Euclidean. Applying these architectures on a regular lat-lon grid causes problems: Cartesian convolutions become distorted at high latitude; 2D FFTs in Fourier neural operators incorrectly assume double periodicity; Cartesian positional encodings in ViTs distort spherical geodesic distances. Recent work moves towards natively spherical primitives, including spherical convolutions (e.g., DeepSphere or DISCO), Spherical Fourier Neural Operators (SFNOs), and geodesic attention. Here we propose Dandelion, a spherical version of Flower, a warp-based neural PDE solver. Layers of Dandelion predict a tangent-plane displacement and transport features along great circles. We obtain a U-Net-like structure by implementing hierarchical pooling entirely in the spherical-harmonic domain. There are thus no convolutions: spatial mixing is achieved only through spherical coordinate changes, or warps. To compare Dandelion with existing spherical architectures, we release an evolving benchmark suite of challenging, natively-spherical PDE datasets including a modified Galewsky jet, anomalous chained turbulence, Cahn-Hilliard decomposition, spherical Riemann shocks, Held-Suarez dry atmospheric transport and global ocean dynamics. This new benchmark fills the gap in existing spherical datasets which are either too small and stylized, or much too large (ERA5) for model iteration. Dandelion is best or second-best on every dataset, and the gap to non-warp baselines widens with resolution: at $256\times 512$, Dandelion and Flower2D occupy the top two slots in both single-step prediction and rollout.
1 Introduction
Dandelion addresses the mismatch between Euclidean neural architectures and spherical planetary dynamics with native spherical warps. The paper also introduces a diverse, tractable benchmark suite for evaluating spherical PDE solvers.
- Spherical planetary dynamics conflict with Euclidean grids, whose convolutions, tokenizers, and positional encodings distort geometry and weaken rotational robustness.
- Dandelion adapts Flower by predicting tangent-space displacements and transporting features along spherical geodesics, while replacing strided pooling with spectral coarsening.
- Existing spherical evaluations range from small toy problems to prohibitively large ERA5-scale datasets, leaving no comparable suite for method development.
- Six new physically diverse spherical PDE datasets, combined with planetswe, form a seven-dataset benchmark spanning instability, turbulence, phase separation, shocks, atmospheric dynamics, and ocean dynamics.
- The datasets are small or moderate enough for models to train from scratch in hours, enabling repeated testing and debugging before operational-data scaling.
2 Methodology: Dandelion
Dandelion replaces grid-sensitive Euclidean operations with intrinsic spherical warps and harmonic-domain resolution changes. Its layers transport features along great circles without convolutional spatial mixing.
- Warping is the network’s only spatial-mixing primitive; feature transport uses input-dependent source locations at linear cost in grid size.
- Tangent warping: Lat–lon offsets produce latitude-dependent geodesic distances, motivating Dandelion’s intrinsic tangent-space displacement representation.
- Tangent warping: Dandelion predicts tangent-plane displacements and maps them to spherical sampling points through Rodrigues’s formula along great circles.The tangent-vector norm equals the traveled geodesic arc length.
- Spectral coarsening: The architecture retains Flower’s U-Net-like multiscale structure while replacing grid-sensitive pooling with spherical-harmonic coarsening and refinement.
- Spectral coarsening: Spherical-harmonic resolution changes are intrinsic and commute with SO(3) rotations, preserving rotational symmetry across resolutions.
3 Datasets
The benchmark suite covers diverse spherical PDE regimes, from instabilities and turbulence to phase separation, shocks, atmospheric circulation, and ocean dynamics. Its rollout figures highlight challenges in preserving evolving dynamics and long-horizon stability.
- Earth-like Shallow Water: The suite includes planetswe, a forced rotating shallow-water system designed to test long-horizon stability under continuous forcing and time-dependent periodicity.
- Rotated Double Galewsky: Barotropic-jets tests rotational and geometric reasoning through mirrored jets, cross-equatorial propagation, and random global rotations.
- Anomalous Chained Turbulence: Active-turbulence includes an anomalous phase with finite-size vortices forming percolating antiferromagnetic chains, a closed-manifold phenomenon without a flat analogue.
- Spinodal Decomposition: Cahn–Hilliard decomposition tests sensitivity to high-frequency modes, smoothness constraints, and conservation of total composition.
- The dry-atmosphere-3D and global-ocean-3D datasets test parameterized atmospheric circulation and seasonal, non-autonomous ocean dynamics with complex boundary conditions.
- In active-turbulence rollouts, only Dandelion captures increasing vorticity magnitude, while other models remain near the input and decay.
4 Results
Dandelion is evaluated against spherical, Euclidean, Fourier, and neighborhood-transformer baselines on seven datasets using matched next-step and rollout tasks. It performs strongly across datasets, while high-resolution timing results show that transformer baselines incur substantial computational costs.
- Dandelion is strictly best on 6/7 next-step cells and 5/7 rollout cells.
- Dandelion cuts Flower2D’s rollout error by 1.6× on planetswe, 1.5× on Shocks, 1.8× on barotropic-jets, and 4.5× on active-turbulence.
- Neighborhood transformers are competitive on some problems but perform worse on finer discretizations, where Flower2D and Dandelion rank ahead on cited datasets.
- At 128 × 256 and above, LR2Transformer training epochs are 5–6× longer and LS2Transformer epochs 30–40× longer than FNO.
- Throughput should be weighed alongside accuracy because parameter matching alone can misrepresent the computational cost of local transformers.
5 Conclusion
Dandelion combines tangent-plane warps with spherical-harmonic resolution changes, leaving warps as its only spatial-mixing operator. It performs strongly across the benchmark suite, while the study remains limited by moderate resolution, short rollouts, and synthetic data.
- Dandelion replaces Flower’s strided convolutions with spherical-harmonic resolution changes, leaving warps as the network’s only spatial mixing operator.
- Dandelion is best or second-best on every dataset and strictly best on a majority of problems.
- 1.5–4.5×: Dandelion cuts Flower2D’s rollout error on planetswe, barotropic-jets, shock-caps, and active-turbulence.
- The released suite contains six spherical PDE problems spanning shallow-water instabilities, turbulence, phase separation, shocks, dry atmospheric circulation, and global ocean dynamics.
- The benchmark caps at 256 × 512, measures at most two autoregressive steps in the loss, and uses synthetic data rather than operational weather data.
A Extended Dataset Descriptions
The appendix summarizes the benchmark datasets’ physical features and defines how fields and varied physical parameters are reported. It points readers to a consolidated summary table and detailed dataset specifications.
- Table 4 summarizes each dataset’s target fields and the physical parameters modified across runs.
- For each physical-parameter combination, multiple runs use different random seeds.
- The appendix provides detailed specifications for the spherical PDE datasets introduced in the work.
A.1 planetswe
planetswe is a forced, hyper-viscous shallow-water benchmark on a rotating sphere with time-varying daily and annual forcing. It is simulated spectrally on a 256 × 512 equiangular grid and provides 40 ERA5-initialized trajectories.
- planetswe solves forced, hyper-viscous shallow water on a rotating unit sphere and is the suite’s only externally generated dataset.
- The forcing varies daily and annually through the moving day/night terminator and seasonally modulated latitude.
- The simulations use Dedalus with a 256 × 512 equiangular polar grid and hourly snapshots.
- Figure 5 compares vorticity evolution along a barotropic-jets trajectory on spherical and latlon-grid visualizations using shared color limits.
- Each retained trajectory spans three model years and contains 3024 snapshots.
- The available fields are surface-height deviation h and the two components of tangent velocity u.
- The published ensemble contains 40 trajectories differing only in ERA5-derived initial conditions, with no physical-constant sweep.
- Flower’s U-Net uses multi-head SelfWarp blocks, strided 2×2 convolutions, transposed-convolution upsampling, and skip connections.
B.2 Dandelion
Dandelion retains Flower’s warp-based U-Net while replacing grid-sensitive operations with tangent-plane geodesic transport and spherical-harmonic coarsening. Its implementation uses equiangular spherical grids and compares against spectral and transformer baselines.
- Dandelion replaces grid-sensitive primitives while retaining Flower’s U-Net skeleton and head-based mixing.
- Tangent SelfWarp predicts local tangent-plane vectors and maps them onto the sphere with Rodrigues’ formula before sampling along great circles.
- Pole-crossing displacements require neither special longitude wrapping nor polar reflection.
- Spectral2xDown performs resolution changes through spherical-harmonic transforms, low-pass truncation, a 1×1 convolution, and GELU adjustment.
- Dandelion remains on the equiangular grid at every level despite residual spherical-harmonic aliasing that is considered small relative to explicit low-pass filtering.
- Inputs use a 3D unit-sphere embedding rather than raw longitude and latitude coordinates.
- FNO uses FFT-based global convolutions, whereas SFNO replaces the FFT with spherical-harmonic transforms.
- The FNO and SFNO implementations hold embedding dimension, layer count, MLP ratio, truncation, normalization, activation, and prediction style fixed.
C Experiment Details
The experiments use a shared next-step forecasting setup with four input frames, trajectory-level data splits, and channel-wise normalization before evaluation in physical units.
- The appendix documents dataset specifications, model architectures, hyperparameters, and the protocol used for every reported number.
- The model receives four consecutive frames stacked along channels and predicts the next frame; rollout advances the four-frame buffer autoregressively.
- Dataset splits are deterministic 80/10/10 train, validation, and test partitions at the trajectory level, preventing frame leakage across splits.
- Inputs and targets are z-score normalized per channel using training-split statistics, then predictions are denormalized before metric evaluation.
C.2 Two-phase training
Training combines single-step supervision with short autoregressive unrolling, using latitude-weighted spherical losses and validation throughout a 25-epoch schedule.
- Training lasts 25 epochs: 20 epochs optimize single-step prediction, followed by 5 epochs of two-step autoregressive training without teacher forcing.
- The autoregressive loss averages errors across both predicted steps, with gradients flowing through the complete two-step unroll.
- Validation runs every epoch with both single-step and 20-step rollout evaluations, while testing uses the final training state rather than a reloaded best checkpoint.
- Cosine-of-latitude weighting makes the discrete loss approximate spherical L2 error so each unit of solid angle contributes equally.
- AdamW uses a one-epoch warmup, cosine annealing over the remaining 24 epochs, weight decay 1 × 10^-4, and global gradient clipping at L2-norm 1.0.
- For each model–dataset cell, the base learning rate is selected from {10^-4, 5×10^-4, 10^-3} using the lowest final single-step validation loss.
C.5 Hardware and distributed setup
Experiments primarily use four-GH200 nodes with distributed training, under a 24-hour compute limit and architecture-specific memory-based batch sizing.
- Most runs use one node with four NVIDIA GH200 Grace–Hopper Superchips and PyTorch DistributedDataParallel across the four GPUs.
- Some planetswe trainings instead run on a single NVIDIA H200.
- All training runs have a 24-hour time limit because of limited compute budget, with larger-dataset transformer runs most likely to approach it.
- Per-GPU batch sizes are tuned through H200 memory probes and scaled to GH200 using the HBM ratio 96/140 ≈ 0.686.
- LocalS2Transformer training did not finish for several higher-resolution datasets within a reasonable timeframe.
C.6 Evaluation
Evaluation measures held-out test performance with latitude-weighted spherical metrics for both one-step prediction and autoregressive rollouts.
- Table 2 reports single-step and rollout results on the held-out test split using latitude-weighted spherical metrics.
- The rollout autoregressively unrolls each test trajectory for min(20, Tout) steps from its first four-frame window and averages metrics over steps.
- Spatial means use the same cos(lat) weights as the training loss, while per-step rollout curves are logged separately from the main table.