Optical Neural Front-End (ONF) System

# SPEC.md — Optical Neural Front-End (ONF) System

> Generated by the Gerolamo compose engine (mode=compose) from meta molecule
> `b595b0c5-544d-4a21-8a60-76664ecec226`. The engine's output was capped at the
> start of Phase 3 (end-to-end training); Phases 3+ continue the same structure.

## Project Overview

This system implements a co-designed optical neural front-end where a reconfigurable phase-change metasurface physically instantiates the first convolutional layer(s) of a vision model, performing feature extraction passively in the optical domain before any digital computation occurs. The metasurface acts as a programmable filter bank — its phase profile encodes learned spatial frequency decompositions — and is trained end-to-end with a lightweight digital backbone using differentiable physical optics simulation. The result is a power-bounded edge vision system where the most compute-intensive front-end operations cost zero digital watts, with the metasurface switchable between task-specific configurations by thermally actuating the phase-change material. This solves the inference power wall on always-on edge vision devices (cameras, drones, IoT sensors) where a full CNN front-end is thermally and energetically untenable.

---

## Architecture

### Component Roles

| Artifact Component | System Role |
|---|---|
| Phase-change metasurface | Physical instantiation of learned conv filters; passive optical compute |
| Free-space optical path | Signal propagation medium; encodes the matrix-vector multiply as wavefront transformation |
| End-to-end differentiable optics | Training bridge between physical optics sim and digital backbone loss |
| Photonic convolution engine | Mathematical core mapping metasurface phase profiles → spatial convolution outputs |
| Edge vision backbone | Lightweight digital CNN consuming metasurface feature maps as input |
| Transfer-function-aware training | Loss augmentation that accounts for metasurface fabrication and switching constraints |

### System Data Flow

```
Scene (photons)
      │
      ▼
┌─────────────────────────────────────────┐
│  METASURFACE LAYER (Physical)           │
│  Phase profile Φ(x,y) encodes W_conv1   │
│  GST/GSST unit cells, ~λ/5 pitch        │
│  Switchable via thermal/electrical pulse │
│  Output: modulated wavefront            │
└─────────────────────────────────────────┘
      │  (free-space propagation, ~mm)
      ▼
┌─────────────────────────────────────────┐
│  SENSOR ARRAY (CMOS/InGaAs)            │
│  Captures intensity |E(x,y)|²          │
│  Each pixel region = one conv output    │
│  Effective: tiled spatial conv result   │
└─────────────────────────────────────────┘
      │  (feature map tensor, not raw pixels)
      ▼
┌─────────────────────────────────────────┐
│  DIGITAL BACKBONE (Edge CPU/NPU)        │
│  Receives pre-extracted features        │
│  MobileNetV3 / EfficientNet-Lite trunk  │
│  Layers 2-N only (Layer 1 is optical)   │
└─────────────────────────────────────────┘
      │
      ▼
   Task Output (classification / detection / segmentation)
```

### Integration Points

1. **Optical↔Digital Interface**: Sensor readout maps to a `(B, C_optical, H', W')` tensor where `C_optical` = number of distinct phase profiles multiplexed (angular, polarization, or time-multiplexed channels).
2. **Training Interface**: `OpticalConvLayer` is a drop-in `nn.Module` whose `forward()` calls the differentiable PSF simulator; at deployment the weights are compiled to phase maps.
3. **Phase Compiler**: Converts trained weight tensor `W ∈ R^{C_out × C_in × k × k}` → phase profile `Φ ∈ [0, 2π]^{N×N}` via inverse design optimization.
4. **Task Switching**: A `MetasurfaceController` sends thermal pulse sequences to transition unit cells between amorphous/crystalline states, loading one of K stored phase profiles.

---

## Tech Stack

### Simulation & Training
- **Python 3.11**, **PyTorch 2.3.0**, **pytorch-lightning 2.2.4**, **numpy 1.26.4**, **scipy 1.13.0** (Fourier optics, angular spectrum), **opencv-python 4.9.0**

### Differentiable Optics
- **dflat 0.4.0** — differentiable flat optics / metasurface simulation (PyTorch-native)
- **neurophox 0.2.2** — photonic mesh / transfer matrix utilities (reference)
- **meent 0.6.0** — RCWA for unit cell S-parameter simulation
- **torch-scatter 2.1.2** — sparse phase map operations

### Phase Compiler & Inverse Design
- **nlopt 2.7.1** (phase retrieval), **shapely 2.0.4** (unit cell geometry), **gdspy 1.6.13** (GDS-II export), **h5py 3.11.0** (phase profile storage)

### Backbone Model
- **timm 0.9.16** (MobileNetV3, EfficientNet-Lite), **torchvision 0.18.0**

### Edge Deployment
- **onnx 1.16.0** + **onnxruntime 1.18.0**, **ai-edge-torch 0.1.1** (TFLite / NPU), **pyserial 3.5** (metasurface controller UART)

### Hardware Interface
- **nidaqmx 0.9.0** (NI DAQ pulse generation), **thorlabs-apt** (motorized mount control)

### Testing & Visualization
- **pytest 8.2.0**, **matplotlib 3.9.0**, **plotly 5.22.0**, **wandb 0.17.0**

---

## Project Structure

```
onf-system/
├── SPEC.md
├── CLAUDE.md
├── README.md
├── pyproject.toml
│
├── onf/
│   ├── optics/
│   │   ├── propagation.py            # Angular spectrum / Fresnel propagation
│   │   ├── metasurface.py            # Phase mask forward model, unit cell library
│   │   ├── psf.py                    # Point spread function computation
│   │   ├── coherence.py              # Partial coherence for broadband sources
│   │   └── rcwa_interface.py         # meent wrapper for unit cell sims
│   ├── layers/
│   │   ├── optical_conv.py           # OpticalConvLayer — differentiable metasurface conv
│   │   ├── sensor.py                 # intensity detection + noise model
│   │   ├── phase_constraint.py       # quantization, smoothness, amplitude regularizers
│   │   └── multiplexing.py           # polarization / angle channel multiplexing
│   ├── backbone/
│   │   ├── truncated_backbone.py     # load timm model, strip layer 1, attach optical head
│   │   ├── full_pipeline.py          # ONFModel: optical layer + digital backbone e2e
│   │   └── task_heads.py
│   ├── compiler/
│   │   ├── phase_retrieval.py        # GS + gradient-based phase retrieval
│   │   ├── unit_cell_lut.py          # phase → geometry LUT for GST cells
│   │   ├── layout_export.py          # phase maps → GDS-II
│   │   └── switching_sequence.py     # thermal pulse sequences for K tasks
│   ├── training/{trainer.py, loss.py, curriculum.py, augmentation.py}
│   ├── controller/{thermal_driver.py, profile_manager.py, mock_controller.py}
│   └── utils/{visualization.py, metrics.py, config.py}
│
├── configs/{base.yaml, imagenet_mobilenet.yaml, coco_detection.yaml, edge_power_constrained.yaml}
├── scripts/{train.py, compile_phase.py, export_backbone.py, eval_optical.py, switch_task.py}
├── data/{unit_cell_library/, phase_profiles/}
├── tests/{test_propagation.py, test_optical_conv.py, test_phase_retrieval.py,
│          test_full_pipeline.py, test_controller_mock.py, fixtures/}
└── notebooks/{01_unit_cell_design.ipynb ... 04_phase_compiler_demo.ipynb}
```

---

## Build Plan

### Phase 1 — Physical Optics Simulation Core
**Objective:** A verified, differentiable free-space optical propagation engine and metasurface phase-mask forward model.

**Tasks:**
1. `onf/optics/propagation.py`: `angular_spectrum_propagate(field, z, wavelength, pixel_pitch)` and Fresnel fallback; PyTorch complex tensors with `requires_grad=True`.
2. `onf/optics/metasurface.py`: `Metasurface` holding `Φ ∈ [0,2π]^{N×N}` as `nn.Parameter`; `apply_phase_mask(...)`; HDF5 unit-cell library loader (GST transmission amplitude+phase).
3. `onf/optics/psf.py`: `compute_psf(...)` intensity PSF at sensor; tiled PSF for spatially-varying aberrations.
4. `onf/optics/coherence.py`: broadband PSF as incoherent sum over wavelength bins (450/550/650nm).
5. Generate `data/unit_cell_library/` via `meent` RCWA over a 5×5 GST geometry sweep.

**Verification:**
```bash
pytest tests/test_propagation.py -v
# Energy conservation within 1%; flat phase → Airy disk; gradcheck passes.
```

### Phase 2 — Optical Conv Layer & Sensor Model
**Objective:** Wrap the optics engine as an `nn.Module` that behaves like a conv layer during training.

**Tasks:**
1. `onf/layers/optical_conv.py`: `OpticalConvLayer(in_channels, out_channels, kernel_size, wavelength, z, sensor_pitch)`; forward = illuminate → metasurface → propagate → sensor intensity; gradients flow through `Φ`.
2. `onf/layers/sensor.py`: shot noise (Poisson), read noise (Gaussian), 8-bit quantization; differentiable via reparameterization.
3. `onf/layers/phase_constraint.py`: quantization (toward 256 levels), TV smoothness (penalize sub-λ/2 variation), amplitude regularizer.
4. `onf/layers/multiplexing.py`: polarization and oblique-incidence channel multiplexers.

**Verification:**
```bash
pytest tests/test_optical_conv.py -v
# (B,1,H,W) → (B,C_out,H',W'); flat phase → uniform output; non-zero ∂/∂Φ.
```

### Phases 3+ (capped by the compose engine — same structure continues)
- **Phase 3 — Backbone integration & end-to-end training:** strip layer 1 of a `timm` backbone, attach the optical head, train the full `ONFModel` with task loss + optical regularizers; curriculum-anneal fabrication constraints.
- **Phase 4 — Phase compiler:** trained `W` → manufacturable `Φ` via phase retrieval + unit-cell LUT; export GDS-II; generate K-task switching sequences.
- **Phase 5 — Hardware bring-up:** `thermal_driver.py` over serial/DAQ; `profile_manager.py` to store/switch K profiles; measured-transfer-function fine-tune.

## Risks & Open Questions
- **Sim-to-real gap:** train against the *measured* metasurface transfer function, not an idealized one, or accuracy collapses on hardware. The fine-tune step in Phase 5 is load-bearing.
- **Switching speed vs. lifetime:** phase-change actuation has finite cycle endurance; budget task-switch frequency against material lifetime.
- **Multiplexing ceiling:** `C_optical` is bounded by available polarization/angle/time channels; the digital backbone must do real work, so set front-end channel count realistically.