Weighted backproject demo standalone script

wbp-napari-demo.py

from typing import Sequence, Tuple, Literal
import einops
import mrcfile
import napari
from magicgui import magicgui
from napari.types import ImageData

import numpy as np
import torch
from torch.nn import functional as F

VOLUME_FILE = 'ribo-16Apx.mrc'

volume = torch.tensor(mrcfile.read(VOLUME_FILE))
volume_shape = torch.tensor(volume.shape)

volume_center = volume_shape // 2
tilt_image_center = volume_center[:2]


def slice_dft(
        dft: torch.Tensor,
        slice_coordinates: torch.Tensor
) -> torch.Tensor:
    """Sample batches of 2D images from a complex cubic volume at specified coordinates.

    Notes
    -----
    - dft should pre-fftshifted to place the origin in Fourier space at the center of the DFT
        i.e. dft should be the result of
            volume -> fftshift(volume) -> fft3(volume) -> fftshift(volume)
    - coordinates should be zyx ordered, match image array dimensions
    - coordinates should be array coordinates, [0, N-1] for a dimension of length N.


    Parameters
    ----------
    dft: torch.Tensor
        (d, h, w) complex valued cubic volume (d == h == w) containing the discrete Fourier transform
        of a cubic volume.
    slice_coordinates: torch.Tensor
        (batch, h, w, zyx) array of coordinates at which `dft` should be sampled.

    Returns
    -------
    samples: torch.Tensor
        (batch, h, w) array of complex valued images sampled from the `dft`
    """

    # cannot sample complex tensors directly with grid_sample
    # c.f. https://github.com/pytorch/pytorch/issues/67634
    # workaround: treat real and imaginary parts as separate channels
    dft = einops.rearrange(torch.view_as_real(dft), 'd h w complex -> complex d h w')
    n_slices = slice_coordinates.shape[0]
    dft = einops.repeat(dft, 'complex d h w -> b complex d h w', b=n_slices)
    slice_coordinates = array_coordinates_to_grid_sample_coordinates(slice_coordinates,
                                                                     array_shape=dft.shape[-3:])
    slice_coordinates = einops.rearrange(slice_coordinates,
                                         'b h w xyz -> b 1 h w xyz')  # add depth dim
    samples = F.grid_sample(
        input=dft,
        grid=slice_coordinates,
        mode='bilinear',  # this is trilinear when input is volumetric
        padding_mode='zeros',
        align_corners=False,
    )
    samples = einops.rearrange(samples, 'b complex 1 h w -> b h w complex')
    samples = torch.view_as_complex(samples.contiguous())
    return samples  # (b, h, w)


def backproject(
        image_stack: torch.Tensor,  # (b, h, w)
        projection_matrices: torch.Tensor,  # (b, 4, 4)
        output_dimensions: Tuple[int, int, int]
) -> torch.Tensor:
    grid_coordinates = get_grid_coordinates(output_dimensions)  # (d, h, w, zyx)
    grid_coordinates = torch.flip(grid_coordinates, dims=(-1,))  # (d, h, w, xyz)
    grid_coordinates = homogenise_coordinates(grid_coordinates)  # (d, h, w, xyzw)
    grid_coordinates = einops.rearrange(grid_coordinates, 'd h w xyzw -> d h w 1 xyzw 1')
    projected_coordinates = projection_matrices[..., :2, :] @ grid_coordinates  # only xy coords
    projected_coordinates = einops.rearrange(
        projected_coordinates, 'd h w img xy 1 -> d h w img xy'
    )
    image_stack_coordinates = add_implied_coordinate_from_dimension(projected_coordinates, dim=3)
    image_stack_coordinates = einops.rearrange(
        image_stack_coordinates, 'd h w img xyz -> img d h w xyz'
    )
    image_stack_coordinates = torch.flip(image_stack_coordinates, dims=(-1,))  # xyz -> zyx
    image_stack_coordinates = array_coordinates_to_grid_sample_coordinates(
        image_stack_coordinates, array_shape=image_stack.shape
    )
    n_images = image_stack_coordinates.shape[0]
    image_stack = einops.repeat(image_stack, 'b h w -> img 1 b h w',
                                img=n_images)  # (b, c, d, h, w) for sampling
    samples = F.grid_sample(
        input=image_stack,
        grid=image_stack_coordinates,
        mode='bilinear',  # this is trilinear when input is volumetric
        padding_mode='zeros',
        align_corners=False,
    )
    reconstruction = einops.reduce(samples, 'img 1 d h w -> d h w', reduction='sum')
    return reconstruction


def Rx(angles_degrees: torch.Tensor) -> torch.Tensor:
    """4x4 matrices for a rotation of homogenous coordinates (xyzw) around the X-axis."""
    angles_degrees = torch.tensor(angles_degrees).reshape(-1)
    angles_radians = torch.deg2rad(angles_degrees)
    c = torch.cos(angles_radians)
    s = torch.sin(angles_radians)
    matrices = einops.repeat(torch.eye(4), 'i j -> n i j', n=len(angles_degrees)).clone()
    matrices[:, 1, 1] = c
    matrices[:, 1, 2] = -s
    matrices[:, 2, 1] = s
    matrices[:, 2, 2] = c
    return torch.squeeze(matrices)


def Ry(angles_degrees: torch.Tensor) -> torch.Tensor:
    """4x4 matrices for a rotation of homogenous coordinates (xyzw) around the Y-axis."""
    angles_degrees = torch.tensor(angles_degrees).reshape(-1)
    angles_radians = torch.deg2rad(angles_degrees)
    c = torch.cos(angles_radians)
    s = torch.sin(angles_radians)
    matrices = einops.repeat(torch.eye(4), 'i j -> n i j', n=len(angles_degrees)).clone()
    matrices[:, 0, 0] = c
    matrices[:, 0, 2] = s
    matrices[:, 2, 0] = -s
    matrices[:, 2, 2] = c
    return torch.squeeze(matrices)


def Rz(angles_degrees: torch.Tensor) -> torch.Tensor:
    """4x4 matrices for a rotation of homogenous coordinates (xyzw) around the Z-axis."""
    angles_degrees = torch.tensor(angles_degrees).reshape(-1)
    angles_radians = torch.deg2rad(angles_degrees)
    c = torch.cos(angles_radians)
    s = torch.sin(angles_radians)
    matrices = einops.repeat(torch.eye(4), 'i j -> n i j', n=len(angles_degrees)).clone()
    matrices[:, 0, 0] = c
    matrices[:, 0, 1] = -s
    matrices[:, 1, 0] = s
    matrices[:, 1, 1] = c
    return torch.squeeze(matrices)


def S(shifts: torch.Tensor) -> torch.Tensor:
    """4x4 matrices for shifts.
    Shifts supplied can be 2D or 3D.
    """
    shifts = torch.tensor(shifts)
    if shifts.shape[-1] == 2:
        shifts = promote_2d_to_3d(shifts)
    shifts = shifts.reshape((-1, 3))
    matrices = einops.repeat(torch.eye(4), 'i j -> n i j', n=shifts.shape[0]).clone()
    matrices[:, 0:3, 3] = shifts
    return torch.squeeze(matrices)


def get_grid_coordinates(grid_dimensions: Sequence[int]) -> torch.Tensor:
    indices = torch.tensor(np.indices(grid_dimensions)).float()  # (coordinates, *grid_dimensions)
    return einops.rearrange(indices, 'coordinates ... -> ... coordinates')


def promote_2d_to_3d(shifts: torch.Tensor) -> torch.Tensor:
    """Promote 2D coordinates to 3D with zeros in the last dimension.

    Last dimension of array should be of length 2.
    """
    shifts = F.pad(torch.tensor(shifts), pad=(0, 1), mode='constant', value=0)
    return torch.squeeze(shifts)


def homogenise_coordinates(coords: torch.Tensor) -> torch.Tensor:
    """3D coordinates to 4D homogenous coordinates with ones in the last dimension.

    Last dimension of array should be of length 3.
    """
    return F.pad(torch.Tensor(coords), pad=(0, 1), mode='constant', value=1)


def generate_rotated_slice_coordinates(rotations: torch.Tensor, n: int) -> torch.Tensor:
    """Generate an array of rotated central slice coordinates for sampling a 3D image.

    Notes
    -----
    - rotation matrices rotate coordinates ordered xyz
    - coordinates returned are ordered zyx to match volumetric array indices

    Parameters
    ----------
    rotations: torch.Tensor
        (batch, 3, 3) array of rotation matrices which rotate xyz coordinates.
    n: int
        sidelength of cubic grid for which coordinates are generated.

    Returns
    -------
    coordinates: torch.Tensor
        (batch, n, n, zyx) array of coordinates.
    """
    # generate [x, y, z] coordinates for a central slice
    # the slice spans the XY plane with origin on DFT center
    x = y = torch.arange(n) - (n // 2)
    xx = einops.repeat(x, 'w -> h w', h=n)
    yy = einops.repeat(y, 'h -> h w', w=n)
    zz = torch.zeros(size=(n, n))
    xyz = einops.rearrange([xx, yy, zz], 'xyz h w -> 1 h w xyz 1')

    # rotate coordinates
    rotations = einops.rearrange(rotations, 'b i j -> b 1 1 i j')
    xyz = einops.rearrange(rotations @ xyz, 'b h w xyz 1 -> b h w xyz')

    # recenter slice on DFT center and flip to zyx
    xyz += n // 2
    zyx = torch.flip(xyz, dims=(-1,))
    return zyx


def add_implied_coordinate_from_dimension(
        coordinates: torch.Tensor, dim: int, prepend_new_coordinate: bool = False
) -> torch.Tensor:
    """Make an implicit coordinate in a multidimensional arrays of coordinates explicit.

    For an array of coordinates with shape (n, t, 3), this function produces an array of
    shape (n, t, 4). The values in the new column reflect the position of the coordinate in `dim`.
    `prepend_new_coordinate` controls whether the new coordinate is prepended
    (`prepend_new_coordinate=True`) or appended (`prepend_new_coordinate=False`) to the existing
    coordinates.

    Parameters
    ----------
    coordinates: torch.Tensor
        (..., d) array of coordinates where d is the dimensionality of coordinates.
    dim: int
        dimension from which the value of the new coordinate will be inferred.
    prepend_new_coordinate: bool
        controls whether the new coordinate is prepended or appended to existing coordinates.

    Returns
    -------
    coordinates: torch.Tensor
        (..., d+1)
    """
    if prepend_new_coordinate is True:
        pad, new_coordinate_index = (1, 0), 0
    else:  # append
        pad, new_coordinate_index = (0, 1), -1
    output = F.pad(coordinates, pad=pad, mode='constant', value=0)
    output[..., new_coordinate_index] = torch.arange(coordinates.shape[dim])
    return output


def _array_coordinates_to_grid_sample_coordinates_1d(
        coordinates: torch.Tensor, dim_length: int
) -> torch.Tensor:
    return (coordinates / (0.5 * dim_length - 0.5)) - 1


def _grid_sample_coordinates_to_array_coordinates_1d(
        coordinates: torch.Tensor, dim_length: int
) -> torch.Tensor:
    return (coordinates + 1) * (0.5 * dim_length - 0.5)


def array_coordinates_to_grid_sample_coordinates(
        array_coordinates: torch.Tensor, array_shape: Sequence[int]
) -> torch.Tensor:
    """Generate coordinates for use with torch.nn.functional.grid_sample from array coordinates.

    Notes
    -----
    - array coordinates are from [0, N-1] for N elements in each dimension.
        - 0 is at the center of the first element
        - N is the length of the dimension
    - grid sample coordinates are from [-1, 1]
        - if align_corners=True, -1 and 1 are at the edges of array elements 0 and N-1
        - if align_corners=False, -1 and 1 are at the centers of array elements 0 and N-1
    - generated coordinates are
    """
    coords = [
        _array_coordinates_to_grid_sample_coordinates_1d(array_coordinates[..., idx], dim_length)
        for idx, dim_length
        in enumerate(array_shape)
    ]
    return einops.rearrange(coords[::-1], 'xyz ... -> ... xyz')


def grid_sample_coordinates_to_array_coordinates(coordinates: torch.Tensor,
                                                 array_shape: Sequence[int]) -> torch.Tensor:
    indices = [
        _grid_sample_coordinates_to_array_coordinates_1d(coordinates[..., idx], dim_length)
        for idx, dim_length
        in enumerate(array_shape[::-1])
    ]
    return einops.rearrange(indices[::-1], 'zyx b h w -> b h w zyx')


def simulate_single_axis_tilt_series(start_angle: float, end_angle: float,
                                     num_images: int) -> ImageData:
    s0 = S(-volume_center)
    r1 = Ry(torch.linspace(start_angle, end_angle, steps=num_images))
    s2 = S(tilt_image_center)
    projection_matrices = s2 @ r1 @ s0
    rotation_matrices = einops.rearrange(projection_matrices[:, :3, :3], 'b i j -> b j i')
    slice_coordinates = generate_rotated_slice_coordinates(rotation_matrices, n=volume_shape[0])
    dft = torch.fft.fftshift(volume, dim=(0, 1, 2))
    dft = torch.fft.fftn(dft, dim=(0, 1, 2))
    dft = torch.fft.fftshift(dft, dim=(0, 1, 2))
    slices = slice_dft(dft, slice_coordinates)
    image_shape = slices.shape[-2:]
    image_center = torch.tensor(image_shape) // 2
    r_max = volume_shape[0] // 2
    ramp_filter = torch.linalg.norm(
        get_grid_coordinates(image_shape) - image_center, dim=-1
    ) / r_max
    slices *= ramp_filter
    projections = torch.fft.ifftshift(slices, dim=(1, 2))
    projections = torch.fft.ifftn(projections, dim=(1, 2))
    projections = torch.fft.ifftshift(projections, dim=(1, 2))
    projections = torch.real(projections)
    return np.array(projections)


@magicgui(
    auto_call=True,
    max_angle={'widget_type': 'Slider', 'min': 0, 'max': 90},
    num_images={'widget_type': 'Slider', 'min': 1, 'max': 100}
)
def simulate_tomogram(max_angle: float, num_images: int) -> ImageData:
    s0 = S(-volume_center)
    r1 = Ry(torch.linspace(-max_angle, max_angle, steps=num_images))
    s2 = S(tilt_image_center)
    projection_matrices = s2 @ r1 @ s0
    tilt_series = simulate_single_axis_tilt_series(-max_angle, max_angle, num_images)
    reconstruction = backproject(
        image_stack=torch.tensor(tilt_series),
        projection_matrices=projection_matrices,
        output_dimensions=volume_shape,
    )
    reconstruction -= torch.mean(reconstruction)
    reconstruction /= torch.std(reconstruction)
    return np.array(reconstruction)


viewer = napari.Viewer(ndisplay=3)
volume_layer = viewer.add_image(np.array(volume), name='original 3D volume')
viewer.window.add_dock_widget(simulate_tomogram, name='WBP simulator')
napari.run()