tomodrgn.lattice.Lattice#

class Lattice(boxsize: int, extent: float = 0.5, ignore_dc: bool = True, device: device | None = None)[source]#

Bases: object

Methods

get_circular_mask

Return a binary mask for self.coords which restricts coordinates to a centered circular lattice

get_downsample_coords

Return a 2-D lattice of coordinates representing a downsampled (fourier-cropped) copy of the original lattice.

get_square_mask

Return a binary mask for self.coords which restricts coordinates to a centered square lattice

rotate

Resample a stack of images on the lattice grid rotated in-plane counterclockwise by a batch of theta angles.

translate_ft

Translate an image by phase shifting its Fourier transform.

translate_ht

Translate an image by phase shifting its Hartley transform H'(k) = cos(2*pi*k*t0)H(k) + sin(2*pi*k*t0)H(-k) img must be 1D unraveled image, symmetric around DC component
get_circular_mask(diameter: int) → Tensor[source]#

Return a binary mask for self.coords which restricts coordinates to a centered circular lattice

Parameters:

diameter – number of grid points to include in the mask along each dimension

Returns:

binary mask, shape (lattice.boxsize ** 2)

get_downsample_coords(boxsize_new: int) → Tensor[source]#

Return a 2-D lattice of coordinates representing a downsampled (fourier-cropped) copy of the original lattice.

Parameters:

boxsize_new – number of grid points along each dimension in the downsampled lattice. Should be odd.

Returns:

coordinates of the downsampled lattice, shape (boxsize_new ** 2, 3)

get_square_mask(sidelength: int) → Tensor[source]#

Return a binary mask for self.coords which restricts coordinates to a centered square lattice

Parameters:

sidelength – number of grid points to include in the mask along each dimension

Returns:

binary mask, shape (lattice.boxsize ** 2)

rotate(images: Tensor, theta: Tensor) → Tensor[source]#

Resample a stack of images on the lattice grid rotated in-plane counterclockwise by a batch of theta angles.

Parameters:

  • images – stack of images to rotate, shape (B,Y,X)

  • theta – batch of angles in radians, shape (Q)

Returns:

rotated images, shape (B,Q,Y,X)

translate_ft(images: Tensor, trans: Tensor, mask: ndarray | Tensor | None = None) → Tensor[source]#

Translate an image by phase shifting its Fourier transform. F’(k) = exp(-2*pi*k*x0)*F(k) Note: shape img_dims below can either be 2D or 1D (unraveled image)

Parameters:

  • images – fourier transform of image, shape (B, img_dims, 2)

  • trans – shift in pixels, shape (B, T, 2)

  • mask – optional mask for lattice coords (img_dims, 1)

Returns:

translated images, shape (B, T, img_dims, 2)

translate_ht(images: Tensor, trans: Tensor, mask: ndarray | Tensor | None = None) → Tensor[source]#

Translate an image by phase shifting its Hartley transform H’(k) = cos(2*pi*k*t0)H(k) + sin(2*pi*k*t0)H(-k) img must be 1D unraveled image, symmetric around DC component

Parameters:

  • images – hartley transform of image, shape (B, boxsize_ht)

  • trans – shift in pixels, shape (B, T, 2)

  • mask – mask for lattice coords, shape (boxsize_ht, 1)

Returns:

translated images, shape (B, T, boxsize_ht**2)