torchoptics

TorchOptics is an open-source Python library for simulating optical systems using Fourier optics, built on PyTorch. It provides GPU-accelerated, fully differentiable wave optics simulations, enabling end-to-end optimization of optical hardware jointly with machine learning models.

Learn more in our paper on arXiv.

Key Features

• 🌊 Differentiable Wave Optics: Model, analyze, and optimize optical systems using Fourier optics.
• 🔥 Built on PyTorch: GPU acceleration, batch processing, and automatic differentiation.
• 🛠️ End-to-End Optimization: Joint optimization of optical hardware and machine learning models.
• 🔬 Optical Elements: Lenses, phase/amplitude modulators, detectors, polarizers, and more.
• 🖼️ Spatial Profiles: Hermite-Gaussian, Laguerre-Gaussian, Zernike modes, gratings, and others.
• 🔆 Polarization and Coherence: Simulate polarized light and fields with arbitrary spatial coherence.

Installation

pip install torchoptics

Documentation

Full documentation is available at torchoptics.readthedocs.io.

Examples

Wave Propagation

Simulate free-space propagation of an octagonal aperture (full example):

import torch
import torchoptics
from torchoptics import Field
from torchoptics.profiles import octagon

device = "cuda" if torch.cuda.is_available() else "cpu"
torchoptics.set_default_spacing(10e-6)
torchoptics.set_default_wavelength(700e-9)

field = Field(octagon(shape=500, radius=150e-5)).to(device)

for z in torch.linspace(0, 2, 11):
    field.propagate_to_z(z).visualize(title=f"z = {z:.2f} m")

4f Imaging System

Simulate a 4f system with a high-pass spatial filter (full example):

import torch
import torchoptics
from torchoptics import Field, System
from torchoptics.elements import AmplitudeModulator, Lens
from torchoptics.profiles import checkerboard, circle

device = "cuda" if torch.cuda.is_available() else "cpu"
torchoptics.set_default_spacing(10e-6)
torchoptics.set_default_wavelength(700e-9)

shape = 500
f = 50e-3

input_field = Field(checkerboard(shape, tile_length=200e-6, num_tiles=15)).to(device)

system = System(
    Lens(shape, f, z=1 * f),
    AmplitudeModulator(1 - circle(shape, radius=200e-6), z=2 * f),
    Lens(shape, f, z=3 * f),
).to(device)

for i in range(5):
    system.measure_at_z(input_field, z=i * f).visualize(title=f"z={i}f", vmax=1)

Inverse Design

Train a diffractive optical system to convert a Gaussian beam into a petal beam (full example):

import torch
import torchoptics
from torch.nn import Parameter
from torchoptics import Field, System
from torchoptics.elements import PhaseModulator
from torchoptics.profiles import gaussian, laguerre_gaussian

device = "cuda" if torch.cuda.is_available() else "cpu"
torchoptics.set_default_spacing(10e-6)
torchoptics.set_default_wavelength(700e-9)

shape = 250
waist_radius = 300e-6

input_field = Field(gaussian(shape, waist_radius=waist_radius), z=0).to(device)

petal_profile = laguerre_gaussian(shape, p=0, l=4, waist_radius=waist_radius)
petal_profile += laguerre_gaussian(shape, p=0, l=-4, waist_radius=waist_radius)
target_field = Field(petal_profile, z=0.8).normalize().to(device)

system = System(
    PhaseModulator(Parameter(torch.zeros(shape, shape)), z=0.2),
    PhaseModulator(Parameter(torch.zeros(shape, shape)), z=0.4),
    PhaseModulator(Parameter(torch.zeros(shape, shape)), z=0.6),
).to(device)

optimizer = torch.optim.Adam(system.parameters(), lr=0.05)
for iteration in range(100):
    optimizer.zero_grad()
    output_field = system.measure_at_z(input_field, 0.8)
    loss = 1 - output_field.inner(target_field).abs().square()
    loss.backward()
    optimizer.step()

For more examples, see the examples gallery.

Contributing

Contributions are welcome! See the Contributing Guide for details.

Citing TorchOptics

If you use TorchOptics in your research, please cite our paper:

@misc{filipovich2024torchoptics,
      title={TorchOptics: An open-source Python library for differentiable Fourier optics simulations},
      author={Matthew J. Filipovich and A. I. Lvovsky},
      year={2024},
      eprint={2411.18591},
      archivePrefix={arXiv},
      primaryClass={physics.optics},
      url={https://arxiv.org/abs/2411.18591},
}

License

Distributed under the MIT License. See LICENSE for details.