CODEBASE MAP#

Developer-facing map of the library layout, runtime flow, and edit boundaries. Use this for orientation. Use tests and the source for exact behavior.

Top-Level Layout#

  • src/fouriax/__init__.py: top-level re-exports for the public API

  • src/fouriax/optics/: core optics runtime

  • src/fouriax/analysis.py: Fisher information and sensitivity helpers

  • src/fouriax/optim.py: batching, loss helpers, and Optax training loops used by examples

  • src/fouriax/fft.py: FFT-based convolution helpers

  • src/fouriax/linop.py: lightweight linear-operator wrapper

  • examples/scripts/: source-of-truth examples

  • examples/notebooks/: synced notebooks

  • tests/: pytest suite

Optics Package#

  • src/fouriax/optics/grid.py Defines Grid, AngularGrid, and Spectrum.

  • src/fouriax/optics/field.py Defines Field and Intensity. Also holds field-shape, domain, and broadcasting helpers shared by layers.

  • src/fouriax/optics/interfaces.py Defines the abstract surfaces: OpticalLayer, IncoherentLayer, Sensor, Monitor.

  • src/fouriax/optics/layers.py Contains scalar-domain layers and coherent composition: OpticalModule, masks, ThinLens, FourierTransform, InverseFourierTransform.

  • src/fouriax/optics/polarization.py Jones-specific layers: JonesMatrixLayer, KJonesMatrixLayer.

  • src/fouriax/optics/imaging.py Incoherent imaging runtime: IncoherentImager, PSF/OTF convolution, cached linear operators.

  • src/fouriax/optics/propagation.py Concrete propagators and propagation planning: RSPropagator, ASMPropagator, KSpacePropagator, plan_propagation, build_na_mask, regime helpers.

  • src/fouriax/optics/sensors.py Readout layers: Detector, DetectorArray, detector linear operators.

  • src/fouriax/optics/noise.py Analytic sensor-noise models.

  • src/fouriax/optics/meta_atoms.py LUT-backed meta-atom interpolation layers.

  • src/fouriax/optics/monitors.py Non-mutating observation helpers for traces and debugging.

Main Runtime Flows#

Coherent Path#

Field -> OpticalModule -> layers/polarization/propagation -> Field

  • Spatial and k-space are explicit.

  • Domain conversion only happens through FourierTransform, InverseFourierTransform, Field.to_kspace(), or Field.to_spatial().

Incoherent Path#

Field.to_intensity() -> IncoherentImager -> Intensity -> DetectorArray

  • IncoherentImager is shift-invariant and scalar-only.

  • Detector-grid resampling belongs in DetectorArray, not in the imager.

Planning Path#

plan_propagation(...) -> concrete propagator instance

  • plan_propagation resolves method selection and precomputed transfer data.

  • NA masking belongs in propagation.py.

Ownership Rules#

  • Put sampling primitives in grid.py.

  • Put field containers and field-shape helpers in field.py.

  • Put scalar modulation and module composition in layers.py.

  • Put Jones-only logic in polarization.py.

  • Put incoherent imaging logic in imaging.py.

  • Put propagators and propagation selection in propagation.py.

  • Put detector readout and detector-grid resampling in sensors.py.

  • Keep analysis.py and optim.py outside optics/; they are helpers built on top of the runtime.

Invariants#

  • Domains are explicit: no hidden spatial/k-space inference.

  • Shape order is: *batch, wavelength, [jones], y, x

  • Scalar fields use trailing shape (wavelength, y, x).

  • Jones fields use trailing shape (wavelength, 2, y, x).

  • Layers validate domain explicitly and fail on the wrong domain.

  • Propagation method selection is explicit at construction time, not hidden inside unrelated layers.

Public API Mapping#

  • fouriax re-exports the main user-facing surface from fouriax.optics, plus analysis, optim, fftconvolve, and LinearOperator.

  • fouriax.optics is the internal package boundary for the runtime.

  • Internal file layout may move; preserve the public fouriax.* surface unless the change is intentionally breaking.

Edit Guide#

  • If changing field shape, domain, or metadata: read src/fouriax/optics/field.py, tests/test_field_domain.py, tests/test_jones_polarization.py

  • If changing coherent layer composition or domain transitions: read src/fouriax/optics/layers.py, tests/test_optical_module.py, tests/test_k_layers.py

  • If changing Jones behavior: read src/fouriax/optics/polarization.py, tests/test_jones_polarization.py

  • If changing propagation or NA behavior: read src/fouriax/optics/propagation.py, tests/test_propagation_selection.py

  • If changing incoherent imaging: read src/fouriax/optics/imaging.py, tests/test_incoherent_propagation.py

  • If changing detector behavior: read src/fouriax/optics/sensors.py, tests/test_camera_sensor.py

  • If changing meta-atom interpolation: read src/fouriax/optics/meta_atoms.py, tests/test_meta_atoms.py

Example Workflow Boundary#

  • Treat examples/scripts/*.py as source of truth.

  • If an example has a synced notebook, update the script first and sync the notebook.

  • Example-facing training helpers belong in src/fouriax/optim.py, not in optics/.