Skip to content

Element

svetlanna.elements.Element 

Element(simulation_parameters: SimulationParameters)

Bases: Module

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:

forward abstractmethod 

forward(incident_wavefront: Wavefront) -> Wavefront

Forward propagation through the optical element.

make_buffer 

make_buffer(
    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:

  • 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, default: False ) –

    See torch docs on buffers, by default False.

Returns:

  • _T –

    The value passed to the method.

process_parameter 

process_parameter(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:

  • 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.