A spec-to-bitmap compiler for diffractive optical elements (DOEs). Input is a design spec (wavelength, input beam, target intensity distribution, substrate, fab-grid pitch); output is a fab-ready 8- or 16-bit grayscale phase bitmap plus a JSON metadata sidecar. Single-process Dash app with in-process JAX compute; runs on CPU or a CUDA12 GPU.

v0 prototype. The pipeline runs end-to-end for three element classes, at different levels of maturity:

1. Axicon (Bessel beam, ring focus). Closed-form analytical. Production-ready.
2. Top-hat beam shaper. Optimal-transport init + MRAF polish. Working.
3. 9×9 spot array. Multi-seed Gerchberg-Saxton. Placeholder algorithm; commercial-grade upgrade pending.

See Capabilities for measured numbers.

In:

• Three element classes (axicon, top-hat, 9×9 spot array).
• Spec validation, routing, design, propagation verification, quantization, fab-grid resampling, bitmap output, JSON metadata, far-field preview.
• Single-process Dash UI with in-process JAX compute; docker compose for GPU.

Out:

• GDSII export (direct-write consumes bitmap; can add via gdstk later if needed).
• Tolerancing beyond a lite ±decentre / ±defocus sweep.
• CGH null-test for an arbitrary asphere.
• Full inverse-design library training.

What the three element classes do today and how well, measured on the canonical specs at 8-level quantization:

Quality numbers come from the bench regression suite. See Known limitations for the spot-array caveat.

Built from scratch rather than forking an existing library. Mature primitives (numpy, JAX, Pillow, Dash) are used directly. The reference projects below are read for algorithms and used as dev-time cross-checks, not taken as runtime dependencies; where a function is materially adapted from one, it is attributed in NOTICE.md.

The compute path is device-agnostic JAX: CPU by default, CUDA12 GPU via an install extra or the Docker image. Start on CPU to try it; move to GPU for production-scale compiles.

Clone the repo, then from its root:

pip install -e '.[ui]'      # core pipeline + Dash dashboard, CPU JAX
python -m ui.dashboard.app  # open http://localhost:8051

Pick a canonical spec from the dropdown and hit Compile: out comes a fab-ready grayscale bitmap, a far-field preview, and verification metrics. DRAFT compiles of small specs run in seconds on CPU.

CPU caveats (this is why the GPU path exists): STANDARD / HIGH presets and large grids take minutes-to-hours on CPU; large-output specs can OOM; and on Windows the JAX persistent compile-cache writer fails (Failed to serialize CpuAotCompilationResult), so every fresh shape pays a full cold compile. For real compute, use the GPU path below (pip install -e '.[ui,cuda]' on a CUDA-capable host, or the Docker image).

Two documented run modes. Same Docker image (Dockerfile.jax-gpu); pick by what you're doing.

Both require nvidia-container-toolkit and a CUDA-capable GPU. JAX is a hard dependency; the compute path runs JAX-CUDA12 inside the container.

On Windows: Docker runs inside WSL2 Ubuntu (Docker Desktop has a chronic broken-Unix-socket bug that no Windows tool can clear). See docs/docker-wsl-setup.md for one-time setup.

Both services share Dockerfile.jax-gpu. Build the image once:

docker compose --profile gpu build gpu

# Dashboard with GPU compile_spec
docker compose --profile gpu up dashboard
# http://localhost:8051

# Long-lived shell for bench scripts and algorithm dev
docker compose --profile gpu up -d gpu
docker compose exec gpu bash
# inside the container:
python scripts/bench_iter_sweep.py

Cold PTX-JIT on Blackwell sm_120 is 8-13 min per distinct trace; the persistent cache makes every subsequent invocation of an already-warmed trace instantaneous. A trace is keyed by (function bytecode, input shapes/dtypes, static args, JAX/XLA version, GPU model). What this means in practice:

• Changing scalar spec values (wavelength, beam diameter, output size in mm, working distance, uniformity target, quantization 8 vs 16) → same trace, instant cache hit.
• Changing fab grid size (pitch or aperture), preset, element type, axicon mode, or top-hat output_shape → new trace, cold compile first time.
• Updating an algorithm function (e.g. polish.py) → invalidates affected traces; recompile next call.

Cache lives at ${HOME}/.cache/jax on the WSL2 host, bind-mounted into the container at /cache/jax.

Prewarm at startup. When the dashboard boots it fires DRAFT axicon

• DRAFT tophat-rect on tiny placeholder grids in a background thread to warm the most-clicked paths. STANDARD and HIGH presets are not prewarmed (they take 5+ min each, would block startup); first click on those pays cold compile, then is warm forever. Set DOE_COMPILER_PREWARM=0 to skip prewarm entirely.

Bulk-warm the full matrix. scripts/warm_cache.py iterates every (canonical spec, preset) combination and runs compile_spec on each once. After this lands, every spec in the catalog dropdown clicks instantly across DRAFT/STANDARD/HIGH. Run once after a fresh checkout or after an algorithm change:

docker compose exec gpu python scripts/warm_cache.py
# 42 combos at ~10 min cold each on Blackwell first time;
# subsequent runs are seconds.

• Spot-array ships a placeholder algorithm. The dashboard exposes a spot-array element with an end-to-end compile path, but the underlying algorithm is plain Wyrowski-mixed-region GS, which lands at ~10-30% spot-to-spot RMS vs the <5% commercial vendors quote. Architecture is correct (canonical specs, dashboard wiring, verify metrics); commercial-grade WGS upgrade is the next big algorithm push.
• Padded top-hat MRAF trades ~3 pp efficiency for uniformity. Tight (small-output) specs run MRAF on a 2× zero-padded grid to close the natural-vs-zoom uniformity gap. That routes ~18% of input energy into the noise region / unfabricated annulus, so delivered signal efficiency on those specs is ~0.82 rather than the ~0.85 the unpadded design suggests. Uniformity is honest (measured on the real fab grid by verify); the efficiency cost is real. Recovering it (aperture-aware target, or a tighter MRAF schedule) is future algorithm work.
• Single compile path. Axicon, top-hat, and spot-array all route through compile_spec. The dashboard's compile_axicon calls compile_spec and layers axicon-specific extras (axial scan, far-field diffraction-order separation, closed-form theory overlay) from src/doe_compiler/axicon_analytics.py on top.
• Extreme designs (effective_n >= 16384) on a ~10 GB GPU stay at full quality. All FFT-bearing kernels (angular_spectrum, MRAF polish, IFTA propagator pair, antialias-lowpass) dispatch through the three-tier tiled-FFT primitive in src/doe_compiler/tiled_fft.py: in-VRAM single-shot for small grids, GPU-resident tiled with donate_argnums for grids that fit on device but blow kernel fusion, host-streamed tiled for grids past the device single-buffer ceiling. complex128 throughout. Sub-micron-pitch (1 um pitch / 10 mm aperture, design_n=16384) specs compile end-to-end without the earlier precision compromise. Telemetry surfaces under metadata["design"]["polish"]["fft_backend"]; the dashboard's RUN block shows a "tiled-FFT path" badge when a non-trivial tier engaged.
• Axicon also pipes through a jax-resident chain. design_axicon builds the phase on the device; antialias, resample, quantize, and the verify-side propagation all dispatch numpy-vs-jax on isinstance so the buffer stays on device through the whole chain (fab-grid phase_quant materializes to numpy at the end for the bitmap save). On the RTX 5070, N=8192 design grid compiles in ~9 s. N=32768 exceeds VRAM (single complex128 buffer is 16 GB); a tiled design + antialias path would be the follow-up if sub-um pitch / >10 mm aperture parts are needed. Top-hat / spot-array stay on the numpy path.
• Run database persists across compiles. Every successful and failed compile lands a row in runs/runs.db (in-tree, gitignored). Schema in ui/dashboard/runs_db.py. Dashboard's catalog modal + Compile button popover both query this DB; Session History drawer too. Rows accumulate without pruning (they are small, since the DB stores metrics and form state, not phase masks).

Read for reference, not forked. Where a license permits, individual functions may be adapted with attribution in NOTICE.md.

pip install -e '.[dev,ui]'   # adds pytest, ruff, scipy, LightPipes
pytest                       # full suite (CPU), ~4 min
ruff check && ruff format

AGPL-3.0-only. See LICENSE.