Core

Core classes for optical simulations


from svetlanna import Wavefront, SimulationParameters, Parameter

Classes

Wavefront

Inherits: torch.Tensor

Class that represents wavefront. It is a subclass of torch.Tensor with additional properties and methods for wavefront analysis and generation.

Properties

intensity

Intensity of the wavefront.

max_intensity

Maximum intensity of the wavefront.

phase

Phase of the wavefront.

Factory methods

plane_wave classmethod


plane_wave(cls, simulation_parameters: SimulationParameters, distance: float = 0.0, wave_direction: Any = None, initial_phase: float = 0.0) -> Self

Create a plane wave wavefront defind by the formula

E(x, y) = \exp\left( i \left( k_x x + k_y y + k_z z + \phi_0 \right) \right)

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`distance`	`float, optional`	Free wave propagation distance $z$ , by default 0.
`wave_direction`	`Any, optional`	Three component tensor-like vector with ( $d_x$ , $d_y$ , $d_z$ ) coordinates, so $\vec{k} = k \frac{\vec{d}}{\\|\\|\vec{d}\\|\\|}$ The resulting field propagates along the vector, by default the wave propagates along z direction.
`initial_phase`	`float, optional`	Additional phase offset ( $\phi_0$ ), by default 0.

📤 Returns

Wavefront

Plane wave field.

gaussian_beam classmethod


gaussian_beam(cls, simulation_parameters: SimulationParameters, waist_radius: float, distance: float = 0.0, dx: float = 0.0, dy: float = 0.0) -> Self

Generates the Gaussian beam wavefront defined by the formula

E(x, y) = \frac&#123;w_0&#125;&#123;w(z)&#125; \exp\left( -\frac&#123;(x - d_x)^2 + (y - d_y)^2&#125;&#123;w(z)^2&#125; \right) \newline \cdot \exp\left( i \left( k z + k\frac&#123;(x - d_x)^2 + (y - d_y)^2&#125;&#123;2 R(z)&#125; - \zeta(z) \right) \right)

where $w(z) = w_0 \sqrt{1 + \left( \frac{z}{z_R} \right)^2}$ , $R(z) = z \left( 1 + \left( \frac{z_R}{z} \right)^2 \right)$ , $\zeta(z) = \arctan\left( \frac{z}{z_R} \right)$ , and $z_R = \frac{\pi w_0^2}{\lambda}$ is the Rayleigh range.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`waist_radius`	`float`	Beam waist radius ( $w_0$ ).
`distance`	`float, optional`	Free wave propagation distance $z$ , by default 0.
`dx`	`float, optional`	Horizontal offset of the beam center ( $d_x$ ), by default 0.
`dy`	`float, optional`	Vertical offset of the beam center ( $d_y$ ), by default 0.

📤 Returns

Wavefront

Gaussian beam field in the oXY plane.

spherical_wave classmethod


spherical_wave(cls, simulation_parameters: SimulationParameters, distance: float, initial_phase: float = 0.0, dx: float = 0.0, dy: float = 0.0) -> Self

Generate wavefront of the spherical wave

E(x, y) = \frac&#123;1&#125;&#123;r&#125; \exp\left( i \left( k r + \phi_0 \right) \right)

where $r = \sqrt{(x - d_x)^2 + (y - d_y)^2 + z^2}$ is the distance from the point source to the point $(x, y)$ in the oXY plane.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`distance`	`float`	Distance from the point source to the oXY plane ( $z$ ).
`initial_phase`	`float, optional`	Phase offset at the source ( $\phi_0$ ), by default 0.
`dx`	`float, optional`	Horizontal position of the point source ( $d_x$ ), by default 0.
`dy`	`float, optional`	Vertical position of the point source ( $d_y$ ), by default 0.

📤 Returns

Wavefront

Spherical wave field in the oXY plane.

hermite_gauss classmethod


hermite_gauss(cls, simulation_parameters: SimulationParameters, waist_radius: float, distance: float = 0.0, dx: float = 0.0, dy: float = 0.0, m: int = 0, n: int = 0) -> Self

Generates the Hermite-Gaussian mode wavefront defined by the formula

E(x, y) = \frac&#123;w_0&#125;&#123;w(z)&#125; H_m\left(\frac&#123;\sqrt&#123;2&#125;(x-d_x)&#125;&#123;w_0&#125;\right) H_n\left(\frac&#123;\sqrt&#123;2&#125;(y-d_y)&#125;&#123;w_0&#125;\right) \exp\left( -\frac&#123;(x - d_x)^2 + (y - d_y)^2&#125;&#123;w(z)^2&#125; \right) \newline \cdot \exp\left( i \left( k z + k\frac&#123;(x - d_x)^2 + (y - d_y)^2&#125;&#123;2 R(z)&#125; - \zeta(z) \right) \right)

where $w(z) = w_0 \sqrt{1 + \left( \frac{z}{z_R} \right)^2}$ , $R(z) = z \left( 1 + \left( \frac{z_R}{z} \right)^2 \right)$ , $\zeta(z) = (m+n+1)\arctan\left( \frac{z}{z_R} \right)$ , and $z_R = \frac{\pi w_0^2}{\lambda}$ is the Rayleigh range.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`waist_radius`	`float`	Beam waist radius ( $w_0$ ).
`distance`	`float, optional`	Free wave propagation distance $z$ , by default 0.
`dx`	`float, optional`	Horizontal offset of the mode center ( $d_x$ ), by default 0.
`dy`	`float, optional`	Vertical offset of the mode center ( $d_y$ ), by default 0.
`m`	`int, optional`	Mode index in the x direction, by default 0 (fundamental mode).
`n`	`int, optional`	Mode index in the y direction, by default 0 (fundamental mode).

📤 Returns

Wavefront

Hermite-Gaussian field.

Methods

fwhm


fwhm(self, simulation_parameters: SimulationParameters) -> tuple[float, float]

Full width at half maximum (FWHM) of the wavefront intensity.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.

📤 Returns

tuple[float, float]

FWHM along x and y axes.

SimulationParameters

Inherits: nn.Module

Properties

axis_names

Get names of non-scalar axes (those with length > 1).

device

Get the device where all axes are stored.

axes

names

Factory methods

from_ranges classmethod


from_ranges(cls, **additional_axes: torch.Tensor | float) -> Self

Create SimulationParameters from coordinate ranges.

📥 Parameters

Parameter	Type	Description
`x_range`	`tuple[float, float]`	(min, max) range for x-axis. Use `ureg` for units.
`x_points`	`int`	Number of points along x-axis.
`y_range`	`tuple[float, float]`	(min, max) range for y-axis. Use `ureg` for units.
`y_points`	`int`	Number of points along y-axis.
`wavelength`	`torch.Tensor	float`
`**additional_axes`	`torch.Tensor	float`

Examples

>>> from svetlanna.units import ureg —
>>> params = SimulationParameters.from_ranges( —
… x_range=(-1ureg.mm, 1ureg.mm), x_points=256, —
… y_range=(-1ureg.mm, 1ureg.mm), y_points=256, —
… wavelength=632.8ureg.nm* —
… ) —

from_dict classmethod


from_dict(cls, axes_dict: Mapping[str, torch.Tensor | float]) -> Self

Create SimulationParameters from a dictionary.

📥 Parameters

Parameter	Type	Description
`axes_dict`	`Mapping[str, torch.Tensor	float]`

Methods

__init__ constructor


__init__() -> None

__init__ constructor


__init__(**additional_axes: torch.Tensor | float) -> None

__init__ constructor


__init__(**kwaxes: torch.Tensor | float) -> None

Simulation parameters. Manages coordinate systems and physical parameters for optical simulations. Required axes: x, y, wavelength. Additional axes can be added.

Inherits from nn.Module so that axes are registered as buffers and participate in automatic device management when used as submodules of Elements.

Note

Axes are registered as non-persistent buffers (persistent=False). This means they are not included in state_dict() and will not be saved during checkpointing. The simulation grid must be provided when constructing the model; it does not need to be restored from a checkpoint.

Examples

Let’s define simalation grid of width and height of 1 mm with 512 points for both axes (Nx=Ny=512) and wavelength of 632.8 nm: —
```python —
import svetlanna as sv —
from svetlanna.units import ureg —
import torch —
sim_params = sv.SimulationParameters( — x=torch.linspace(-0.5, 0.5, 512) * ureg.mm, y=torch.linspace(-0.5, 0.5, 512) * ureg.mm, wavelength=632.8 * ureg.nm,
) —
``` —
You can make wavelength an array for polychromatic simulations: —
```python hl_lines=“4” —
sim_params = sv.SimulationParameters( — x=torch.linspace(-0.5, 0.5, 512) * ureg.mm, y=torch.linspace(-0.5, 0.5, 512) * ureg.mm, wavelength=torch.linspace(600, 800, 10) * ureg.nm,
) —
``` —
The order of axes matters! It defines the order of dimensions in wavefront tensors. —
In first case above, all optical elements will expect wavefront tensors with shape (..., Ny, Nx), —
while in the second case, the expected shape will be (..., Nwavelength, Ny, Nx). —
... means any number of leading dimensions (e.g., for batch). —
If you change the order: —
```python hl_lines=“3 4” —
sim_params = sv.SimulationParameters( — x=torch.linspace(-0.5, 0.5, 512) * ureg.mm, wavelength=torch.linspace(600, 800, 10) * ureg.nm, y=torch.linspace(-0.5, 0.5, 512) * ureg.mm,
) —
``` —
the expected order of axes is ('y', 'wavelength', 'x'), so all optical elements will expect wavefront tensors with shape (..., Ny, Nwavelength, Nx). —
You can add custom axes as needed: —
```python hl_lines=“2 4” —
sim_params = sv.SimulationParameters( — t=torch.linspace(0, 1, 5) * ureg.s, # time axis x=torch.linspace(-0.5, 0.5, 512) * ureg.mm, wavelength=632.8 * ureg.nm, y=torch.linspace(-0.5, 0.5, 512) * ureg.mm,
) —
``` —
In this case, the expected order of axes is ('y', 'x', 't') as wavelength is scalar, so all optical elements will expect wavefront tensors with shape (..., Ny, Nx, Nt). —

clone


clone(self) -> 'SimulationParameters'

Create a deep copy of the SimulationParameters instance.

📤 Returns

SimulationParameters

A new instance with cloned axes.

equal


equal(self, value: SimulationParameters) -> bool

Check equality with another SimulationParameters instance. The comparison between tensor axes is based on torch.equal, see documentation for more details. Comparing instances on diffrent devices will raise RuntimeError because torch.equal requires tensors to be on the same device.

📥 Parameters

Parameter	Type	Description
`value`	`SimulationParameters`	SimulationParameters instance to compare with.

📤 Returns

bool

True if all axes are equal, False otherwise.

meshgrid


meshgrid(self, x_axis: str, y_axis: str) -> tuple[torch.Tensor, torch.Tensor]

Create a coordinate meshgrid from two axes.

📥 Parameters

Parameter	Type	Description
`x_axis`	`str`	Name of the axis for x-coordinates (typically ‘x’).
`y_axis`	`str`	Name of the axis for y-coordinates (typically ‘y’).

📤 Returns

tuple[torch.Tensor, torch.Tensor]

2D coordinate grids with ‘xy’ indexing convention.

Examples

```python —
import svetlanna as sv —
import torch —
sim_params = sv.SimulationParameters( — x=torch.linspace(-0.5, 0.5, 10), y=torch.linspace(-0.5, 0.5, 12), wavelength=1,
) —
X, Y = sim_params.meshgrid(“x”, “y”) —
print(X.shape) # torch.Size([12, 10]) —
``` —

axis_sizes


axis_sizes(self, axs: tuple[str, ...] | None = None) -> torch.Size

Get the size of specified axes in order (cached for performance).

📥 Parameters

Parameter	Type	Description
`axs`	`tuple[str, …]	None`

📤 Returns

torch.Size

Size object with lengths of specified axes.

Examples

```python —
import svetlanna as sv —
import torch —
sim_params = sv.SimulationParameters( — x=torch.linspace(-0.5, 0.5, 10), y=torch.linspace(-0.5, 0.5, 12), wavelength=1,
) —
print(sim_params.axis_sizes((‘y’, ‘x’))) # torch.Size([12, 10]) —
``` —

index


index(self, name: str) -> int

Get the negative index of an axis in tensors.

📥 Parameters

Parameter	Type	Description
`name`	`str`	Name of the axis.

📤 Returns

int

Negative index for use in tensor operations.

⚠️ Raises

AxisNotFound — If the axis doesn’t exist or is scalar.

cast


cast(self, tensor: torch.Tensor, *axes: str) -> torch.Tensor

Cast tensor to match simulation parameters axes for broadcasting.

Reshapes tensor so it can be broadcast with wavefront tensors. Scalar axes are skipped (they don’t affect tensor shape).

📥 Parameters

Parameter	Type	Description
`tensor`	`torch.Tensor`	Input tensor whose trailing dimensions correspond to `axes`.
`*axes`	`str`	Axes names corresponding to tensor’s trailing dimensions.

📤 Returns

torch.Tensor

Tensor reshaped for broadcasting with wavefront.

Examples

```python —
import svetlanna as sv —
import torch —
sim_params = sv.SimulationParameters( — x=torch.linspace(-0.5, 0.5, 3), y=torch.linspace(-0.5, 0.5, 2), wavelength=torch.linspace(1, 2, 5),
) —
# axes: (wavelength, y, x) —
print(sim_params.axis_sizes((“wavelength”, “y”, “x”))) # torch.Size([5, 2, 3]) —
a = torch.rand(2, 3) # y, x —
a = sim_params.cast(a, “y”, “x”) —
print(a.shape) # torch.Size([1, 2, 3]) —
# a is now ready to broadcast with tensor of shape (5, 2, 3) —
``` —

axes_size


axes_size(self, *args, **kwargs)

Parameter

Inherits: torch.Tensor

torch.Parameter-like tensor with an internal storage module.

This class is used to keep a trainable torch.nn.Parameter inside a torch.nn.Module while presenting a torch.Tensor-like interface.

Properties

inner_parameter

Methods

__init__ constructor


__init__(self, data: Any, requires_grad: bool = True)

📥 Parameters

Parameter	Type	Description
`data`	`Any`	Initial value, should be a tensor or convertible to a tensor.
`requires_grad`	`bool, optional`	Whether the parameter requires gradients, by default True.

Examples

You can use Parameter as a trainable parameter in any SVETlANNa —
element when it is typed as —
[OptimizableFloat][svetlanna.parameters.OptimizableFloat] or —
[OptimizableTensor][svetlanna.parameters.OptimizableTensor]: —
```python —
import svetlanna as sv —
import torch —
sim_params = sv.SimulationParameters(…) —
diffractive_layer = sv.elements.DiffractiveLayer( — simulation_parameters=sim_params, mask=sv.Parameter(2 * torch.pi * torch.rand(Ny, Nx)),
) —
``` —

ConstrainedParameter

Inherits: Parameter

Parameter constrained to a bounded range.

The constraint is implemented by applying bound_func to the inner parameter, mapping it to $[0, 1]$ , and then scaling and shifting it to (min_value, max_value).

Properties

min_value

max_value

bound_func

inv_bound_func

value

Constrained parameter value.

Methods

__init__ constructor


__init__(self, data: Any, min_value: Any, max_value: Any, bound_func: Callable[[torch.Tensor], torch.Tensor] = torch.sigmoid, inv_bound_func: Callable[[torch.Tensor], torch.Tensor] = torch.logit, requires_grad: bool = True)

📥 Parameters

Parameter	Type	Description
`data`	`Any`	Initial parameter value.
`min_value`	`Any`	Minimum allowed value.
`max_value`	`Any`	Maximum allowed value.
`bound_func`	`Callable[[torch.Tensor], torch.Tensor], optional`	Function that maps $\mathbb{R}\to[0,1]$ , by default `torch.sigmoid`.
`inv_bound_func`	`Callable[[torch.Tensor], torch.Tensor], optional`	Inverse of `bound_func`, by default `torch.logit`. It is used once to compute the initial inner parameter value from `data`.
`requires_grad`	`bool, optional`	Whether the parameter requires gradients, by default True.

Examples

You can use ConstrainedParameter as a trainable parameter in any —
SVETlANNa element when it is typed as —
[OptimizableFloat][svetlanna.parameters.OptimizableFloat] or —
[OptimizableTensor][svetlanna.parameters.OptimizableTensor]: —
```python —
import svetlanna as sv —
import torch —
sim_params = sv.SimulationParameters(…) —
diffractive_layer = sv.elements.DiffractiveLayer( — simulation_parameters=sim_params, mask=sv.ConstrainedParameter( 2 * torch.pi * torch.rand(Ny, Nx), min_value=0, max_value=2 * torch.pi, )
) —
``` —

LinearOpticalSetup

Inherits: nn.Module

Linear optical network composed of [Element][svetlanna.elements.Element] instances. It works the same way as a torch.nn.Sequential module, but with some additional features.

Methods

__init__ constructor


__init__(self, elements: Iterable[Element]) -> None

📥 Parameters

Parameter	Type	Description
`elements`	`Iterable[Element]`	Optical elements that make up the setup. Elements are evaluated in the provided order.

Examples

```python —
import svetlanna as sv —
setup = sv.LinearOpticalSetup( — elements=[ element1, element2, element3, ]
) —
output_wavefront = setup(input_wavefront) —
``` —

forward


forward(self, input_wavefront: Wavefront) -> Wavefront

stepwise_forward


stepwise_forward(self, input_wavefront: Wavefront) -> tuple[Wavefront, ...]

Apply elements step-by-step and collect intermediate wavefronts.

📥 Parameters

Parameter	Type	Description
`input_wavefront`	`Wavefront`	A wavefront that enters the optical network.

📤 Returns

tuple[Wavefront, ...]

A tuple of wavefronts showing the propagation through the setup. The first wavefront is the input wavefront, and the last one is the output wavefront after propagation through all elements.

reverse


reverse(self, Ein: Tensor) -> Tensor

Reverse propagation through the setup. All elements in the setup must have a reverse method. If any element lacks this method, a TypeError is raised.

📥 Parameters

Parameter	Type	Description
`Ein`	`Tensor`	Input wavefront to reverse propagate.

📤 Returns

Tensor

Output wavefront after reverse propagation.

⚠️ Raises

TypeError — If reverse propagation is not supported by all elements in the setup.

to_specs


to_specs(self) -> Iterable[ParameterSpecs | SubelementSpecs]

Detector

Inherits: Element

Object that plays a role of a physical detector in an optical system: (1) func=‘intensity’ transforms incident field to intensities for further image analysis (2) …

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, func = 'intensity')

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters for a further optical network.
`func`	`str`	A parameter that defines a function that will be applied to an incident field to obtain a detector image. (1) func=‘intensity’ – detector returns intensities (2) …

forward


forward(self, input_field: Wavefront) -> torch.Tensor

Method that returns the image obtained from the incident field by a detector using self.func. in the simplest case the image on a detector is an intensities image) …

📥 Parameters

Parameter	Type	Description
`input_field`	`Wavefront`	A tensor (Wavefront) of an incident field on a detector.

📤 Returns

detector_output : torch.Tensor

The image on a detector (according to self.func).