Optical elements

Every optical element inherits from torch.nn.Module and from the Element base class.

The common interface


from svetlanna.elements import ThinLens
from svetlanna.units import ureg
 
lens = ThinLens(simulation_parameters=params, focal_length=100*ureg.mm)
 
# forward propagation
wf_out = lens(wf_in)
 
# backward propagation, where the element supports it
wf_back = lens.reverse(wf_out)
 
# the complex transmission function (a property, for elements that have one)
t = lens.transmission_function

FreeSpace

Free-space propagation.


from svetlanna.elements import FreeSpace
 
propagate = FreeSpace(
    simulation_parameters=params,
    distance=50*ureg.mm,
    method='zpASM',
)
 
wf_out = propagate(wf_in)

method is required. Valid values are 'ASM', 'zpASM', 'RSC' and 'zpRSC'. FreeSpace has no reverse(): to propagate backwards, create a second FreeSpace with a negative distance.

Propagation methods

ASM / zpASM

Angular Spectrum Method — the field is decomposed into plane waves, each of which acquires its own propagation phase:

$E_{\text{out}} = \mathcal{F}^{-1}\left[\mathcal{F}[E_{\text{in}}] \cdot H\right], \quad H = \exp\left(ikz\sqrt{1 - \lambda^2 f_x^2 - \lambda^2 f_y^2}\right)$

Accurate for short and moderate distances. zpASM zero-pads the field first, which removes the wrap-around artefacts of the plain FFT at the cost of memory and time.


# extra padding, if the defaults are not enough
propagate = FreeSpace(
    params, distance=1*ureg.m, method='zpRSC',
    total_paddings_x=1024, total_paddings_y=1024,
)

FreeSpace warns when the chosen method is close to its sampling limit for the given distance and grid. Take those warnings seriously — they signal aliasing.

ThinLens

A thin lens: a quadratic phase transmission function.


from svetlanna.elements import ThinLens
 
lens = ThinLens(
    simulation_parameters=params,
    focal_length=100*ureg.mm,
    radius=10*ureg.mm,   # optional lens aperture, default: infinite
)
 
wf_out = lens(wf_in)
wf_back = lens.reverse(wf_out)

Physics

$T(x, y) = \exp\left(-i\frac{k}{2f}(x^2 + y^2)\right) \cdot M(r)$

where $M(r)$ is the aperture mask if radius is given.

Parameter	Type	Description
`focal_length`	float or `Parameter`	Focal length; may be trainable
`radius`	float	Aperture radius (default `inf`)

Aperture, RoundAperture, RectangularAperture

Amplitude masks.

RoundAperture


from svetlanna.elements import RoundAperture
 
aperture = RoundAperture(
    simulation_parameters=params,
    radius=5*ureg.mm
)

RectangularAperture


from svetlanna.elements import RectangularAperture
 
aperture = RectangularAperture(
    simulation_parameters=params,
    width=10*ureg.mm,
    height=8*ureg.mm
)

Aperture


from svetlanna.elements import Aperture
import torch
 
# an arbitrary mask on the simulation grid
mask = torch.zeros(512, 512)
mask[100:400, 100:400] = 1.0
 
aperture = Aperture(
    simulation_parameters=params,
    mask=mask
)

DiffractiveLayer

A passive phase element — a diffractive optical element (DOE), and the basic building block of a diffractive neural network.


from svetlanna.elements import DiffractiveLayer
import torch
 
phase_mask = torch.rand(512, 512) * 2 * torch.pi
 
layer = DiffractiveLayer(
    simulation_parameters=params,
    mask=phase_mask,
    mask_norm=2*torch.pi,   # normalisation, default 2π
)
 
wf_out = layer(wf_in)
wf_back = layer.reverse(wf_out)

Transmission function

$T(x, y) = \exp\left(i \cdot 2\pi \cdot \frac{\text{mask}}{\text{mask\_norm}}\right)$

Wrapping mask in torch.nn.Parameter makes the layer trainable.

SpatialLightModulator (SLM)

A programmable modulator. Unlike DiffractiveLayer it has its own physical size and pixel grid, which may differ from the simulation grid; the mask is interpolated onto the simulation grid and acts as an aperture outside its active area.


from svetlanna.elements import SpatialLightModulator
import torch
 
slm = SpatialLightModulator(
    simulation_parameters=params,
    mask=torch.rand(256, 256) * 2 * torch.pi,
    height=8*ureg.mm,
    width=10*ureg.mm,
    center=(0.0, 0.0),   # position of the SLM centre, (x, y)
    mode='nearest',      # interpolation mode
)
 
wf_out = slm(wf_in)

Parameter	Description
`mask`	Phase mask on the SLM pixel grid, shape `(Ny_mask, Nx_mask)`
`height`, `width`	Physical size of the modulator
`center`	Position of its centre in simulation coordinates
`lut_function`	Lookup table applied to the mask values (quantisation, gamma)
`mode`	Interpolation mode: `'nearest'`, `'bilinear'`, `'bicubic'`, `'area'`, `'nearest-exact'`

Phase quantisation

A real SLM offers a discrete set of phase levels. QuantizerFromStepFunction builds a differentiable quantiser from a smooth one-step function, so a quantised modulator can still be trained:


import svetlanna as sv
from svetlanna.elements.slm import QuantizerFromStepFunction, one_step_tanh
 
slm = SpatialLightModulator(
    params,
    mask=mask,
    height=8*ureg.mm,
    width=8*ureg.mm,
    lut_function=sv.PartialWithParameters(
        QuantizerFromStepFunction(
            N=256,                     # number of levels
            max_value=2*torch.pi,
            one_step_function=one_step_tanh,
        ),
        alpha=torch.tensor(1.0),       # steepness of each step
    ),
)

See the SLM tutorial for a worked example.

NonlinearElement

An element defined by an arbitrary response function mapping a wavefront to a wavefront.


import torch
from svetlanna.elements import NonlinearElement
 
# E_out = sqrt(|E_in|) * exp(i arg(E_in))
nl_element = NonlinearElement(
    simulation_parameters=params,
    response_function=lambda E: torch.polar(torch.sqrt(E.abs()), E.angle()),
)
 
wf_out = nl_element(wf_in)

The response function may carry trainable parameters:


import svetlanna as sv
 
def response_function(E, a, b):
    return torch.polar(E.abs()**a, b * E.angle())
 
nl_element = NonlinearElement(
    simulation_parameters=params,
    response_function=sv.PartialWithParameters(
        response_function,
        a=sv.ConstrainedParameter(0.5, min_value=0.0, max_value=1.0),
        b=sv.Parameter(1.0),
    ),
)

The response function can be any nn.Module — including a neural network — which makes this the hook for hybrid optical/electronic architectures.

Element summary

Element	Description	Key parameters	`reverse()`
`FreeSpace`	Propagation	`distance`, `method`	❌ (use a negative distance)
`ThinLens`	Thin lens	`focal_length`, `radius`	✅
`RoundAperture`	Circular aperture	`radius`	❌
`RectangularAperture`	Rectangular aperture	`width`, `height`	❌
`Aperture`	Arbitrary mask	`mask`	❌
`DiffractiveLayer`	DOE / phase mask	`mask`, `mask_norm`	✅
`SpatialLightModulator`	SLM	`mask`, `height`, `width`, `lut_function`	✅
`NonlinearElement`	Nonlinearity	`response_function`	❌

Trainable parameters

Numeric parameters can be made trainable:


from svetlanna import Parameter, ConstrainedParameter
 
# a trainable focal length
focal = Parameter(data=100e-3, requires_grad=True)
lens = ThinLens(params, focal_length=focal)
 
# constrained to [50, 200] mm
focal_bounded = ConstrainedParameter(
    data=100e-3,
    min_value=50e-3,
    max_value=200e-3,
    requires_grad=True
)
lens = ThinLens(params, focal_length=focal_bounded)

ConstrainedParameter reparametrises the value through a sigmoid, so the optimiser can move freely while the physical value stays inside the bounds.

Writing your own element


import torch
from svetlanna.elements import Element
 
class CustomElement(Element):
    def __init__(self, simulation_parameters, my_param):
        super().__init__(simulation_parameters)
        # process_parameter registers trainable parameters and buffers correctly
        self.my_param = self.process_parameter("my_param", my_param)
 
    @property
    def transmission_function(self):
        X, Y = self.simulation_parameters.meshgrid(x_axis='x', y_axis='y')
        return torch.exp(1j * self.my_param * (X**2 + Y**2))
 
    def forward(self, incident_wavefront):
        return incident_wavefront * self.transmission_function
 
    def reverse(self, transmission_wavefront):
        return transmission_wavefront * self.transmission_function.conj()