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.1 --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. The checkpoint contains model weights and configuration for inference and evaluation. Optimizer state is not included.

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

The tables list results presented in the paper, also available in paper_results.json. The validation partition was also used to select the checkpoint. Pulse MAE uses target alignment and the normalization described in the paper.

The paper also evaluates ten experimental TG-FROG measurements, averaged across paired seeds 11, 29, and 47:

Method	Mean trace NRMSE ↓	Mean trace SSIM ↑	Mean amplitude roughness ↓
COPRA	0.013481	0.910208	0.179120
PnP-COPRA	0.014060	0.925819	0.159540
COPRA Flow	0.014705	0.919202	0.087340

Input domain

The model is trained on synthetic TG-FROG traces at 800 nm with 1% Gaussian trace noise and nonnegative clipping. Use the documented physical grid and preprocessing when preparing measurements.

Citation

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