| --- |
| 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). |
|
|
