atal-narayan/copra-flow

COPRA Flow

COPRA Flow refines ultrafast laser pulse reconstructions from transient-grating frequency-resolved optical gating (TG-FROG) measurements. A conditional flow model takes a COPRA reconstruction and the measured trace, then produces a refined complex electric field.

• Code and installation: atalnarayan/copra-flow
• Training data: COPRA Flow TG-FROG
• Weights: copra-flow.pt, PyTorch checkpoint, approximately 7 MB
• License: MIT

Run the model

Install the code following the GitHub README. From the code repository, download this checkpoint with the Hugging Face CLI:

hf download atal-narayan/copra-flow copra-flow.pt --revision v0.1.0 --local-dir checkpoints
uv run copra-flow example --out data/example.npz
uv run copra-flow infer --checkpoint checkpoints/copra-flow.pt \
  --input data/example.npz --out-dir runs/example

The command performs 400 COPRA iterations and 64 Euler flow refinement steps. It writes prediction.npz and metrics.json. No training data or ground-truth pulse is required for inference. See the inference guide for preparing experimental measurements or refining an existing COPRA field.

Inputs use 256 temporal samples at 2.5 fs spacing. The measurement must follow the documented frequency and delay conventions. This is a custom scientific PyTorch model; load it through copra-flow, rather than a Transformers pipeline.

Model and training

The model has 1,754,770 parameters. It predicts a conditional vector field using two-channel real/imaginary pulse fields and a peak-normalized, square-root-transformed TG-FROG trace. The architecture uses 128 hidden channels, eight residual blocks, a 256-dimensional condition embedding, and a 128-dimensional time embedding.

Training used 270,000 source–target pairs from 9,000 synthetic pulses, with 30,000 pairs from 1,000 disjoint pulses for validation. All source candidates for a pulse stay in the same partition. The run used AdamW, learning rate 0.0003, batch size 64, and 50 epochs on an NVIDIA A100 40 GB GPU with mixed precision. The released checkpoint was selected by validation loss at epoch 46, update 194,074; its validation flow loss was 0.017534218672072208.

config.json contains the training and architecture settings. splits.json contains the fixed pulse split, and training_history.json records all 50 epochs. This release preserves every model tensor exactly, removes optimizer state, and replaces machine-specific paths with portable relative paths. It is intended for inference and evaluation, not exact training resumption. Source and release hashes are recorded in release.json and SHA256SUMS.

Evaluation

The paper's synthetic benchmark evaluates the final COPRA source for each of 1,000 validation pulses with 64 Euler steps:

Method	Mean pulse MAE ↓	Mean trace SSIM ↑
COPRA	0.4685	0.9709
COPRA + direct Hann	0.3345	0.8207
COPRA Flow	0.0674	0.9103

These are validation results, not results on an independent held-out test set. The same validation partition selected the checkpoint. Pulse errors use the synthetic benchmark's target alignment and normalization. A complete rerun matched the six reference metric means within an absolute tolerance of 0.0001; see validation_results.json and the code's benchmark script.

Scope and limitations

Training uses synthetic TG-FROG traces at 800 nm with 1% additive Gaussian trace noise and nonnegative clipping. Performance on other pulse distributions, noise conditions, acquisition systems, or physical grids is not established. Improved pulse error does not guarantee improved trace consistency; check the reconstructed trace when interpreting a result. Smoothness alone does not establish pulse accuracy.

The experimental inference workflow executes, but the historical experimental metric table does not reproduce exactly, including when rerunning the original research scripts. See the reproducibility notes for the observed differences and verification limits. This release does not claim an exact reproduction of that table.

Citation

Citation metadata is maintained in the code repository's CITATION.cff.