TorchOptics

In torchoptics, coherent light (like a laser) is modeled by a single complex field, whereas partially coherent light (like LEDs or thermal sources) requires a statistical description called the mutual coherence function ($\Gamma$).

The SpatialCoherence class provides this statistical description, allowing for physically accurate simulations of sources that cannot be modeled by a single wavefront.

Key Concepts:

• Mutual Coherence Function ($\Gamma$): Encodes the statistical correlation between two spatial points $(x_1, y_1)$ and $(x_2, y_2)$.
• Intensity ($I$): The diagonal of the coherence function, $\Gamma(x, y, x, y)$, represents the time-averaged intensity.
• Memory Complexity: The coherence matrix is stored as a 4D complex tensor of shape (H, W, H, W). Because this scales as $O(N^4)$, it is recommended to keep grid sizes small (typically 20–50 points per dimension) or use GPU acceleration for larger grids.

The torchoptics.System class models an optical system as an ordered sequence of elements along the optical axis. It behaves similarly to torch.nn.Sequential, where propagation between elements is handled automatically. Elements are positioned along the axis based on their z attribute. During a forward pass, the system propagates a field through every element in order of increasing z and returns the field immediately after the last element. Elements with a z position strictly less than the input field's starting z are skipped, while elements at exactly the starting z are applied.

import torchoptics
from torchoptics import Field, System
from torchoptics.elements import Lens

# Elements are ordered by their z position
system = System(
    Lens(shape=(1000, 1000), f=200e-3, z=1 * 200e-3),
    Lens(shape=(1000, 1000), f=200e-3, z=3 * 200e-3),
)

input_field = Field(torchoptics.profiles.gaussian(1000, waist_radius=1e-3), z=0)
output_field = system(input_field)

Polarized fields are represented by data tensors with the shape (..., 3, H, W), where the dimension of size 3 corresponds to the $x$, $y$, and $z$ polarization components. Polarization elements in TorchOptics operate on these fields by applying a $3 \times 3$ Jones matrix $J$ at every grid point:

$$\begin{pmatrix} E_x' \ E_y' \ E_z' \end{pmatrix} = J(x, y) \begin{pmatrix} E_x \ E_y \ E_z \end{pmatrix}$$

To use polarized elements, you must first construct a polarized field (e.g., by initializing a tensor with shape (3, H, W) and wrapping it in a Field object).

import torch
from torchoptics import Field

# Example: Creating an x-polarized Gaussian beam
shape = (512, 512)
polarized_data = torch.zeros(3, *shape, dtype=torch.cdouble)
# Assume gaussian() is available to populate the x-component
polarized_data[0] = gaussian(shape, waist_radius=1.5e-3)
polarized_field = Field(polarized_data)

To handle polarized light, TorchOptics extends the Field class to represent the electric field vector using three complex components corresponding to the $x$, $y$, and $z$ directions.

In paraxial optics, the $z$ component is typically zero, and energy is carried by the transverse $x$ and $y$ components. A polarized field is represented by a data tensor of shape (3, H, W), where:

• Index 0 is the $x$-component.
• Index 1 is the $y$-component.
• Index 2 is the $z$-component.

import torch
import torchoptics
from torchoptics import Field
from torchoptics.profiles import gaussian

torchoptics.set_default_spacing(10e-6)
torchoptics.set_default_wavelength(700e-9)
shape = 200

# Create an x-polarized Gaussian beam (only the x-component is non-zero)
data = torch.zeros(3, shape, shape, dtype=torch.cdouble)
data[0] = gaussian(shape, waist_radius=500e-6)
x_pol = Field(data)