README.md

---
license: mit
tags:
  - physics
  - electromagnetic-simulation
  - flow-matching
  - photonics
  - silicon-photonics
  - fdtd
  - pde
library_name: pytorch
---

# PIC-Flow

A physics-embedded flow-matching neural surrogate that replaces FDTD for full-field
electromagnetic prediction of silicon photonic devices. Given a permittivity map
ε(x,y), a source-port mask, and a free-space wavelength λ, PIC-Flow generates the
complex field E_z in a single multi-step ODE integration — typically in well under
a second on a single A100, vs. seconds-to-minutes for CPU FDTD.

This repo hosts the **FM + phase + residual** checkpoint from epoch 300 (the headline
model from the paper). All training code, dataset-generation tooling, and inference
notebooks live in the GitHub repo:
[Rizzo-Integrated-Photonic-Systems-Lab/PIC-Flow](https://github.com/Rizzo-Integrated-Photonic-Systems-Lab/PIC-Flow).

## Files

| Path | Description |
|---|---|
| `checkpoints/phase_residual_300.pt` | FM+phase+residual U-Net, epoch 300, ~1 GB. |

## Quick usage

```bash
pip install huggingface_hub torch numpy
```

```python
from huggingface_hub import hf_hub_download
import torch

ckpt_path = hf_hub_download(
    "RizzoLab/PIC-Flow",
    "checkpoints/phase_residual_300.pt",
)
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
# ckpt["state_dict"] -> model weights (real-valued U-Net, 63.3M params)
# ckpt["stats"]      -> field/permittivity normalization stats
# ckpt["args"]       -> training hyperparameters
```

End-to-end inference (load model, build conditioning, run flow-matching sampler) is
covered by [`tools/predict_parametric_device.py`](https://github.com/Rizzo-Integrated-Photonic-Systems-Lab/PIC-Flow/blob/main/tools/predict_parametric_device.py)
and [`notebooks/03_inference.ipynb`](https://github.com/Rizzo-Integrated-Photonic-Systems-Lab/PIC-Flow/blob/main/notebooks/03_inference.ipynb)
in the GitHub repo.

## Model

- **Architecture**: real-valued U-Net, 63.3M parameters. Real and imaginary E_z
  components enter as separate input channels; the permittivity and source-mask maps
  are visible at every layer; the flow-matching integration time `t` and the wavelength
  `λ` enter as scalar conditioning inputs.
- **Generative framework**: conditional flow matching (Lipman et al., 2023). Inference
  integrates a learned velocity field from Gaussian noise to a physically valid E_z
  using Euler or Heun ODE steps.
- **Physics constraint**: masked Helmholtz residual loss `L_res` (PML, source, and
  dielectric-interface pixels excluded), with a per-sample compliance metric
  `ρ_R = sqrt(L_res) × 100%`.

## Training data

- 22,500 Meep FDTD simulations at λ = 1.55 µm
- Three device families: 2×2 MMIs, Y-branches, directional couplers (7,500 each)
- Latin-hypercube parameter sweeps over geometric variables per family
- 18,000 / 2,250 / 2,250 train / val / test split

Training: 300 epochs on 12 NVIDIA V100 GPUs, identical hyperparameters across the three
ablation runs (FM only, FM+phase, FM+phase+residual).

## Performance

On the held-out test split (200-step Heun sampler):

| Device family | ρ_R |
|---|---|
| 2×2 MMI | 2.7% |
| Y-branch | 2.5% |
| Directional coupler | 2.2% |

Out-of-distribution (same checkpoint, geometries never seen during training):

| Device | ρ_R |
|---|---|
| Aggressive Euler S-bend (tight R, large offset) | 12% |
| Short, steep taper | 4.0% |
| Long, wide taper | 3.6% |
| Cascaded 1×3 Y-branch (new device class) | 9.1% |

Wall clock on a single NVIDIA A100 (fp16 autocast, vs. 16-thread Meep FDTD on the
same node):

| Sampler | Wall time | Speedup | ρ_R |
|---|---|---|---|
| FDTD (reference) | 5.61 s | 1.0× | (reference) |
| Euler 100 step | 2.19 s | 2.6× | 1.9% |
| Euler 20 step | 440 ms | 12.7× | 3.0% |
| Euler 5 step | 110 ms | 50.6× | 5.5% |

## Citation

```bibtex
@article{Quaratiello2026PICFlow,
  author  = {Joseph Quaratiello and Anthony Rizzo},
  title   = {A Physics-Embedded Flow-Matching Model for Electromagnetic Prediction
             of Silicon Photonic Devices},
  journal = {arXiv},
  year    = {2026}
}
```

## License

MIT. See [LICENSE](https://github.com/Rizzo-Integrated-Photonic-Systems-Lab/PIC-Flow/blob/main/LICENSE)
in the GitHub repo.