Models Module

Neural networks models

@author: Antony Vamvakeros

class nDTomo.pytorch.models_torch.CNN1D(nch_in=1, nch_out=1, nfilts=32, nlayers=4, norm_type='layer', activation='Linear')

Bases: object

A 1D convolutional neural network for sequential or spectral data processing with optional normalization and residual connections.

Parameters:

• nch_in (int) – Number of input channels.

• nch_out (int) – Number of output channels.

• nfilts (int) – Number of filters in the convolutional layers.

• nlayers (int) – Number of intermediate convolutional blocks.

• norm_type (str or None) – Type of normalization: ‘batch’, ‘layer’, or None.

• activation (str) – Final activation type. If ‘Sigmoid’, a Sigmoid activation is appended after the last layer.

• Forward

• -------

• x (torch.Tensor) – Input tensor of shape (batch_size, nch_in, sequence_length).

• residual (bool) – If True, adds input x to the output (residual connection).

Returns:

Output tensor of shape (batch_size, nch_out, sequence_length).

Return type:

torch.Tensor

add_norm_layer(layers, nfilts, norm_type)

forward(x, residual=False)

class nDTomo.pytorch.models_torch.CNN2D(npix, nch_in=1, nch_out=1, nfilts=32, nlayers=4, norm_type='layer', activation='Linear')

Bases: object

A configurable 2D Convolutional Neural Network (CNN) for image processing tasks.

Parameters:

• npix (int) – The spatial size (height and width) of the input images.

• nch_in (int, optional, default=1) – Number of input channels (e.g., 1 for grayscale images, 3 for RGB).

• nch_out (int, optional, default=1) – Number of output channels.

• nfilts (int, optional, default=32) – Number of filters (channels) in the intermediate convolutional layers.

• nlayers (int, optional, default=4) – Number of intermediate convolutional layers.

• norm_type (str or None, optional, default='layer') – Normalization type to apply after convolutions: - ‘batch’: Batch normalization. - ‘layer’: Layer normalization. - None: No normalization.

• activation (str, optional, default='Linear') – Final activation function to apply: - ‘Sigmoid’: Applies a sigmoid activation. - ‘Linear’: No activation applied.

forward(x, residual=False):

Performs a forward pass through the network. - residual: If True, adds the input x to the output.

add_norm_layer(layers, nfilts, norm_type)

Adds a normalization layer to the list of layers.

Parameters:

• layers (list) – List of layers to which the normalization layer is appended.

• nfilts (int) – Number of channels for normalization.

• norm_type (str) – Type of normalization (‘batch’ or ‘layer’).

forward(x, residual=False)

Performs a forward pass through the network.

Parameters:

• x (torch.Tensor) – Input tensor of shape (batch_size, nch_in, npix, npix).

• residual (bool, optional, default=False) – If True, adds the input x to the output.

Returns:

Output tensor of shape (batch_size, nch_out, npix, npix).

Return type:

torch.Tensor

class nDTomo.pytorch.models_torch.CNN3D(npix, nch_in=1, nch_out=1, nfilts=32, nlayers=4, norm_type='layer', activation='Linear')

Bases: object

A 3D convolutional neural network for volumetric data processing with optional normalization and configurable depth.

Parameters:

• npix (int) – Size of the 3D cube (assumes cube of shape npix x npix x npix).

• nch_in (int) – Number of input channels.

• nch_out (int) – Number of output channels.

• nfilts (int) – Number of filters in convolutional layers.

• nlayers (int) – Number of intermediate convolutional blocks.

• norm_type (str or None) – Type of normalization: ‘batch’, ‘layer’, or None.

• activation (str) – Final activation type. If ‘Sigmoid’, a Sigmoid activation is appended.

• Forward

• -------

• x (torch.Tensor) – Input tensor of shape (batch_size, nch_in, D, H, W).

• residual (bool) – If True, adds the input x to the output (residual connection).

Returns:

Output tensor of shape (batch_size, nch_out, D, H, W).

Return type:

torch.Tensor

add_norm_layer(layers, nfilts, norm_type)

forward(x, residual=False)

class nDTomo.pytorch.models_torch.ConvBlock(in_channels, out_channels, spatial_dims, norm_type='layer')

Bases: object

A basic convolutional block with two Conv2D layers, LayerNorm (by default), and ReLU activation.

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of output channels.

• spatial_dims (tuple[int, int]) – Spatial dimensions of the input (height, width).

• norm_type (str) – Normalization type: ‘layer’ (default) or ‘batch’.

Returns:

Output tensor of shape (batch_size, out_channels, H, W).

Return type:

torch.Tensor

forward(x)

class nDTomo.pytorch.models_torch.Crop2D(top=0, bottom=0, left=0, right=0)

Bases: object

Crop a 4D tensor along the height and width dimensions.

Parameters:

• top (int, optional) – Number of pixels to remove from the top. Default is 0.

• bottom (int, optional) – Number of pixels to remove from the bottom. Default is 0.

• left (int, optional) – Number of pixels to remove from the left. Default is 0.

• right (int, optional) – Number of pixels to remove from the right. Default is 0.

Notes

• Expects input tensors in NCHW format: (batch, channels, height, width).

• If any crop value is 0, that side is left unchanged.

• This layer performs a simple slice, without resizing or interpolation.

Examples

>>> crop = Crop2D(top=1, bottom=2, left=0, right=1)
>>> x = torch.randn(2, 3, 10, 10)
>>> y = crop(x)
>>> y.shape
torch.Size([2, 3, 7, 9])

forward(x)

class nDTomo.pytorch.models_torch.DownBlock(in_channels, out_channels, spatial_dims, norm_type=None)

Bases: object

A downsampling block consisting of a Conv2D with stride=2 followed by a ConvBlock.

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of output channels.

• spatial_dims (tuple[int, int]) – Spatial dimensions after downsampling.

• norm_type (str or None) – Type of normalization to use in the ConvBlock.

Returns:

Output tensor of shape (batch_size, out_channels, H/2, W/2).

Return type:

torch.Tensor

forward(x)

class nDTomo.pytorch.models_torch.PeakFitCNN(nch_in=1, nch_out=1, nfilts=32, upscale_factor=4, norm_type='instance', activation='Linear', padding='same', npix=None)

Bases: object

A 2D convolutional neural network designed for upsampling and refining spectral or spatial peak data, optionally doubling or quadrupling the input resolution using bilinear interpolation and CNN blocks.

Parameters:

• nch_in (int) – Number of input channels.

• nch_out (int) – Number of output channels.

• nfilts (int) – Number of filters in the intermediate convolution layers.

• upscale_factor (int) – Upscaling factor for the input. Supported values: 2 or 4.

• norm_type (str) – Type of normalization to apply: ‘instance’, ‘batch’, or ‘layer’.

• activation (str) – Final activation function: ‘Linear’, ‘ReLU’, ‘Sigmoid’, or ‘LeakyReLU’.

• padding (str) – Padding mode for convolutions (‘same’ or ‘valid’).

• npix (int) – Number of pixels in the input (required for LayerNorm).

• Forward

• -------

• x (torch.Tensor) – Input tensor of shape (batch_size, nch_in, H, W).

Returns:

Output tensor of shape (batch_size, nch_out, H * upscale_factor, W * upscale_factor).

Return type:

torch.Tensor

forward(x)

class nDTomo.pytorch.models_torch.PeakModel(prms, device='cuda')

Bases: object

A peak fitting model that represents parameterized 1D functions (e.g. Gaussian or Pseudo-Voigt) using learnable normalized parameters constrained to [0, 1]. Converts normalized parameters to their physical range before evaluating the function.

Parameters:

• prms (dict) – Dictionary containing: - ‘val’: torch.Tensor of shape (n_params, npix, npix), the initial normalized parameters. - ‘min’: dict of minimum values for each parameter. - ‘max’: dict of maximum values for each parameter.

• device (str) – Device to place parameters on (‘cuda’ or ‘cpu’).

• Forward

• -------

• x (torch.Tensor) – Input tensor of shape (N, 1) representing the x-axis values for function evaluation.

• model (str) – Peak function model to use: ‘Gaussian’ or ‘PseudoVoigt’.

Returns:

y – Output tensor of shape (N, 1), the evaluated function for each pixel.

Return type:

torch.Tensor

forward(x, model='Gaussian')

class nDTomo.pytorch.models_torch.PrmCNN2D(npix, nch_in=1, nch_out=1, nfilts=32, nlayers=4, norm_type='layer', prms_layer=True, cnn_layer=True, tensor_vals='random', tensor_initial=None, padding='same')

Bases: object

A flexible 2D model that combines a trainable tensor (image parameterization) with an optional CNN-based processing module. Can operate in three modes: - Pure parameterization (learned image). - CNN only (applies CNN to input). - Parameterization + CNN (CNN applied to learned image).

Parameters:

• npix (int) – Image resolution (assumes square images).

• nch_in (int) – Number of input channels.

• nch_out (int) – Number of output channels.

• nfilts (int) – Number of filters in CNN layers.

• nlayers (int) – Number of intermediate CNN blocks (excluding first and last layers).

• norm_type (str) – Type of normalization: ‘layer’, ‘instance’, or ‘batchnorm’.

• prms_layer (bool) – If True, a learnable tensor is used as input.

• cnn_layer (bool) – If True, a CNN processes the input or parameter tensor.

• tensor_vals (str) – Initialization mode for the learned tensor: ‘random’, ‘zeros’, ‘ones’, ‘mean’, ‘random_positive’, or ‘custom’.

• tensor_initial (torch.Tensor or None) – Custom tensor to use if tensor_vals == ‘custom’.

• padding (str) – Padding mode for convolutions (‘same’ or ‘valid’).

• Forward

• -------

• x (torch.Tensor) – Input tensor if cnn_layer=True and prms_layer=False. Ignored otherwise.

Returns:

Output tensor of shape (1, nch_out, npix, npix).

Return type:

torch.Tensor

forward(x)

class nDTomo.pytorch.models_torch.ResNet2D(npix, nch_in=1, nch_out=1, nfilts=32, n_res_blocks=4, norm_type='layer', activation='Linear')

Bases: object

A 2D ResNet-style convolutional network for image processing tasks with configurable residual blocks and optional normalization.

Parameters:

• npix (int) – Spatial size of the input (assumed square).

• nch_in (int) – Number of input channels.

• nch_out (int) – Number of output channels.

• nfilts (int) – Number of filters in convolutional layers.

• n_res_blocks (int) – Number of ResNet blocks in the model.

• norm_type (str or None) – Type of normalization: ‘batch’, ‘layer’, or None.

• activation (str) – Final activation function: ‘Linear’ or ‘Sigmoid’.

Returns:

Output tensor of shape (batch_size, nch_out, npix, npix).

Return type:

torch.Tensor

add_norm_layer(nfilts, norm_type)

forward(x, residual=False)

class nDTomo.pytorch.models_torch.ResNetBlock(nfilts, npix, kernel_size=3, stride=1, padding=1, norm_type=None)

Bases: object

A basic 2D residual block consisting of two convolutional layers with optional normalization and ReLU activation, followed by a skip connection.

Parameters:

• nfilts (int) – Number of filters in the convolutional layers.

• npix (int) – Spatial dimension (assumes square input for LayerNorm).

• kernel_size (int) – Size of convolutional kernel.

• stride (int) – Stride for the convolution.

• padding (int) – Padding to apply to convolution.

• norm_type (str or None) – Type of normalization: ‘batch’, ‘layer’, or None.

Returns:

Output tensor of same shape as input.

Return type:

torch.Tensor

add_norm_layer(nfilts)

forward(x)

class nDTomo.pytorch.models_torch.SD2I(npix: int, factor: int = 8, upsample: bool = True)

Bases: object

SD2I (Single Digit to Image) reconstruction network.

This is a PyTorch reimplementation of the Keras/TensorFlow SD2I model. It reconstructs a 2D image from a single scalar input using a sequence of dense layers, optional progressive upsampling, and convolutional refinement.

Parameters:

• npix (int) – Target output image size (height = width = npix).

• factor (int, optional) – Number of feature channels after the dense-to-feature-map reshape. Default is 8.

• upsample (bool, optional) – If True, use the progressive upsampling branch (two stages of ×2 upsampling). If False, use the direct full-resolution branch without upsampling. Default is True.

• (upsample=True) (Architecture)

• ----------------------------

• stack (3. Conv) – ceil(npix / 4) × ceil(npix / 4) × factor features.

• (factor (2. Reshape to)

• H0

• W0)

• 4). (where H0 = W0 = ceil(npix /)

• neighbor). (5. Upsample ×2 (nearest) –

  • Conditional crop: remove 1 pixel from bottom & right if npix % 4 in {1, 2}.

• block (6. Conv)

• neighbor). –

  • Conditional crop: remove 1 pixel from top & left if npix is odd.

• block

• activation). (7. Final 3×3 conv to 1 channel (linear)

• (upsample=False) (Architecture)

• -----------------------------

• stack

• (factor

• npix

• npix).

• stack –

  • 3×3 conv, 128 filters, ReLU

  • 3×3 conv, 128 filters, ReLU

  • 3×3 conv, 64 filters, ReLU

  • 3×3 conv, 1 filter, linear

Notes

• All convolutions use “same” padding (padding=1 for kernel_size=3).

• Upsampling uses nearest-neighbor to match Keras UpSampling2D.

• Weights are initialized with a RandomNormal(mean=0.0, std=0.05) to match Keras kernel_initializer=’random_normal’.

Examples

>>> model = SD2I(npix=128, factor=8, upsample=True)
>>> inp = torch.randn(4, 1)  # batch of scalars
>>> out = model(inp)
>>> out.shape
torch.Size([4, 1, 128, 128])

forward(x)

class nDTomo.pytorch.models_torch.UNet2D(nch_in, nch_out, npix, base_nfilts=64, num_blocks=4, norm_type=None, activation='Linear')

Bases: object

A U-Net style 2D convolutional neural network with downsampling and upsampling paths and skip connections. Supports configurable depth and normalization.

Parameters:

• nch_in (int) – Number of input channels.

• nch_out (int) – Number of output channels.

• npix (int) – Input spatial dimension (assumes square input).

• base_nfilts (int) – Number of filters in the base layer.

• num_blocks (int) – Number of downsampling and upsampling blocks.

• norm_type (str or None) – Normalization type: ‘batch’, ‘layer’, or None.

• activation (str) – Final activation function: ‘Sigmoid’ or ‘Linear’.

• Forward

• -------

• x (torch.Tensor) – Input tensor of shape (batch_size, nch_in, npix, npix).

Returns:

Output tensor of shape (batch_size, nch_out, npix, npix).

Return type:

torch.Tensor

forward(x)

class nDTomo.pytorch.models_torch.UpBlock(in_channels, bridge_channels, out_channels, spatial_dims, norm_type='layer')

Bases: object

An upsampling block with transposed convolution and concatenation with a skip connection (bridge), followed by a ConvBlock.

Parameters:

• in_channels (int) – Number of channels to upsample.

• bridge_channels (int) – Number of channels in the bridge tensor (from encoder).

• out_channels (int) – Number of output channels after convolution.

• spatial_dims (tuple[int, int]) – Spatial dimensions of the output.

• norm_type (str) – Normalization type to use in ConvBlock: ‘layer’ or ‘batch’.

Returns:

Output tensor of shape (batch_size, out_channels, H*2, W*2).

Return type:

torch.Tensor

forward(x, bridge)

class nDTomo.pytorch.models_torch.VolumeModel(npix, num_slices, vol=None, device='cuda')

Bases: object

A simple learnable 3D volume model where each slice in the volume is trainable. Supports differential updates or direct use of the internal volume.

Parameters:

• npix (int) – Number of pixels in each spatial dimension (assumes square slices).

• num_slices (int) – Number of slices along the depth axis.

• vol (torch.Tensor or None) – Optional initial volume of shape (num_slices, npix, npix). If None, initializes to zeros.

• device (str) – Device to place the model parameters on (‘cuda’ or ‘cpu’).

• Forward

• -------

• input_volume (torch.Tensor) – External volume to add to the internal volume (shape: num_slices, npix, npix).

• diff (bool) – If True, output is input_volume + self.volume. If False, output is self.volume only.

Returns:

Transformed volume of shape (num_slices, npix, npix).

Return type:

torch.Tensor

forward(input_volume, diff=True)