Prototypes for the third ML-RT paper: methods barely used in astrophysical radiative transfer yet, evaluated for comparability against papers 1–2 (MLP / CVAE / CGAN / LSTM + custom) and studied for how they optimise.
ML-RT2/ is this paper-3 project; the earlier (unfinished) PINN exploration is kept for reference
under archive/. The ML-RT/ project (papers 1–2) is left untouched.
STARDUST is a deterministic map:
| input | output |
|---|---|
| 8 physical parameters (log halo mass, redshift, source age, QSO α, QSO efficiency, stellar escape fraction, IMF slope, log IMF min-mass) | 4 radial profiles — x_HII, x_HeII, x_HeIII, T — each 1500 points over radius |
Dataset 053 (41,081 samples) is used for direct comparison with papers 1–2. Transforms
match those papers: parameters min-max scaled to [0,1]; fraction profiles log10(x+1e-6);
temperature log10(T); we additionally z-score each channel on the train split (stored on the
Normalizer for inversion). Split is 80/10/10 at seed 42.
The input is only 8 scalars, so the "advanced" payoff is in the decoder / generative / physics side, not attention-over-input. Three frontiers, plus stretch ideas:
| theme | model(s) | why it's a fit / novel for astro-RT | status |
|---|---|---|---|
| A. Neural operators | FNO-1D, DeepONet | outputs are smooth 1D fields with a sharp front; spectral / basis operators are resolution-free and cheap at inference | prototypes done |
| B. Physics-informed operator | PINO | operator + a differentiable grey ionisation-equilibrium residual (ported recombination fits); data-anchored, avoids the old PINN failure; most novel | prototype done |
| C. Generative / probabilistic | conditional flow matching (diffusion later) | successor to CVAE/CGAN; calibrated ensembles + error bars; few-step sampling keeps inference cheap | prototype done |
| D. Attention decoder | profile/signal transformer | long-range coupling (front position ↔ integrated absorption), inter-species attention | prototype done |
| E. Stretch | Neural-ODE-in-r, JEPA, CNP | very novel for astro; physical structure / shared latent / lightweight uncertainty | prototypes done |
All models share the data module, metrics, trainer, and history format, so the cross-architecture optimisation study (loss curves, per-channel val error, gradient norms, epoch wall-time, sample-efficiency sweeps) is apples-to-apples.
ML-RT2/ (this project = paper 3; the earlier PINN work is in archive/)
environment.yml conda env (CUDA 12.1; spans 3090 → H200)
STATUS.md shareable progress summary
data/ dataset 053 / 057 (see data/053_data_set/get_data.sh)
apptainer/emulators.def portable image for the heterogeneous cluster
docs/
methods.tex/.pdf shareable methods brief (formulas, refs, pipeline + per-model diagrams)
refs.bib
diagrams/
*.tex/.pdf standalone per-architecture TikZ schematics (8) (+ README)
archive/ earlier PINN exploration (legacy_pinn/, study_material/)
src/
data.py 053 loader, transforms, Normalizer, synthetic sanity task
config.py dataclass config, minimal CLI
metrics.py encoded MSE + per-channel + physical rel-L2 + front-position error
physics.py torch RT physics for PINO (recombination fits + spatial residual)
trainer.py bf16 AMP, grad-accum, checkpoint/resume, diagnostics -> history.json
train.py entry point (dispatches by --model)
sweep.py Optuna search (Hyperband/ASHA pruning) + 1-D sensitivity sweeps
analyze.py load a run -> test metrics + result plots
plots.py profile (emulator vs STARDUST) + training-dynamics figures
sanity_check.py shape + overfit-a-batch on the synthetic task
models/
fno.py 1D Fourier Neural Operator (self-contained)
deeponet.py DeepONet with Fourier-feature trunk
flow.py conditional flow matching (generative, gives ensembles)
pino.py physics-informed neural operator (operator + RT residual)
transformer.py profile/signal transformer decoder
node.py Neural-ODE decoder in radius
jepa.py joint-embedding predictive architecture
cnp.py conditional neural process (per-radius uncertainty)
Local sanity (CPU, no data):
cd src && python sanity_check.py # shape + overfit-a-batch for each modelGet the data (into ML-RT2/data/053_data_set): run data/053_data_set/get_data.sh, which
fetches the dataset from Google Drive via gdown (included in the conda env). See the comments at
the top of the script for the one-time host setup (the Drive file id).
Train (cluster):
python train.py --model fno --run_name fno_053 --epochs 300
python train.py --model deeponet --run_name deeponet_053 --epochs 300 --latent 256
python train.py --model flow --run_name flow_053 --epochs 400 # generative
python train.py --model pino --run_name pino_053 --epochs 300 --physics_weight 0.1
# resume after a wall-time kill:
python train.py --model fno --run_name fno_053 --resume output/fno_053/last.ptTune (parallel workers share one study):
python sweep.py --model fno --study fno_053 --storage sqlite:///sweeps/fno.db \
--trials 40 --epochs 60 --subset 8000 # run on several nodes at once
python sweep.py --model fno --sensitivity modes --values 16,24,32,48,64 --epochs 80Analyse a run (test metrics + plots into <run_dir>/plots/):
python analyze.py --run_dir output/fno_053- 4-day wall-time cap → the trainer writes resumable checkpoints (
best.pt,last.pt) and a livehistory.json;--resumecontinues seamlessly. - Heterogeneous GPUs (3090 / L40S / A100 / H200) → bf16 autocast everywhere, and batch size ×
--grad_accumlets the same run fit 24 GB or 141 GB. The Apptainer image pins the CUDA userspace so one.sifruns on every node.
All eight architectures (FNO, DeepONet, PINO, flow matching, transformer, Neural-ODE, JEPA, CNP)
are implemented, sanity-verified (shape + overfit-a-batch), and run end-to-end on real 053 data. The
Optuna sweep harness (search + sensitivity), analysis/plotting, the methods brief (docs/methods.pdf
- per-architecture diagrams), and the cluster env/Apptainer image are all in place. Next: launch
the cluster runs. Deferred by choice: the
table_ion.cHeII physics fix (awaiting paper/author review) and extending the PINO residual (helium equilibrium / heating term).