Elements

Optical elements: lenses, apertures, diffractive layers, SLMs and more


from svetlanna.elements import ...

Classes

ThinLens

Inherits: Element

Properties

transmission_function

The tensor representing the transmission function of the element

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, focal_length: OptimizableFloat, radius: float = torch.inf)

Thin lens element.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`focal_length`	`OptimizableFloat`	The focal length of the lens. $\text{focal\_length} > 0$ for a converging lens.
`radius`	`float`	The radius of the thin lens. The field outside the radius ( $x^2 + y^2 > \text{radius}^2$ ) will propagate with no change in phase. Default is infinity, meaning that the lens has no aperture and the field will propagate with a phase change everywhere.

forward


forward(self, incident_wavefront: Wavefront) -> Wavefront

reverse


reverse(self, transmission_wavefront: Wavefront) -> Wavefront

to_specs


to_specs(self) -> Iterable[ParameterSpecs]

NonlinearElement

Inherits: Element

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, response_function: Callable[[Wavefront], Wavefront])

Nonlinear optical element with a given response function. The response function takes an incident wavefront and returns the modified wavefront.

Examples

Suppose the response function is defined as —
$E^\text{out} = \sqrt{|E^\text{in}|}e^{i \arg(E^\text{in})}$ : —
```python hl_lines=“8” —
import svetlanna as sv —
import torch —
sim_params = sv.SimulationParameters(…) —
sv.elements.NonlinearElement( — simulation_parameters=sim_params, response_function = lambda E: torch.polar(torch.sqrt(E.abs()), E.angle())
) —
``` —
If you want to optimize the parameters of the response function, you can use —
svetlanna.PartialWithParameters to wrap the response function with trainable parameters. —
For example, if the response function is defined as —
$E^\text{out} = |E^\text{in}|^a e^{i b \arg(E^\text{in})}$ , where $0<a<1$ and —
$b$ are trainable, you can define the nonlinear element as follows: —
```python hl_lines=“1 2 6-10” —
def response_function(E, a, b): — return torch.polar(E.abs()**a, b * E.angle())
sv.elements.NonlinearElement( — simulation_parameters=sim_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), ),
) —
``` —
You can also train a neural network inside the nonlinear element! —
```python hl_lines=“3” —
sv.elements.NonlinearElement( — simulation_parameters=sim_params, response_function=my_neural_network,
) —
``` —

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`response_function`	`Callable[[Wavefront], Wavefront]`	Function that describes the nonlinear response of the element.

forward


forward(self, incident_wavefront: Wavefront) -> Wavefront

SpatialLightModulator

Inherits: Element Generic[_F]

Properties

transmission_function

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, mask: OptimizableTensor, height: float, width: float, lut_function: _F = identity, center: Tuple[float, float] = (0.0, 0.0), mode: Literal['nearest', 'bilinear', 'bicubic', 'area', 'nearest-exact'] = 'nearest')

Spatial Light Modulator (SLM) element implementation. SLM supports pixel size that differs from the simulation grid size. The lookup table function (lut_function) allows applying a non-linear transformation to the mask values, for example, to implement quantization.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`mask`	`OptimizableTensor`	Mask tensor of the shape `(Ny_mask, Nx_mask)`, where `Ny_mask` and `Nx_mask` are the height and width of the mask in pixels. It can be different from the simulation grid shape `(Ny, Nx)`; interpolation is applied to fit the mask to the SLM area.
`height`	`float`	Height of the SLM.
`width`	`float`	Width of the SLM.
`lut_function`	`_F, optional`	Lookup table function applied to the mask values, by default `identity`.
`center`	`Tuple[float, float], optional`	Center coordinate `(x, y)` of the SLM in the simulation grid coordinates, by default `(0.0, 0.0)`.
`mode`	`Literal[ 'nearest', 'bilinear', 'bicubic', 'area', 'nearest-exact' ], optional`	Interpolation mode for resizing the mask, by default `'nearest'`. See `torch.nn.functional.interpolate` documentation for more details.

forward


forward(self, incident_wavefront: Wavefront) -> Wavefront

reverse


reverse(self, transmission_wavefront: Wavefront) -> Wavefront

DiffractiveLayer

Inherits: Element

Properties

transmission_function

The tensor representing the transmission function of the element

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, mask: OptimizableTensor, mask_norm: float = 2 * torch.pi)

Diffractive layer defined by a phase mask. The field after propagating through the layer is calculated as:

E^\text&#123;out&#125;_&#123;xyw...&#125; = E^\text&#123;in&#125;_&#123;xyw...&#125; \cdot \exp\left(2\pi i \frac&#123;\text&#123;mask&#125;_&#123;xy&#125;&#125;&#123;\text&#123;mask\_norm&#125;&#125;\right)

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`mask`	`OptimizableTensor`	Two-dimensional tensor representing the aperture mask of shape `(H, W)`.
`mask_norm`	`float, optional`	Mask normalization factor.

forward


forward(self, incident_wavefront: Wavefront) -> Wavefront

reverse


reverse(self, transmission_wavefront: Wavefront) -> Wavefront

to_specs


to_specs(self) -> Iterable[ParameterSpecs]

_BufferedValueContainer

Inherits: tuple

Internal class that marks buffered values.

It is used to prevent double setattr calls with the same value in patterns like self.x = self.make_buffer('x', x_value). Inheriting from tuple is used for performance reasons, hence __slots__. This approach was identified by GPT as the fastest one.

Element

Inherits: nn.Module

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters) -> None

This is the abstract class for all optical elements in SVETlANNa. It is inherited from torch.nn.Module, so it is PyTorch-compatible. Each element takes an incident wavefront and produces a transmitted wavefront.

📥 Parameters

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

forward


forward(self, incident_wavefront: Wavefront) -> Wavefront

Forward propagation through the optical element.

to_specs


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

make_buffer


make_buffer(self, name: str, value: _T, persistent: bool = False) -> _T

Make buffer for internal use.

Use case in __init__ method:


self.mask = self.make_buffer('mask', some_tensor)

This allow torch to properly process the .to method on the element, since the buffer maask will be transferred to the required device along with simulation parameters. This allows torch to properly process the .to method on the element, since the buffer mask will be transferred to the required device along with simulation parameters.

📥 Parameters

Parameter	Type	Description
`name`	`str`	Name of the new buffer (it is more convenient to use the name of the new attribute).
`value`	`_T`	Tensor to be buffered.
`persistent`	`bool, optional`	See torch docs on buffers, by default `False`.

📤 Returns

_T

The value passed to the method.

process_parameter


process_parameter(self, name: str, value: _V) -> _V

Process element parameter passed by user. Automatically registers buffer for non-parametric tensors.

Use case in __init__ method:


class SomeElement(Element):
    def __init__(self, simulation_parameters, mask, a):
        super().__init__(simulation_parameters)
 
        self.mask = self.process_parameter('mask', mask)
        self.a = self.process_parameter('a', a)
 
        ...

📥 Parameters

Parameter	Type	Description
`name`	`str`	Name of the new buffer (it is more convenient to use the name of the new attribute).
`value`	`_V`	The value of the element parameter.

📤 Returns

_V

The value passed to the method.

AbstractMulElement

Inherits: Element ABC

Class that generalize all apertures with $E^\text{out} = \hat{T}E^\text{in}$ like forward function, where $\hat{T}$ is transmission function.

Properties

transmission_function

The tensor representing transmission function of the element, $\hat{T}$ .

Methods

forward


forward(self, incident_wavefront: Wavefront) -> Wavefront

Aperture

Inherits: AbstractMulElement

Properties

transmission_function

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, mask: OptimizableTensor)

Aperture defined by mask tensor. Commonly, the mask is a tensor with values of either 0 or 1, where 0 represents blocked light and 1 represents allowed light.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`mask`	`torch.Tensor`	Two-dimensional tensor representing the aperture mask of shape `(Ny, Nx)`. The mask works as following:
`None`	`$$E^\text{out}_{xyw...} = \text{mask}_{xy} E^\text{in}_{xyw...}$$`
`None`	`In this case 0 blocks light and 1 allows light go through.`

to_specs


to_specs(self) -> Iterable[ParameterSpecs]

RectangularAperture

Inherits: AbstractMulElement

Properties

transmission_function

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, height: float, width: float)

Rectangular aperture. Through the rectangular area of defined height and width located in the center the light is allowed to pass, otherwise blocked.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`height`	`float`	Aperture height.
`width`	`float`	Aperture width.

to_specs


to_specs(self) -> Iterable[ParameterSpecs]

RoundAperture

Inherits: AbstractMulElement

Properties

transmission_function

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, radius: float)

Round-shaped aperture. Through the round area of defined radius located in the center the light is allowed to pass, otherwise blocked.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters.
`radius`	`float`	Radius of the round-shaped aperture.

to_specs


to_specs(self) -> Iterable[ParameterSpecs]

FreeSpace

Inherits: Element

A class that describes a propagation of the wavefront in free space between two optical elements

Methods

__init__ constructor


__init__(self, simulation_parameters: SimulationParameters, distance: OptimizableFloat, method: Literal['ASM', 'zpASM', 'RSC', 'zpRSC'], total_paddings_x: int | None = None, total_paddings_y: int | None = None)

Init method for FreeSpace class. Defines the parameters and precomputes the parameters for the chosen method of propagation.

📥 Parameters

Parameter	Type	Description
`simulation_parameters`	`SimulationParameters`	Simulation parameters of the system. Contains the information about the spatial grid, wavelength, etc.
`distance`	`OptimizableFloat`	The propagation distance along the optical axis.
`method`	`Literal["ASM", "zpASM", "RSC", "zpRSC"]`	The method used for propagation. 1. ASM - Angular Spectrum Method 2. zpASM - zero-padded Angular Spectrum Method 3. RSC - Rayleigh-Sommerfeld Convolution 4. zpRSC - zero-padded Rayleigh-Sommerfeld Convolution
`total_paddings_x`	`int	None, optional`
`total_paddings_y`	`int	None, optional`

forward


forward(self, incident_wavefront: Wavefront) -> Wavefront

Calculates the wavefront after propagating in the free space

📥 Parameters

Parameter	Type	Description
`incident_wavefront`	`Wavefront`	Wavefront before propagation in free space

📤 Returns

Wavefront

Wavefront after propagation in free space

⚠️ Raises

ValueError — Occurs when a non-existent direct distribution method is chosen

to_specs


to_specs(self) -> Iterable[ParameterSpecs]

Method which determining the specific parameters of the element for visualization in the widget

📤 Returns

Iterable[ParameterSpecs]

Sequence of ParameterSpecs objects containing the parameters of the element

Functions

`one_step_tanh(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor`

A one-step function that can be used in [QuantizerFromStepFunction][svetlanna.elements.slm.QuantizerFromStepFunction]. This function is defined as

$f(x) = \dfrac{\tanh(\alpha(x-0.5))}{2\tanh(\alpha/2)} + \frac{1}{2}$

The parameter $\alpha \in [0, +\infty)$ controls the steepness of the transition. $\alpha=0$ corresponds to a linear function.

📥 Parameters

Parameter	Type	Description
`x`	`torch.Tensor`	Input tensor with values from 0 to 1.
`alpha`	`torch.Tensor`	Steepness control parameter.

📤 Returns

torch.Tensor

Output tensor with values from 0 to 1.

`one_step_cos(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor`

A one-step function that can be used in [QuantizerFromStepFunction][svetlanna.elements.slm.QuantizerFromStepFunction]. This function is defined as

$f(x) = \dfrac{1 - \cos(\pi x^{\alpha + 1})}{2}$

The parameter $\alpha \in [0, +\infty)$ controls the steepness of the transition. $\alpha=0$ corresponds to a function with a smooth transition, but not linear, and as $\alpha$ increases, the transition becomes steeper.

📥 Parameters

Parameter	Type	Description
`x`	`torch.Tensor`	Input tensor with values from 0 to 1.
`alpha`	`torch.Tensor`	Steepness control parameter.

📤 Returns

torch.Tensor

Output tensor with values from 0 to 1.

`QuantizerFromStepFunction(N: int, max_value: float, one_step_function: Callable[Concatenate[torch.Tensor, Params], torch.Tensor]) -> Callable[Concatenate[torch.Tensor, Params], torch.Tensor]`

Create a quantizer function from a given one-step function. The resulting quantizer function takes a tensor of values from 0 to max_value and returns a tensor of values from 0 to max_value with N quantization levels. Each level is defined as the output of the one-step function at the fractional part of the input value, divided by max_value and multiplied by N.

📥 Parameters

Parameter	Type	Description
`N`	`int`	Number of quantization levels.
`max_value`	`float`	Maximum value of the input tensor.
`one_step_function`	`Callable[Concatenate[torch.Tensor, Params], torch.Tensor]`	The one-step function that takes a tensor of values from 0 to 1 as the first argument and returns a tensor of values from 0 to 1. The function should have a steep transition from 0 to 1, and the steepness can be controlled by an additional parameter (for example, alpha). The function should satisfy the following conditions: $f(0) = 0$ , $f(1)=1$ , and $f'(0) = f'(1)$ to ensure that the quantizer function is continuous and smooth.

Examples

```python —
import svetlanna as sv —
from svetlanna.elements.slm import QuantizerFromStepFunction, one_step_tanh —
sv.elements.SpatialLightModulator( — lut_function=sv.PartialWithParameters( QuantizerFromStepFunction( N=256, max_value=2 * torch.pi, one_step_function=one_step_tanh ), alpha=torch.tensor(1.0), ), …
) —
``` —

📤 Returns

Callable[Concatenate[torch.Tensor, Params], torch.Tensor]

Quantizer function that takes the same parameters as the one-step function and applies quantization to the input tensor.

Elements

Classes

ThinLens

Properties

Methods

📥 Parameters

NonlinearElement

Methods

📥 Parameters

SpatialLightModulator

Properties

Methods

📥 Parameters

DiffractiveLayer

Properties

Methods

📥 Parameters

_BufferedValueContainer

Element

Methods

📥 Parameters

📥 Parameters

📤 Returns

📥 Parameters

📤 Returns

AbstractMulElement

Properties

Methods

Aperture

Properties

Methods

📥 Parameters

RectangularAperture

Properties

Methods

📥 Parameters

RoundAperture

Properties

Methods

📥 Parameters

FreeSpace

Methods

📥 Parameters

📥 Parameters

📤 Returns

⚠️ Raises

📤 Returns

Functions

one_step_tanh(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor

📥 Parameters

📤 Returns

one_step_cos(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor

📥 Parameters

📤 Returns

QuantizerFromStepFunction(N: int, max_value: float, one_step_function: Callable[Concatenate[torch.Tensor, Params], torch.Tensor]) -> Callable[Concatenate[torch.Tensor, Params], torch.Tensor]

📥 Parameters

📤 Returns

identity(phase: torch.Tensor) -> torch.Tensor

`one_step_tanh(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor`

`one_step_cos(x: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor`

`QuantizerFromStepFunction(N: int, max_value: float, one_step_function: Callable[Concatenate[torch.Tensor, Params], torch.Tensor]) -> Callable[Concatenate[torch.Tensor, Params], torch.Tensor]`

`identity(phase: torch.Tensor) -> torch.Tensor`