README.md

---
language:
  - en
tags:
  - "project:genesis"
  - "project:ArtIMis"
  - "type:model"
  - "science:materials"
  - "science:physics-simulation"
  - "risk:general"
  - "custom-code"
license: other
license_name: heat-public-release-unlimited-distribution
license_link: https://doi.org/10.25583/2571471
library_name: transformers
pipeline_tag: image-to-image
---

# Janus-first2last

Janus-first2last is a conditional DDPM/DDIM model for HEAT-PLI `av_density` field prediction. This checkpoint is the `first-to-last prediction` variant: condition on the first HEAT-PLI frame at simulation time 0.0 and predict the last HEAT-PLI frame at simulation time 25.0.

This Hugging Face release packages the model as custom Transformers code with `AutoModel` support. The original PyTorch `.pt` checkpoint was converted to `model.safetensors` so the public artifact contains model weights only, not optimizer state.

*Last Updated*: **2026-06-24**


## Model Changelog

+ **2026-06-24** Hugging Face release package with `AutoModel` support and `model.safetensors` weights.
+ **2026-05-22** Initial Janus-first2last model card for the DDPM checkpoint package.

## Model short description

Conditional DDPM model with DDIM sampling for HEAT-PLI density-field first-to-last prediction.

## Model description

Janus-first2last uses a conditional UNet diffusion model implemented in `modeling_janus_ddpm.py`. The model conditions on one HEAT-PLI `av_density` frame and generates the corresponding target frame by DDIM sampling.

The model operates directly on density fields rather than VAE latents. At inference time, the input field is resized to the model internal resolution `(560, 200)`, min-max normalized per sample to `[-1, 1]`, used as the condition channel, and decoded back to the input field resolution after sampling.


## Finetuned from model (optional)

Not applicable. This checkpoint is not documented as a fine-tune, adapter, quantized model, or merge of an existing base model.

## Model Type

Conditional DDPM/DDIM UNet for scientific image-to-image field prediction. The repository includes custom Transformers code and requires `trust_remote_code=True` when loaded through `AutoModel`.

## Inputs and outputs

- Input field: HEAT-PLI `av_density`, provided as a NumPy array of shape `(H, W)` or as an `.npz` file containing the `av_density` field.
- Demo data shape in this package: `(1120, 400)` with `float16` storage; three demo cases are included under `examples/id01040`, `examples/id01045`, and `examples/id01050`.
- Public training data source: `https://oceans11.lanl.gov/heat/pli/pli240420_fp16_full/`.
- Internal model resolution: `(560, 200)`.
- Output field: predicted `av_density` array with the same spatial resolution as the input.
- Direction: condition on the first HEAT-PLI frame at simulation time 0.0; predict the last HEAT-PLI frame at simulation time 25.0.

## Compute Infrastructure

The reported training and evaluation runs were performed on an 8x NVIDIA A100 GPU server.

### Hardware

Use A100 GPUs with at least 30 GB of GPU memory for comparable training and evaluation runs. Inference can run on CPU or a single GPU, but GPU inference is recommended.

### Software

```txt
numpy>=1.24
torch>=2.1
transformers>=4.39
huggingface_hub>=0.20
safetensors>=0.4
```

## Papers and Scientific Outputs

```bibtex
@article{Chu2026DiffUNet2,
  title = {DiffUNet^2: Bidirectional Prediction, Probabilistic Generation and Collaborative Visual Discovery for Scientific Data},
  author = {Chu, Mengdi and Yang, Jiaxin and Forbes, Angus G. and Debardeleben, Nathan and Lawrence, Earl and Biswas, Ayan and Shen, Han-Wei},
  year = {2026},
  doi = {https://doi.org/10.48550/arXiv.2606.03926}
}
```

Associated paper: https://doi.org/10.48550/arXiv.2606.03926

## Model License

Other. 

## Contact Info and Model Card Authors

Mengdi Chu, `chu.752@osu.edu`

# Intended Uses

This section describes intended research uses for the Janus-first2last checkpoint and documents use cases that are out of scope.

## Intended Use

This model is intended for research use as a learned surrogate for HEAT-PLI density-field first-to-last prediction. It can be used for fast exploratory prediction, uncertainty sampling through repeated DDIM runs, and model-card/API demonstration on the packaged HEAT-PLI demo data.

### Primary Intended Users

Researchers and engineers studying scientific machine learning, surrogate modeling, visual analytics, and uncertainty-aware prediction for physics simulation data.

### Mission Relevance

This project is funded by ArtIMis.

## Out-of-Scope Use Cases

This model is not intended for operational safety decisions, real-time hardware control, or extrapolation to HEAT configurations, fields, timesteps, or parameter ranges not represented in the training and evaluation data. Generated fields should be validated against domain constraints and high-fidelity simulations before scientific or engineering conclusions are drawn.

# How to use

Use `AutoModel.from_pretrained(..., trust_remote_code=True)` or the packaged `api_ddpm.py` CLI.

## Install Instructions

```bash
pip install numpy torch transformers huggingface_hub safetensors
```

## Training configuration

Training configuration is reconstructed from the model implementation, source checkpoint metadata, and the private training notebook used to train this checkpoint. The notebook itself is not included in this Hugging Face release.

- Source checkpoint: `DDPM_first2last/first2last_latest.pt`
- Published weights: `model.safetensors`
- Public source dataset: `https://oceans11.lanl.gov/heat/pli/pli240420_fp16_full/`
- Training field: `av_density`
- Pairing strategy: first available frame in each PLI case folder as condition, last available frame as target
- Data split: folder-level random split, `70%` train, `15%` validation, `15%` test
- Split seed: `42`
- Input preprocessing: load `av_density`, cast to `float32`, resize from `(1120, 400)` to `(560, 200)` with bilinear interpolation, per-sample min-max normalize to `[-1, 1]`
- Conditional UNet input channels: `2`
- Conditional UNet output channels: `1`
- Base channels: `64`
- Timestep embedding dimension: `256`
- DDPM training timesteps: `1000`
- Beta schedule: linear from `1e-4` to `0.02`
- Batch size: `8`
- Planned training epochs: `800`
- Published checkpoint epoch: `533`
- Optimizer: `AdamW`
- Learning rate: `1e-4`
- Learning-rate scheduler: `CosineAnnealingLR`, `T_max=800`
- Gradient clipping: global norm `1.0`
- Checkpointing: save every `2` epochs and whenever validation loss improves
- Training objective: MSE between predicted noise and sampled Gaussian noise
- Conditioning field: HEAT-PLI `av_density`
- Target field: HEAT-PLI `av_density`

## Inference configuration

- Default sampling method: DDIM
- Default sampling steps: `50`
- Default `eta`: `1.0`
- Optional ensemble sampling: `n_samples > 1`
- Optional fixed random seed: `seed`
- Recommended loading mode: `AutoModel.from_pretrained(repo_id, trust_remote_code=True)`

# Code snippets of how to use the model  

Library usage from the Hub:

```python
import numpy as np
from transformers import AutoModel

repo_id = "YOUR_USERNAME/Janus-first2last"
model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
model = model.to("cuda")  # optional

sample = np.load("examples/id01040/pli240420_id01040_pvi_idx00000.npz")
density = sample["av_density"]
prediction = model.predict(density, num_steps=50, eta=1.0, seed=42)
mean, std = model.predict_ensemble(density, n_samples=5, num_steps=50, eta=1.0)
```

Command-line inference after cloning or downloading the repository:

```bash
python api_ddpm.py \
  --model . \
  --input examples/id01040/pli240420_id01040_pvi_idx00000.npz \
  --output outputs/id01040_first2last_pred.npz \
  --field av_density \
  --num_steps 50 \
  --eta 1.0
```

Command-line inference directly from the Hub:

```bash
python api_ddpm.py \
  --model YOUR_USERNAME/Janus-first2last \
  --input examples/id01040/pli240420_id01040_pvi_idx00000.npz \
  --output outputs/id01040_first2last_pred.npz \
  --field av_density \
  --num_steps 50 \
  --eta 1.0
```

# Limitations

## Risks

This is a domain-specific scientific surrogate model. The primary risks are over-trusting generated fields, applying the model outside the HEAT-PLI distribution, or using predictions without domain validation. 

## Limitations

- The model predicts `av_density` only.
- The API uses per-sample min-max normalization, so absolute calibration should be checked for downstream analysis.
- The model resizes fields to `(560, 200)` internally and resizes predictions back to the input resolution.
- Stochastic DDIM sampling can produce different outputs unless a fixed seed is provided.
- Performance outside the represented HEAT-PLI first/last-frame prediction setting has not been established in this card.

# Training details

In this section, training details are documented from the available checkpoint metadata and model implementation.

## Training data

The model was trained on HEAT-PLI `av_density` endpoint pairs from the HEAT dataset. The public PLI fp16 dataset is available at `https://oceans11.lanl.gov/heat/pli/pli240420_fp16_full/`; the general HEAT dataset page is `https://oceans11.lanl.gov/heat/`, with DOI `10.25583/2571471`.

The training code scanned case folders containing `.npz` files, kept folders with at least two frames, and used the first sorted frame as the conditioning frame and the last sorted frame as the target frame. The packaged demo data include three small case folders, `examples/id01040`, `examples/id01045`, and `examples/id01050`, each with two frames: `idx00000` at simulation time 0.0 and `idx00100` at simulation time 25.0. These examples are included only for smoke testing and inference demonstrations; use the public HEAT URL for the full training data.

## Training Procedure

The source checkpoint metadata records:

- epoch: `533`
- train loss: `0.000961887405287`
- validation loss: `0.00133784566235`
- best validation loss: `0.00101556984417`
- parameters: `10,722,177`
- model state tensors: `118`

The architecture uses a conditional UNet with two input channels, one output channel, base channel width 64, sinusoidal timestep embeddings with dimension 256, three downsampling blocks, a middle attention block, and three upsampling blocks. The scheduler uses 1000 DDPM timesteps with beta range `1e-4` to `0.02`; inference uses DDIM sampling, defaulting to 50 steps.

During training, each batch samples a random diffusion timestep, adds Gaussian noise to the normalized target frame, and trains the UNet to predict that noise conditioned on the normalized first frame. Validation uses the same denoising loss procedure without gradient updates.

### Reproducibility Information (optional)

- Random seed used for data split: `42`
- Machine/environment info: GPUs with at least 30 GB of GPU memory for comparable training runs.
- Public training data URL: `https://oceans11.lanl.gov/heat/pli/pli240420_fp16_full/`
- Link to evaluation or inference pipeline: `api_ddpm.py`
- Published weights: `model.safetensors`
- Demo data: `examples/id01040`, `examples/id01045`, and `examples/id01050`
- Exact train/validation/test folder lists are not included in this release; reproduce the split by applying the documented seed and 70/15/15 folder-level split to the public PLI fp16 dataset.

## Pre-training information

- Hyperparameter tuning: not specified in checkpoint metadata.
- Model initialization: random initialization from the PyTorch module defaults used by the training notebook.
- Fine-tuning, chain-of-thought, and n-shot learning: not applicable.
- Optimizer information: `AdamW`, learning rate `1e-4`; optimizer state existed in the source `.pt` checkpoint but is intentionally not included in the Hugging Face release artifact.
- Loss function: denoising diffusion training loss, reported as training and validation MSE-style loss in checkpoint metadata.
- Stopping criterion: planned training length was `800` epochs; this card reports the published source checkpoint at epoch `533` and best validation loss stored in that checkpoint.
- Number of training epochs: `533` checkpoint epoch from a planned `800` epoch run.
- Number of training steps: not specified in checkpoint metadata.
- Batch size: `8`.
- Training speed: not specified in checkpoint metadata.
- Learning-rate scheduler: `CosineAnnealingLR`, `T_max=800`.
- Gradient clipping: global norm `1.0`.
- Training loss: `0.000961887405287`.
- Prompting templates/functions: not applicable.
- Model optimization techniques: no distillation, quantization, pruning, or sparsity method is documented for this checkpoint.

# Evaluation details

Evaluation details are documented from the packaged demo examples and checkpoint validation metadata.

## Evaluation data  

The packaged demo examples use HEAT-PLI folders `id01040`, `id01045`, and `id01050`, with `av_density` arrays of shape `(1120, 400)`. The input example is `examples/id01040/pli240420_id01040_pvi_idx00000.npz`, and the corresponding target endpoint is `examples/id01040/pli240420_id01040_pvi_idx00100.npz`.

## Evaluation Procedure

Run the model on the first frame and compare the generated result to the last frame. The training notebook used denoising validation loss during training and also included exploratory DDIM evaluations with PSNR and SSIM for sampled predictions. DDIM step counts explored in the notebook included `10`, `25`, `50`, `70`, `100`, and `150`; the published API defaults to `50` steps.

## Uncertainty Quantification.

Uncertainty/variability can be estimated by running `predict_ensemble` or the CLI `--n_samples` option. The API returns an ensemble mean and standard deviation across stochastic DDIM samples. The default single-sample inference does not provide uncertainty estimates.

## Evaluation results

Checkpoint validation MSE from source training metadata:

| Model | Direction | Validation loss | Best validation loss |
| --- | --- | ---: | ---: |
| Janus-first2last | first-to-last prediction | 0.00133784566235 | 0.00101556984417 |

Evaluation runtime depends on device, field size, `--num_steps`, and `--n_samples`. The reported runs used an A100 GPU.

# More Information (optional)

None.