svetlanna.networks package

Submodules

svetlanna.networks.autoencoder module

class svetlanna.networks.autoencoder.LinearAutoencoder(sim_params: SimulationParameters, encoder_elements_list: list[Element] | Iterable[Element], decoder_elements_list: list[Element] | Iterable[Element], to_return: Literal['wf', 'amps'] = 'wf', device: str | device = device(type='cpu'))

Bases: Module

A simple autoencoder network consisting of consistent encoder and decoder for a simultaneous training.

.encode() - forward propagation through encoder elements .decode() - forward propagation through decoder elements

Comment: Works only for a single wavelength, input wavefront [‘h’, ‘w’]

decode(wavefront_encoded)

Propagation through the decoder part – decode an encoded image.

Returns

wavefront_decodedWavefront

A decoded wavefront.

property device: str | device | int

encode(wavefront_in)

Propagation through the encoder part – encode an image wavefront (input).

Returns

wavefront_encodedWavefront

An encoded input wavefront.

forward(wavefront_in)

Parameters

wavefront_in: Wavefront(‘bs’, ‘H’, ‘W’)

Returns

wavefront_encoded, wavefront_decodedtorch.Wavefront

Encoded and decoded wavefronts.

to(device: str | device | int) → LinearAutoencoder

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters

and buffers in this module

dtype (torch.dtype): the desired floating point or complex dtype of

the parameters and buffers in this module

tensor (torch.Tensor): Tensor whose dtype and device are the desired

dtype and device for all parameters and buffers in this module

memory_format (torch.memory_format): the desired memory

format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

svetlanna.networks.diffractive_conv module

class svetlanna.networks.diffractive_conv.ConvDiffNetwork4F(sim_params: SimulationParameters, network_elements_list: list, focal_length: float, conv_phase_mask: Tensor, learnable_mask: bool = False, max_phase: float = 6.283185307179586, fs_method: Literal['fresnel', 'AS'] = 'AS', device: str | device = device(type='cpu'))

Bases: Module

A simple convolutional network with a 4f system as an optical convolutional layer.

Comment: -> [4f system (convolution)] -> [some system of elements] ->

property device: str | device | int

forward(wavefront_in)

Parameters

wavefront_in: Wavefront(‘bs’, ‘H’, ‘W’)

Input wavefront or a batch of Wavefronts.

Returns

: torch.Tensor | Wavefront

Output after a Wavefront propagation through a Convolutional layer and the Other Part of all Network.

to(device: str | device | int) → ConvDiffNetwork4F

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters

and buffers in this module

dtype (torch.dtype): the desired floating point or complex dtype of

the parameters and buffers in this module

tensor (torch.Tensor): Tensor whose dtype and device are the desired

dtype and device for all parameters and buffers in this module

memory_format (torch.memory_format): the desired memory

format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

class svetlanna.networks.diffractive_conv.ConvLayer4F(sim_params: SimulationParameters, focal_length: float, conv_diffractive_mask: Tensor, learnable_mask: bool = False, max_phase: float = 6.283185307179586, fs_method: Literal['fresnel', 'AS'] = 'AS')

Bases: Module

Diffractive convolutional layer based on a 4f system.

forward(input_wf: Wavefront)

Forward propagation through a convolutional diffractive system based on a 4f system.

Parameters

input_wf: Wavefront(‘batch_size’, ‘H’, ‘W’)

An input wavefront(s).

Returns

: Wavefront

A wavefront after a propagation through a system.

get_conv_layer_4f()

get_diffractive_layer()

Returns a DiffractiveLayer according to pre-defined settings from init. It can be trainable or not according to self.learnable_mask flag.

get_free_space()

Returns a FreeSpace of a focal length for a 4f system.

get_thin_lens()

Returns a ThinLens with a pre-defined focal length.

svetlanna.networks.diffractive_rnn module

class svetlanna.networks.diffractive_rnn.DiffractiveRNN(sim_params: SimulationParameters, sequence_len: int, fusing_coeff: float, read_in_layer: Sequential, memory_layer: Sequential, hidden_forward_layer: Sequential, read_out_layer: Sequential, detector_layer: Sequential, device: str | device = device(type='cpu'))

Bases: Module

A simple recurrent diffractive network of an architecture proposed in the article:

https://www.nature.com/articles/s41566-021-00796-w

property device: str | device | int

forward(subsequence_wf: Wavefront)

Parameters

subsequence_wf: Wavefront(‘batch_size’, ‘sequence_len’, ‘H’, ‘W’)

Wavefronts for a sequence. Comment: works for a single wavelength in SimulationParameters!

to(device: str | device | int) → DiffractiveRNN

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters

and buffers in this module

dtype (torch.dtype): the desired floating point or complex dtype of

the parameters and buffers in this module

tensor (torch.Tensor): Tensor whose dtype and device are the desired

dtype and device for all parameters and buffers in this module

memory_format (torch.memory_format): the desired memory

format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

Module contents