---
license: cc-by-nc-4.0
tags:
  - physics
  - pde
  - neural-operator
  - cfd
  - fluid-dynamics
  - poseidon
  - scOT
library_name: transformers
pipeline_tag: image-to-image
base_model: camlab-ethz/Poseidon-B
datasets:
  - empirischtech/pyfr-cylinder-kh
---

# Poseidon-B — Joint fine-tune (Cylinder + Kelvin–Helmholtz)

A fine-tuned checkpoint of [**Poseidon-B**](https://huggingface.co/camlab-ethz/Poseidon-B)
(scOT, ~158M parameters) extended to **two flow regimes that are absent from Poseidon's
pretraining distribution**, added simultaneously in a single training run:

1. **Cylinder** — wall-bounded, incompressible flow past a circular cylinder (vortex shedding), Re ≈ 62–180.
2. **KHC** — a compressible Navier–Stokes Kelvin–Helmholtz double shear layer.

The goal of this checkpoint is to add these new regimes **without catastrophic forgetting**
of Poseidon's original families. It is trained jointly on the two new regimes together with a
small **replay** buffer of two pretraining families, so that after adaptation the original
families remain near (or below) their original error while the two new regimes are fit well.

> **Research model.** Like the base model, this is intended for research use only.

## Model description

- **Architecture:** `ScOT` (Swin-V2-style multiscale operator transformer, `swinv2` backbone), identical to Poseidon-B — the new regimes were added **without changing the architecture or the input/output channels**.
- **Input/output:** a 4-channel physical state `(ρ, u, v, p)` on a `128 × 128` grid.
- **Conditioning:** lead-time conditioned. The time input is a normalized lead time `t = k / (T − 1)`, where `k` is the frame step and `T` the trajectory length **on the 21-frame cadence the model was trained on** (see Usage).
- **Solid-body encoding:** the cylinder is written **into the density channel** — `ρ = 1` in fluid, `ρ = 0` in solid — so the geometry enters through an existing channel and doubles as a geometry mask. No extra input channels are added.
- **Parameters:** ~158M (full fine-tune; all weights trainable).

Key config values (see `config.json`): `image_size=128`, `num_channels=4`,
`num_out_channels=4`, `embed_dim=96`, `depths=[8,8,8,8]`, `num_heads=[3,6,12,24]`,
`window_size=16`, `use_conditioning=true`.

## Training

Full fine-tuning from `Poseidon-B` with a joint data mixture (new regimes at weight 1.0,
replay families at weight 0.25):

| Dataset | Role | Sampling weight |
|---|---|---|
| Cylinder (train) | new regime | 1.0 |
| KHC (train)      | new regime | 1.0 |
| CE               | replay (retention) | 0.25 |
| RP               | replay (retention) | 0.25 |

**Hyperparameters:**

- Optimizer: AdamW, lr `2e-4`, weight decay `1e-5`
- Schedule: OneCycle, 10% warm-up
- Loss: `L1`, gradient clipping at 1.0
- Batch size: 192, 31 optimizer steps/epoch × 50 epochs = **1,550 updates** (fixed budget)
- Mixed precision (AMP), seed 0
- **All trajectories resampled to a common length of `T = 21` frames**; frame step `k ∈ [1, 12]`
- Checkpoint selection: lowest validation error on 30 held-out cylinder trajectories

Per-family normalization is used: each family is normalized by its own channel statistics
(the exact stats used are shipped in `norm_stats_used.json`). The binary geometry channel is
left un-normalized.

## Evaluation

Reported as per-channel NRMSE (RMSE divided by the channel's value range), with a
no-change **persistence** baseline for reference. All numbers below are 3-seed means.

**New regimes**

| Regime | Metric | Value |
|---|---|---|
| Cylinder | lead-2 NRMSE | **0.018** |
| KHC (holdout) | NRMSE (self-normalized) | **0.0377** |

Both are well below their persistence baselines.

**Retention on Poseidon's pretraining families** (factor relative to the unchanged base model; `< 1.0` = improved, `> 1.0` = degraded):

| Family | In replay buffer? | Relative to pretrained |
|---|---|---|
| CE     | yes  | 0.63× |
| RP     | yes  | 0.69× |
| CRP    | **no** (held-out probe) | 1.16× |
| CE-KH  | **no** (held-out probe) | 0.77× (density channel) |

All four tracked families land between **0.63× and 1.16×** their original error, i.e. the
two new regimes were added while the original families were largely retained or improved —
including families that were never placed in the replay buffer.

**Rollout stability (cylinder).** Autoregressive rollout over held-out trajectories stays
stable: median NRMSE grows only from ~0.007 to ~0.017 over 20 steps, with field correlation
above 0.9 at every step (median ≈ 0.985 at step 20).

## Intended use & limitations

- **Intended use:** research on operator learning, PDE surrogates, transfer/retention, and
  fast approximate rollouts for the two supported regimes (cylinder wake and compressible
  KH shear layers) and Poseidon's original families.
- **Resolution:** all fields are on a `128 × 128` grid; claims are restricted to wake-scale
  dynamics (fine scales are smoothed).
- **Extrapolation:** the cylinder holdout sits adjacent to the training band; accuracy
  degrades smoothly with distance in Reynolds number and is only validated up to ~33% beyond
  the training maximum, and on 2D flow (below the onset of 3D wake instability).
- **Rollout horizon:** validated within ~4 shedding periods. The frame-resampled model has a
  minimum usable lead of `1/20` of the trajectory window; shorter leads degrade sharply.
- **Not a general CFD solver.** It is a regime-specific surrogate, not a replacement for a
  PDE solver outside the tested distributions.

## Usage

The checkpoint is saved in Hugging Face `save_pretrained` format (`config.json` +
`pytorch_model.bin`) and loads with the `scOT` model class. Two conventions matter for
getting the reported accuracy:

1. **Normalize per channel** with the target family's statistics (shipped in
   `norm_stats_used.json`), and de-normalize the output with the same statistics.
2. **Frame cadence.** The model was trained on trajectories **resampled to `T = 21` frames**.
   The time input is the normalized lead `t = k / (T − 1) = k / 20`, where `k` is the frame
   step on that 21-frame cadence. A raw trajectory at a different frame count must first be
   resampled to 21 frames (e.g. by evenly spaced index selection) for the lead time to be
   meaningful. `k = 2` (`t = 0.1`) is the canonical lead this checkpoint was selected on.

```python
import numpy as np
import torch
from scOT.model import ScOT  # from the Poseidon / scOT codebase

model = ScOT.from_pretrained("rwmasood/poseidon-b-joint-cyl-kh").eval().cuda()

# Cylinder channel statistics used at training time (from norm_stats_used.json).
# Channels are (rho, u, v, p). For the cylinder regime rho encodes the body: 1=fluid, 0=solid.
mean = torch.tensor([0.992431640625, 0.9874249696731567,
                     -6.629727431572974e-05, 0.9546216726303101]).view(1, 4, 1, 1).cuda()
std  = torch.tensor([0.08666647970676422, 0.1872677057981491,
                     0.15668894350528717, 0.11970974504947662]).view(1, 4, 1, 1).cuda()

# A raw trajectory of physical fields (rho, u, v, p), shape (n_frames, 4, 128, 128).
traj = np.load("cylinder_case.npy")

# Resample to the 21-frame cadence the model was trained on.
T = 21
idx = np.linspace(0, traj.shape[0] - 1, T).round().astype(int)
traj = torch.from_numpy(traj[idx]).float().cuda()

# Predict frame i + k from frame i, with lead time k / (T - 1).
i, k = 0, 2
x = (traj[i:i + 1] - mean) / std
t = torch.full((1,), k / (T - 1), device=x.device)   # 2 / 20 = 0.1

with torch.no_grad():
    pred_norm = model(pixel_values=x, time=t).output   # scOT returns a ScOTOutput with `.output`

pred = pred_norm * std + mean                          # de-normalize to physical units
```

For the **KHC** regime, swap in that family's statistics instead:

```python
# KHC channel statistics (rho, u, v, p), from norm_stats_used.json
mean = torch.tensor([1.4987393617630005, 0.15815676748752594,
                     -0.00015812707715667784, 11.184250831604004]).view(1, 4, 1, 1)
std  = torch.tensor([0.32741427421569824, 0.6542320847511292,
                     0.4637199640274048, 3.0745689868927]).view(1, 4, 1, 1)
```

Autoregressive rollout: feed the model's de-normalized prediction back in as the next input.

> Sanity check: on the 30 held-out cylinder trajectories at lead 2 (21-frame cadence), this
> checkpoint gives a mean per-channel NRMSE of ≈ 0.017, matching the reported cylinder number.

## Files

- `config.json` — model configuration (`ScOT` / `swinv2`)
- `pytorch_model.bin` — model weights
- `norm_stats_used.json` — per-family normalization statistics (cylinder, KHC, CE, RP)

## Environment notes

- `transformers==4.29.2` (the version the base model config targets).
- The `scOT` package is required to instantiate the model class.
- On recent NVIDIA GPUs (Blackwell / sm_120), install a matching CUDA build of PyTorch.

## Base model & attribution

Fine-tuned from [`camlab-ethz/Poseidon-B`](https://huggingface.co/camlab-ethz/Poseidon-B)
(Herde et al., *Poseidon: Efficient Foundation Models for PDEs*, https://arxiv.org/abs/2405.19101).
Cylinder and KHC training data were generated with the [PyFR](https://www.pyfr.org/) solver.

## License

Released under **CC-BY-NC-4.0**, inheriting the base model's license (research/non-commercial use).