torchoptics 0.4.0

Python library for differentiable Fourier optics simulations with PyTorch.

TorchOptics is a differentiable wave optics simulation library built on PyTorch.

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, modulators, detectors, polarizers, and more.
• 🖼️ Spatial Profiles — Hermite-Gaussian, Laguerre-Gaussian, Zernike modes, and others.
• 🔆 Polarization and Coherence — Simulate polarized light and fields with arbitrary spatial coherence.

Learn more about TorchOptics in our research paper on arXiv.

Installation

TorchOptics is available on PyPI and can be installed with:

pip install torchoptics

Documentation

Read the full documentation at torchoptics.readthedocs.io.

Usage

This example shows how to simulate a 4f imaging system using TorchOptics, computing and visualizing the field at each focal plane along the optical axis:

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

# Set simulation properties
shape = 1000  # Number of grid points in each dimension
spacing = 10e-6  # Spacing between grid points (m)
wavelength = 700e-9  # Field wavelength (m)
focal_length = 200e-3  # Lens focal length (m)
tile_length = 400e-6  # Checkerboard tile length (m)
num_tiles = 15  # Number of tiles in each dimension

# Determine device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Configure default properties
torchoptics.set_default_spacing(spacing)
torchoptics.set_default_wavelength(wavelength)

# Initialize input field with checkerboard pattern
field_data = checkerboard(shape, tile_length, num_tiles)
input_field = Field(field_data).to(device)

# Define 4f optical system with two lenses
system = System(
    Lens(shape, focal_length, z=1 * focal_length),
    Lens(shape, focal_length, z=3 * focal_length),
).to(device)

# Measure field at focal planes along the z-axis
measurements = [
    system.measure_at_z(input_field, z=i * focal_length)
    for i in range(5)
]

# Visualize the measured intensity distributions
for i, measurement in enumerate(measurements):
    measurement.visualize(title=f"z={i}f", vmax=1)

Intensity distributions at different focal planes in the 4f system.

Propagation of the intensity distribution.

For more examples and detailed usage, please refer to the documentation.

Contributing

We welcome contributions! See our 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

TorchOptics is distributed under the MIT License. See the LICENSE file for more details.