Tomography Module

PyTorch functions for 2D and 3D tomography and Radon transform simulation.

This module includes differentiable and non-differentiable implementations of: - Forward and back projection routines (Radon and inverse Radon transforms) in 2D and 3D. - Iterative reconstruction methods such as SIRT and CGLS using functional forward/back projectors. - Sparse matrix-based forward and backward operations (A-matrix formulation) using PyTorch sparse tensors. - Utility functions for constructing sparse system matrices, rotating images, and defining affine transforms.

Main features: - Differentiable 3D forward and backward projectors using torchvision.transforms.functional.rotate. - Support for iterative solvers: SIRT (with normalization) and CGLS. - Conversion utilities for using SciPy sparse matrices in PyTorch (e.g., Amatrix_torch, Sino_torch, Amatrix_rec). - Grid-based rotation of images via affine transformation (imrotate_torch). - Compatible with both CPU and CUDA devices.

Author: Antony Vamvakeros

nDTomo.pytorch.tomo_torch.Amatrix_rec(AtorchT, s, ntr)[source]

Reconstructs an image from a sinogram using the transpose of the A matrix.

Parameters:

AtorchT (torch.sparse.FloatTensor) – Transpose of the system matrix A (shape: [ntr*ntr, npr*ntr]).
s (torch.Tensor) – Sinogram of shape (npr, ntr).
ntr (int) – Number of translation steps (output image side length).

Returns:

rec – Reconstructed image tensor of shape (ntr, ntr).

Return type:

torch.Tensor

nDTomo.pytorch.tomo_torch.Amatrix_sino(Atorch, im, npr, ntr)[source]

Computes the sinogram using matrix multiplication with a sparse A matrix.

Parameters:

Atorch (torch.sparse.FloatTensor) – Sparse system matrix A (shape: [npr*ntr, ntr*ntr]).
im (torch.Tensor) – Flattened image tensor of shape (ntr*ntr, 1).
npr (int) – Number of projections (angles).
ntr (int) – Number of translation steps.

Returns:

stf – Sinogram tensor of shape (npr, ntr).

Return type:

torch.Tensor

nDTomo.pytorch.tomo_torch.Amatrix_torch(A, gpu=True)[source]

Converts a SciPy sparse matrix A to a PyTorch sparse tensor.

Parameters:

A (scipy.sparse matrix) – Input sparse matrix in COO format (or convertible to COO).
gpu (bool, optional) – If True, moves the tensor to CUDA. Default is True.

Returns:

Atorch – PyTorch sparse tensor version of A.

Return type:

torch.sparse.FloatTensor

nDTomo.pytorch.tomo_torch.RotMat(theta)[source]

Creates a 2D rotation matrix for use in affine transformations.

Parameters:: theta (float or torch.Tensor) – Rotation angle in radians.
Returns:: rotmat – A 2×3 affine rotation matrix tensor.
Return type:: torch.Tensor

nDTomo.pytorch.tomo_torch.Sino_torch(Atorch, im, ntr, npr)[source]

Generates a sinogram by applying a torch sparse A matrix to an image.

Parameters:

Atorch (torch.sparse.FloatTensor) – Sparse A matrix in torch format.
im (ndarray) – Input image as a NumPy array of shape (ntr, ntr).
ntr (int) – Number of translation steps (image side length).
npr (int) – Number of projections.

Returns:

s – Sinogram as a tensor of shape (npr, ntr).

Return type:

torch.Tensor

nDTomo.pytorch.tomo_torch.back_project_3D(sinos, angles, npix, nch, device='cuda')[source]

Perform backprojection (inverse Radon transform) of a 3D sinogram using PyTorch.

Parameters:

sinos (torch.Tensor) – Input sinogram of shape (nch, npix, len(angles)), where: - nch is the number of slices or channels, - npix is the number of detector elements, - len(angles) is the number of projection angles.
angles (list or ndarray) – List of projection angles in degrees.
npix (int) – Number of pixels in the output reconstructed image.
nch (int) – Number of slices or channels in the volume.
device (str, optional) – PyTorch device string (default: ‘cuda’).

Returns:

vol – Reconstructed 3D volume of shape (nch, npix, npix). Each channel corresponds to a separate slice, reconstructed via filtered or unfiltered backprojection.

Return type:

torch.Tensor

nDTomo.pytorch.tomo_torch.cgls_pytorch_functional(sinos, angles, npix, nch=1, n_iter=10, device='cuda')[source]

CGLS reconstruction using PyTorch with functional forward and back projectors.

Parameters:

sinos (torch.Tensor) – Input sinogram tensor of shape (nch, npix, n_angles), e.g., (1, 151, 180).
angles (list or ndarray) – List of projection angles in degrees.
npix (int) – Width/height of the reconstructed image.
nch (int) – Number of slices or channels in the volume (default = 1).
n_iter (int) – Number of CGLS iterations.
device (str) – Computation device, e.g., ‘cuda’.

Returns:

Reconstructed volume of shape (nch, npix, npix).

Return type:

torch.Tensor

nDTomo.pytorch.tomo_torch.compute_W_ray(angles, npix, nch, device='cuda')[source]

Compute ray normalization weights for each voxel via forward projection of a constant volume.

This function simulates the accumulation of contributions each detector sees from a uniform volume, useful for SIRT or SART-type normalization.

Parameters:

angles (list or ndarray) – List of projection angles in degrees.
npix (int) – Number of pixels in each dimension of the image (image size).
nch (int) – Number of slices or channels in the 3D volume.
device (str, optional) – PyTorch device string (default: ‘cuda’).

Returns:

W_ray – Weighting map of shape (nch, npix, len(angles)) representing forward projection of ones.

Return type:

torch.Tensor

nDTomo.pytorch.tomo_torch.create_torch_Amat(Acoo, values, indices, shape, device='cuda')[source]

Constructs a PyTorch sparse A matrix from COO components.

Parameters:

Acoo (scipy.sparse.coo_matrix) – Sparse matrix object (for context).
values (ndarray) – Non-zero values of the matrix.
indices (ndarray) – 2×N array of row and column indices.
shape (tuple) – Shape of the matrix.
device (str, optional) – Target device (‘cuda’ or ‘cpu’).

Returns:

Amat – Constructed sparse matrix on the target device.

Return type:

torch.sparse.FloatTensor

nDTomo.pytorch.tomo_torch.forward_project_3D(vol, angles, npix, nch, device='cuda')[source]

Perform forward projection (Radon transform) of a 3D volume using PyTorch.

Parameters:

vol (torch.Tensor) – Input volume of shape (1, nch, npix, npix), where: - 1 is the batch dimension, - nch is the number of channels or slices (e.g., spectral bins or time steps), - npix x npix is the spatial resolution of each slice.
angles (list or ndarray) – List of projection angles in degrees.
npix (int) – Number of pixels along each projection axis (image size).
nch (int) – Number of slices or channels in the volume.
device (str, optional) – PyTorch device string (default: ‘cuda’).

Returns:

sinos – Simulated sinogram of shape (nch, npix, len(angles)), where each slice corresponds to a different channel/slice and each column is a projection.

Return type:

torch.Tensor

nDTomo.pytorch.tomo_torch.imrotate_torch(im, theta, dtype=<class 'nDTomo.pytorch._mock_torch.DummyTensor'>)[source]

Rotates a 2D image (or batch) using an affine transformation.

Parameters:

im (torch.Tensor) – Input tensor of shape (N, C, H, W).
theta (float) – Rotation angle in radians.
dtype (torch dtype, optional) – Data type of the affine matrix and grid. Default is float tensor on GPU.

Returns:

imr – Rotated image of the same shape as input.

Return type:

torch.Tensor

nDTomo.pytorch.tomo_torch.sirt_pytorch_functional(sinos, angles, npix, nch=1, n_iter=20, relax=0.01, epsilon=1e-06, device='cuda')[source]

SIRT reconstruction using PyTorch with function-based forward and backward projectors.

Parameters:

sinos (torch.Tensor) – Input sinogram tensor of shape (nch, npix, n_angles), e.g. (1, 151, 180).
angles (list or ndarray) – List of projection angles in degrees.
npix (int) – Width/height of the reconstructed image.
nch (int) – Number of slices or channels in the volume (default = 1).
n_iter (int) – Number of SIRT iterations.
relax (float) – Relaxation factor (typically small, e.g., 0.01).
epsilon (float) – Small number to avoid division by zero.
device (str) – Computation device, e.g. ‘cuda’.

Returns:

Reconstructed volume of shape (nch, npix, npix).

Return type:

torch.Tensor