fouriax

Differentiable free-space optics for JAX.

Warning

This project is vibecoded. It’s primarily developed by collaborating with AI coding assistants. While we strive for correctness, expect experimental features, rapid iterations, and potentially non-traditional implementation structures.

fouriax is a JAX library for simulating and optimizing coherent and incoherent optical systems with automatic differentiation. It provides composable optics layers, explicit spatial and k-space transforms, propagation planning, Jones polarization, meta-atom lookup tables, and sensor readout for gradient-based inverse design.

The source tree is organized with the core optics runtime in src/fouriax/optics/ and the higher-level helper modules at src/fouriax/analysis.py and src/fouriax/optim.py.

Sections

• Guides: user-facing explanations and conceptual documentation.

• API Reference: generated reference pages for the core library surface.

• Examples: synced example notebooks rendered as documentation pages.

• Development: local setup, testing, and contributor workflow.

Key Features

• Composable optical stacks built around OpticalLayer and OpticalModule

• Explicit spatial and k-space transitions via Fourier transform layers

• Multiple propagation backends with automatic planning

• Jones polarization support

• Meta-atom interpolation layers

• Incoherent imaging and detector models

• Differentiable optimization workflows built on JAX and Optax

• Guides
  • Fouriax User Guide
• API Reference
  • Sections
• Examples
  • Basic Propagation — First Optical Stack
  • Diffractive Lens Optimization — Focusing a Plane Wave
  • 4f Optical Correlator
  • Coherent Hologram — Phase-Only Logo Reconstruction
  • Incoherent Camera Imaging — PSF and OTF Modes
  • Optical Neural Network (ONN) on MNIST — Phase-Mask Classifier
  • 4f Edge Detection — Learning the Vortex Phase Filter
  • Polarization-Multiplexed Holography — Dual Reconstructions
  • Meta-Atom Library Optimization — Focusing with Interpolated Unit Cells
  • Sensitivity Analysis and Fisher Information for a Diffractive Lens
• Development
  • Local Setup
  • Common Commands