TorchOptics is an open-source Python library for differentiable wave optics simulations with PyTorch.
- 🌊 Differentiable Wave Optics: A comprehensive framework for modeling, analyzing, and designing optical systems using differentiable Fourier optics.
- 🔥 Built on PyTorch: Leverages PyTorch for GPU acceleration, batch processing, automatic differentiation, and efficient gradient-based optimization.
- 🛠️ End-to-End Optimization: Enables optimization of optical hardware and deep learning models within a unified, differentiable pipeline.
- 🔬 Optical Elements: Features standard optical elements like modulators, lenses, detectors, and polarizers.
- 🖼️ Spatial Profiles: Provides a wide range of spatial profiles, including Hermite-Gaussian and Laguerre-Gaussian beams.
- 🔆 Polarization & Coherence: Supports simulations of polarized light and optical fields with arbitrary spatial coherence.
Our research paper, available on arXiv, introduces the TorchOptics library and provides a comprehensive review of its features and applications.
Access the latest documentation at torchoptics.readthedocs.io.
To install the latest stable release of TorchOptics from PyPI, run:
pip install torchopticsFor the latest development version, install directly from GitHub:
git clone https://github.com/MatthewFilipovich/torchoptics
cd torchoptics
pip install -e '.[dev]'This installs the library in editable mode, along with additional dependencies for development and testing.
This example demonstrates how to simulate a 4f imaging system using TorchOptics. The field at each focal plane along the z-axis is computed and visualized:
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.
We welcome bug reports, questions, and feature suggestions to improve TorchOptics.
- Found a bug or have a question? Please open an issue on GitHub.
- Want to contribute a new feature? Follow these steps:
- Fork the repository: Go to https://github.com/MatthewFilipovich/torchoptics/fork
- Create a feature branch:
git checkout -b feature/fooBar - Commit your changes:
git commit -am 'Add some fooBar' - Push to the branch:
git push origin feature/fooBar - Submit a Pull Request: Open a Pull Request on GitHub
If you use TorchOptics in your research, please cite our paper:
M.J. Filipovich and A.I. Lvovsky, TorchOptics: An open-source Python library for differentiable Fourier optics simulations, arXiv preprint arXiv:2411.18591 (2024).
TorchOptics is distributed under the MIT License. See the LICENSE file for more details.