README.md

---
license: mit
language:
- en
tags:
- physics
- optics
- ultrafast-optics
- tg-frog
- pulse-retrieval
- flow-matching
- pytorch
datasets:
- atal-narayan/copra-flow-tg-frog
---

# 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](https://github.com/Quantum-Light-Matter-Cooperative-QLMC/copra-flow)
- **Training data:** [COPRA Flow TG-FROG](https://huggingface.co/datasets/atal-narayan/copra-flow-tg-frog)
- **Weights:** [`copra-flow.pt`](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:

```bash
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](https://github.com/Quantum-Light-Matter-Cooperative-QLMC/copra-flow/blob/main/docs/inference.md)
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](https://github.com/Quantum-Light-Matter-Cooperative-QLMC/copra-flow/blob/main/CITATION.cff).