Optical systems

SVETlANNa provides LinearOpticalSetup for assembling sequential optical systems.

LinearOpticalSetup


from svetlanna import LinearOpticalSetup
from svetlanna.elements import ThinLens, FreeSpace, RoundAperture
from svetlanna.units import ureg
 
setup = LinearOpticalSetup([
    RoundAperture(params, radius=2*ureg.mm),
    ThinLens(params, focal_length=100*ureg.mm),
    FreeSpace(params, distance=100*ureg.mm, method='zpASM'),
])
 
# forward propagation
wf_out = setup(wf_in)
 
# step-by-step propagation
steps = setup.stepwise_forward(wf_in)

All elements in a LinearOpticalSetup must share the same SimulationParameters. A mismatch produces a warning at construction time.

Methods

Method	Description
`forward(wf)`	Propagate through every element in order
`stepwise_forward(wf)`	Tuple of wavefronts: the input, then the field after each element
`reverse(wf)`	Backward propagation, if every element supports it

Backward propagation

If every element implements reverse(), the setup builds a reverse network automatically:


wf_back = setup.reverse(wf_out)

Otherwise reverse() raises TypeError. FreeSpace has no reverse(), so a setup containing one must be inverted by hand — build a mirrored setup with negative distances and focal lengths.

Alternative: nn.Sequential

For simple cases the standard nn.Sequential works too:

Basic


import torch.nn as nn
 
system = nn.Sequential(
    RoundAperture(params, radius=2*ureg.mm),
    ThinLens(params, focal_length=100*ureg.mm),
    FreeSpace(params, distance=100*ureg.mm, method='zpASM'),
)
 
wf_out = system(wf_in)

LinearOpticalSetup adds the parameter checks, stepwise_forward, reverse and the specification export on top of that.

Classic configurations

Focusing system


from svetlanna import LinearOpticalSetup
 
def focusing_system(params, focal_length, aperture_radius=None):
    """Aperture (optional) → lens → propagation to the focal plane."""
    elements = []
 
    if aperture_radius is not None:
        elements.append(RoundAperture(params, radius=aperture_radius))
 
    elements.extend([
        ThinLens(params, focal_length=focal_length),
        FreeSpace(params, distance=focal_length, method='zpASM'),
    ])
 
    return LinearOpticalSetup(elements)
 
setup = focusing_system(params, focal_length=100*ureg.mm, aperture_radius=2*ureg.mm)
wf_focus = setup(wf_in)

4f system


from svetlanna.elements import Aperture
 
def four_f_system(params, focal_length, filter_mask=None):
    """
    A 4f system for optical filtering.
    Layout: input → f → L1 → f → [filter] → f → L2 → f → output
    """
    f = focal_length
 
    def prop():
        return FreeSpace(params, distance=f, method='zpASM')
 
    elements = [prop(), ThinLens(params, focal_length=f), prop()]
 
    if filter_mask is not None:
        elements.append(Aperture(params, mask=filter_mask))
 
    elements.extend([prop(), ThinLens(params, focal_length=f), prop()])
 
    return LinearOpticalSetup(elements)

See the 4f system tutorial for the full worked example.

Keplerian telescope


def kepler_telescope(params, f1, f2):
    """Magnification M = -f2 / f1."""
    return LinearOpticalSetup([
        ThinLens(params, focal_length=f1),
        FreeSpace(params, distance=f1 + f2, method='zpASM'),
        ThinLens(params, focal_length=f2),
    ])
 
telescope = kepler_telescope(params, f1=50*ureg.mm, f2=200*ureg.mm)
# magnification: -4×

Diffractive neural networks (DONNs)

A multi-layer D²NN


import torch
import torch.nn as nn
from svetlanna import LinearOpticalSetup
from svetlanna.elements import DiffractiveLayer, FreeSpace
 
class DONN(nn.Module):
    """A diffractive optical neural network."""
 
    def __init__(self, params, n_layers, layer_distance, grid=512):
        super().__init__()
 
        elements = []
        for _ in range(n_layers):
            # a trainable phase mask
            phase = nn.Parameter(torch.rand(grid, grid) * 2 * torch.pi)
            elements.append(DiffractiveLayer(params, mask=phase))
            elements.append(
                FreeSpace(params, distance=layer_distance, method='zpASM')
            )
 
        self.optical = LinearOpticalSetup(elements)
 
    def forward(self, wf):
        return self.optical(wf)
 
donn = DONN(params, n_layers=5, layer_distance=10*ureg.mm)
wf_out = donn(wf_in)

DiffractiveLayer registers the mask through process_parameter, so the masks appear in donn.parameters() and are picked up by any PyTorch optimiser.

With a detector for classification


from svetlanna.detector import Detector, DetectorProcessorClf
 
class DONNClassifier(nn.Module):
    def __init__(self, params, n_layers, n_classes):
        super().__init__()
        self.donn = DONN(params, n_layers=n_layers, layer_distance=10*ureg.mm)
        self.detector = Detector(params, func='intensity')
        self.processor = DetectorProcessorClf(
            num_classes=n_classes,
            simulation_parameters=params,
            segmentation_type='strips'
        )
 
    def forward(self, wf):
        wf = self.donn(wf)
        intensity = self.detector(wf)
        return self.processor.batch_forward(intensity)
 
classifier = DONNClassifier(params, n_layers=5, n_classes=10)

SVETlANNa also ships ready-made architectures in svetlanna.networks: DiffractiveRNN, ConvDiffNetwork4F, LinearAutoencoder and SimpleReservoir.

Intermediate fields

stepwise_forward returns the input wavefront followed by the field after each element, so a setup with three elements yields four wavefronts:


setup = LinearOpticalSetup([
    RoundAperture(params, radius=2*ureg.mm),
    ThinLens(params, focal_length=100*ureg.mm),
    FreeSpace(params, distance=100*ureg.mm, method='zpASM'),
])
 
steps = setup.stepwise_forward(wf_in)
 
# steps[0] — the input
# steps[1] — after the aperture
# steps[2] — after the lens
# steps[3] — in the focal plane
 
import matplotlib.pyplot as plt
titles = ['Input', 'After the aperture', 'After the lens', 'Focal plane']
fig, axes = plt.subplots(1, len(steps), figsize=(4*len(steps), 4))
for ax, wf, title in zip(axes, steps, titles):
    ax.imshow(wf.intensity.cpu(), cmap='hot')
    ax.set_title(title)

stepwise_forward runs under torch.no_grad() — it is a diagnostic tool, not part of a training loop.

Working on the GPU


# option 1: move params first, then build the system
params.to('cuda')
setup = LinearOpticalSetup([...])   # already on the GPU
 
# option 2: move an existing system
setup = LinearOpticalSetup([...])
setup = setup.to('cuda')
 
# the wavefront must be on the GPU too
wf_out = setup(wf.to('cuda'))

Serialisation


import torch
 
torch.save(setup.state_dict(), 'optical_system.pth')
 
setup = LinearOpticalSetup([...])   # rebuild with the same structure
setup.load_state_dict(torch.load('optical_system.pth'))

For a human-readable description of a setup, use the specs module and the svetlanna.visualization helpers (show_specs, show_structure, show_stepwise_forward), which render interactive widgets inside Jupyter.