Skip to content

Visualization

svetlanna.visualization.show_structure

show_structure(*specsable: Specsable)

Display setup structure in an IPython environment.

This helper renders only the hierarchy of elements (without parameter specs) and is useful for quick notebook previews.

Parameters:

  • *specsable (Specsable, default: () ) –

    One or more Specsable objects to display.

Examples:

import svetlanna as sv
import torch
from svetlanna.visualization import show_structure

Nx = Ny = 128
sim_params = sv.SimulationParameters(
    x=torch.linspace(-1, 1, Nx),
    y=torch.linspace(-1, 1, Ny),
    wavelength=0.1,
)

setup = sv.LinearOpticalSetup(
    [
        sv.elements.RectangularAperture(sim_params, width=0.5, height=0.5),
        sv.elements.FreeSpace(sim_params, distance=0.2, method="AS"),
        sv.elements.DiffractiveLayer(sim_params, mask=torch.rand(Ny, Nx), mask_norm=1),
        sv.elements.FreeSpace(sim_params, distance=0.2, method="AS"),
    ]
)

show_structure(setup)
Output (in IPython environment):

svetlanna.visualization.show_specs

show_specs(*specsable: Specsable) -> SpecsWidget

Display setup structure with interactive specs preview.

Returns:

  • SpecsWidget

    Widget with element tree and per-element specs HTML.

Examples:

import svetlanna as sv
import torch
from svetlanna.visualization import show_specs

Nx = Ny = 128
sim_params = sv.SimulationParameters(
    x=torch.linspace(-1, 1, Nx),
    y=torch.linspace(-1, 1, Ny),
    wavelength=0.1,
)

setup = sv.LinearOpticalSetup(
    [
        sv.elements.RectangularAperture(sim_params, width=0.5, height=0.5),
        sv.elements.FreeSpace(sim_params, distance=0.2, method="AS"),
        sv.elements.DiffractiveLayer(sim_params, mask=torch.rand(Ny, Nx), mask_norm=1),
        sv.elements.FreeSpace(sim_params, distance=0.2, method="AS"),
    ]
)

show_specs(setup)
Output (in IPython environment):

svetlanna.visualization.show_stepwise_forward

show_stepwise_forward(
    *specsable: Specsable,
    input: Tensor,
    simulation_parameters: SimulationParameters,
    types_to_plot: tuple[StepwisePlotTypes, ...] = (
        "I",
        "phase",
    ),
    slices_to_plot: (
        Mapping[str, Index | tuple[Index, ...]] | None
    ) = None
) -> StepwiseForwardWidget

Display stepwise wavefront propagation for setup elements.

The function registers forward hooks on torch.nn.Module elements, runs a forward pass for each provided root element, captures intermediate outputs, renders them as images, and returns an interactive widget.

Parameters:

  • input (Tensor) –

    Input wavefront.

  • simulation_parameters (SimulationParameters) –

    Simulation parameters

  • types_to_plot (tuple[StepwisePlotTypes, ...], default: ('I', 'phase') ) –

    Field properties to plot, by default ("I", "phase").

  • slices_to_plot (Mapping[str, Index | tuple[Index, ...]] | None, default: None ) –

    Axis slices to apply before plotting each captured output. Default is None.

Returns:

  • StepwiseForwardWidget

    Widget containing setup structure and captured per-element outputs.

Examples:

Basic usage:

import svetlanna as sv
import torch
from svetlanna.visualization import show_stepwise_forward

Nx = Ny = 128
sim_params = sv.SimulationParameters(
    x=torch.linspace(-1, 1, Nx),
    y=torch.linspace(-1, 1, Ny),
    wavelength=0.1,
)

setup = sv.LinearOpticalSetup(
    [
        sv.elements.RectangularAperture(sim_params, width=0.5, height=0.5),
        sv.elements.FreeSpace(sim_params, distance=0.2, method="AS"),
        sv.elements.DiffractiveLayer(sim_params, mask=torch.rand(Ny, Nx), mask_norm=1),
        sv.elements.FreeSpace(sim_params, distance=0.2, method="AS"),
    ]
)

input_wavefront = sv.Wavefront.plane_wave(sim_params)
show_stepwise_forward(
    setup,
    input=input_wavefront,
    simulation_parameters=sim_params,
    types_to_plot=("I", "phase", "Re"),
)
Output (in IPython environment):

Spatial slicing with boolean masks:

Use named axis slicing to focus on a region of interest. Plot only the central area:

show_stepwise_forward(
    setup,
    input=input_wavefront,
    simulation_parameters=sim_params,
    slices_to_plot={
        "x": (sim_params.x > -0.5) & (sim_params.x < 0.5),  # x_mask
        "y": (sim_params.y > -0.5) & (sim_params.y < 0.5),  # y_mask
    },
)
# Equivalent to: wavefront[y_mask, x_mask] (axis order depends on sim_params)

Integer indexing for named axes:

Select a specific value from a named axis (e.g., wavelength channel).

# Suppose sim_params has multiple wavelengths,
# so input has shape (wavelength, y, x)
show_stepwise_forward(
    setup,
    input=input_wavefront,
    simulation_parameters=sim_params,
    slices_to_plot={
        "wavelength": 0,
    },
)
# Equivalent to: wavefront[0, :, :]

Slicing unnamed axes (batch dimensions):

Use the special key "_" with a tuple of slices for unnamed leading axes.

# If input has shape (batch, channel, y, x)
show_stepwise_forward(
    setup,
    input=batched_wavefront,
    simulation_parameters=sim_params,
    slices_to_plot={
        "_": (0, 2),  # First batch, third channel
    },
)
# Equivalent to: wavefront[0, 2, :, :]

Combining named and unnamed slicing:

You can mix both approaches. Named axes override positional slices from "_".

# If input has shape (batch, wavelength, y, x) where wavelength is named axes in sim_params
show_stepwise_forward(
    setup,
    input=batched_wavefront,
    simulation_parameters=sim_params,
    slices_to_plot={
        "_": (0, 0),        # First batch, first wavelength (from position)
        "wavelength": 1,    # Override wavelength to second
    },
)
# Result: wavefront[0, 1, :, :]