RheOFormer

Code accompanying the arXive paper on RheOFormer, an OFormer-based neural operator for rheological and viscoelastic flow problems.

The repository covers two complementary sets of experiments. Both share the same OFormer backbone (model.py) but use different data pipelines, training scripts, and configs.

Pipeline	Setting	Models
2D (train_2d.py)	2D viscoelastic flow on irregular meshes	IrregSTEncoder2D + IrregSTDecoder2D
1D (train_1d.py, evaluate_1d.py)	1D rheometric protocols (TEVP, Giesekus)	Encoder1D + PointWiseDecoder1D

Repository layout

.
├── model.py                       # shared OFormer backbone (1D + 2D)
├── train_2d.py                    # canonical 2D viscoelastic flow training
├── train_1d.py                    # rheometric 1D training (TEVP / Giesekus)
├── evaluate_1d.py                 # rheometric 1D evaluation on canonical protocols
├── plot_1d_results.py             # rheometric 1D figure generation
├── requirements.txt
├── data/                          # dataset / data-generation modules
│   ├── viscoelastic_2d_data.py    # FlowDataset (2D mesh .npz loader)
│   ├── tevp_data.py               # JAX ODE solver + GRF inputs for TEVP
│   └── giesekus_data.py           # JAX ODE solver + GRF inputs for Giesekus
├── utils/                         # shared helpers (kept separate per pipeline)
│   ├── canonical_utils.py         # checkpoints, 2D losses, mesh visualisation
│   └── rheometric_utils.py        # 1D losses, normalisation, finite differences
└── configs/                       # all experiment configs (1d_* = rheometric, 2d_* = canonical)
    ├── 2d_triangle.yaml           # canonical: flow past a triangular obstacle
    ├── 2d_channel.yaml            # canonical: contraction channel
    ├── 1d_tevp.yaml               # rheometric: TEVP constitutive model
    └── 1d_giesekus.yaml           # rheometric: Giesekus constitutive model

Installation

pip install -r requirements.txt

The rheometric pipeline additionally requires JAX for on-the-fly data generation. Install the appropriate JAX wheel for your CUDA version (see JAX install docs).

Quick start

All commands below should be run from the repo root.

Canonical (2D viscoelastic flow)

python train_2d.py --config configs/2d_triangle.yaml   # flow past a triangular obstacle
python train_2d.py --config configs/2d_channel.yaml    # contraction channel

The configs point to directories of .npz mesh-data files; update train_data / test_data paths to match your local copy.

Rheometric (1D constitutive responses)

Model	Input	Output	Reference
TEVP (Thixotropic Elasto-Viscoplastic)	Scalar shear rate γ̇₁₂(t)	Structure parameter λ(t), shear stress σ₁₂(t)	Larson (2015)
Giesekus (Tensorial viscoelastic)	Velocity gradient ∇u₁₂, ∇u₂₁, ∇u₁₁, ∇u₂₂	Stress tensor σ₁₁, σ₂₂, σ₁₂	Giesekus (1982)

Each 1D config sets a log_dir (e.g. ./TEVP_Results, ./Giesekus_Results) that anchors all of the experiment's outputs:

<log_dir>/
├── checkpoints/<model>_best.pt, <model>_final.pt, <model>_iter*.pt
├── results/<model>_<protocol>.npz
└── figures/<plot>.png

Use a unique log_dir per run to avoid mixing results.

Training (data generated on-the-fly via JAX from Gaussian Random Field shear-rate inputs):

python train_1d.py --config configs/1d_tevp.yaml       # → ./TEVP_Results/checkpoints/
python train_1d.py --config configs/1d_giesekus.yaml   # → ./Giesekus_Results/checkpoints/

Evaluation on canonical test protocols (writes .npz to <log_dir>/results/):

# TEVP
python evaluate_1d.py --checkpoint TEVP_Results/checkpoints/tevp_final.pt --protocol oscillatory

# Giesekus
python evaluate_1d.py --checkpoint Giesekus_Results/checkpoints/giesekus_final.pt --protocol extensional
python evaluate_1d.py --checkpoint Giesekus_Results/checkpoints/giesekus_final.pt --protocol pure_shear
python evaluate_1d.py --checkpoint Giesekus_Results/checkpoints/giesekus_final.pt --protocol simple_shear

# Extrapolation study: GRF inputs scaled beyond the training range
python evaluate_1d.py --checkpoint TEVP_Results/checkpoints/tevp_final.pt --protocol grf --output_scale 2.0

Available --protocol values:

• TEVP: grf, oscillatory, steady_shear, step_down
• Giesekus: grf, extensional, pure_shear, simple_shear

Figure generation (reads <log_dir>/results/, writes <log_dir>/figures/):

python plot_1d_results.py --model tevp --config configs/1d_tevp.yaml
python plot_1d_results.py --model giesekus --config configs/1d_giesekus.yaml

Model architecture

RheOFormer uses an encoder–decoder–propagator OFormer architecture:

1. Encoder — embeds input function values with temporal coordinates, processes them through Fourier-type linear self-attention layers, and projects to a latent space.
2. Decoder — uses Gaussian Fourier feature projection for query coordinates, then cross-attention to transfer information from the encoder's latent representation to query locations.
3. Propagator — residual feed-forward blocks that evolve latent dynamics, approximating the temporal evolution of the underlying PDE in latent space.

Key design choices:

• Fourier-type linear attention (softmax-free) for computational efficiency
• Rotary positional embeddings for resolution-invariant encoding
• OneCycleLR scheduler with warm-up for stable convergence
• Standard normalisation (z-score) of inputs and outputs

Citation

If you use this code, please cite the accompanying paper.

RheOFormer