Torchoptics

Differentiable wave optics simulation library built on PyTorch

Generate Convert Improve

Install / Use

/learn @MatthewFilipovich/Torchoptics

About this skill

Quality Score

0/100

README

</div>

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:

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:

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 = 1000
f = 200e-3

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

system = System(
    Lens(shape, f, z=1 * f),
    AmplitudeModulator(1 - circle(shape, radius=5e-4), 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:

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},
}