Rsg devel

src/libtilt/ctf/ctf_2d.py

@@ -22,6 +22,7 @@ def calculate_ctf(
         image_shape: Tuple[int, int],
         rfft: bool,
         fftshift: bool,
+        device
 ):
     """
 
@@ -56,20 +57,21 @@ def calculate_ctf(
         Whether to apply fftshift on the resulting CTF images.
     """
     # to torch.Tensor and unit conversions
-    defocus = torch.atleast_1d(torch.as_tensor(defocus, dtype=torch.float))
+    assert bool(rfft) + bool(fftshift) <= 1, "Error, only one of `rfft` and `fftshift` may be `True`."
+    defocus = torch.atleast_1d(torch.as_tensor(defocus, dtype=torch.float, device=device))
     defocus *= 1e4  # micrometers -> angstroms
-    astigmatism = torch.atleast_1d(torch.as_tensor(astigmatism, dtype=torch.float))
+    astigmatism = torch.atleast_1d(torch.as_tensor(astigmatism, dtype=torch.float, device=device))
     astigmatism *= 1e4  # micrometers -> angstroms
-    astigmatism_angle = torch.atleast_1d(torch.as_tensor(astigmatism_angle, dtype=torch.float))
+    astigmatism_angle = torch.atleast_1d(torch.as_tensor(astigmatism_angle, dtype=torch.float, device=device))
     astigmatism_angle *= (C.pi / 180)  # degrees -> radians
-    pixel_size = torch.atleast_1d(torch.as_tensor(pixel_size))
-    voltage = torch.atleast_1d(torch.as_tensor(voltage, dtype=torch.float))
+    pixel_size = torch.atleast_1d(torch.as_tensor(pixel_size, device=device))
+    voltage = torch.atleast_1d(torch.as_tensor(voltage, dtype=torch.float, device=device))
     voltage *= 1e3  # kV -> V
     spherical_aberration = torch.atleast_1d(
-        torch.as_tensor(spherical_aberration, dtype=torch.float)
+        torch.as_tensor(spherical_aberration, dtype=torch.float, device=device)
     )
     spherical_aberration *= 1e7  # mm -> angstroms
-    image_shape = torch.as_tensor(image_shape)
+    image_shape = torch.as_tensor(image_shape, device=device)
 
     # derived quantities used in CTF calculation
     defocus_u = defocus + astigmatism
@@ -79,10 +81,10 @@ def calculate_ctf(
     k2 = C.pi / 2 * spherical_aberration * _lambda ** 3
     k3 = torch.tensor(np.deg2rad(phase_shift))
     k4 = -b_factor / 4
-    k5 = np.arctan(amplitude_contrast / np.sqrt(1 - amplitude_contrast ** 2))
+    k5 = torch.arctan(amplitude_contrast / torch.sqrt(1 - amplitude_contrast ** 2))
 
     # construct 2D frequency grids and rescale cycles / px -> cycles / Å
-    fftfreq_grid = _construct_fftfreq_grid_2d(image_shape=image_shape, rfft=rfft)  # (h, w, 2)
+    fftfreq_grid = _construct_fftfreq_grid_2d(image_shape=image_shape, rfft=rfft, device=device)  # (h, w, 2)
     fftfreq_grid = fftfreq_grid / einops.rearrange(pixel_size, 'b -> b 1 1 1')
     fftfreq_grid_squared = fftfreq_grid ** 2
 

src/libtilt/fft_utils.py

@@ -22,10 +22,11 @@ def dft_center(
     device: torch.device | None = None,
 ) -> torch.LongTensor:
     """Return the position of the DFT center for a given input shape."""
+    _rfft_shape = rfft_shape(image_shape)
     fft_center = torch.zeros(size=(len(image_shape),), device=device)
     image_shape = torch.as_tensor(image_shape).float()
     if rfft is True:
-        image_shape = torch.tensor(rfft_shape(image_shape))
+        image_shape = torch.tensor(_rfft_shape, device=device)
     if fftshifted is True:
         fft_center = torch.divide(image_shape, 2, rounding_mode='floor')
     if rfft is True:
@@ -438,17 +439,19 @@ def fftfreq_to_dft_coordinates(
     coordinates: torch.Tensor
         `(..., d)` array of coordinates into a fftshifted DFT.
     """
+    _image_shape = image_shape
     image_shape = torch.as_tensor(
-        image_shape, device=frequencies.device, dtype=frequencies.dtype
+        _image_shape, device=frequencies.device, dtype=frequencies.dtype
     )
+    _rfft_shape = rfft_shape(_image_shape)
     _rfft_shape = torch.as_tensor(
-        rfft_shape(image_shape), device=frequencies.device, dtype=frequencies.dtype
+        _rfft_shape, device=frequencies.device, dtype=frequencies.dtype
     )
     coordinates = torch.empty_like(frequencies)
     coordinates[..., :-1] = frequencies[..., :-1] * image_shape[:-1]
     if rfft is True:
         coordinates[..., -1] = frequencies[..., -1] * 2 * (_rfft_shape[-1] - 1)
     else:
         coordinates[..., -1] = frequencies[..., -1] * image_shape[-1]
-    dc = dft_center(image_shape, rfft=rfft, fftshifted=True, device=frequencies.device)
+    dc = dft_center(_image_shape, rfft=rfft, fftshifted=True, device=frequencies.device)
     return coordinates + dc

src/libtilt/grids/central_slice_grid.py

@@ -43,7 +43,7 @@ def rotated_central_slice_grid(
         device=device,
     )  # (h, w, 3)
     if rotation_matrix_zyx is False:
-        grid = torch.flip(grid, dims=(-1,))
+        grid = torch.flip(grid, dims=(-1,)) #TODO: This operation is slow since it is copying the full tensor
     rotation_matrices = einops.rearrange(rotation_matrices, '... i j -> ... 1 1 i j')
     grid = einops.rearrange(grid, 'h w coords -> h w coords 1')
     grid = rotation_matrices @ grid

src/libtilt/interpolation/interpolate_dft_3d.py

@@ -49,7 +49,8 @@ def sample_dft_3d(
     samples = torch.view_as_complex(samples.contiguous())  # (b, )
 
     # pack data back up and return
-    [samples] = einops.unpack(samples, pattern='*', packed_shapes=ps)
+    # [samples] = einops.unpack(samples, pattern='*', packed_shapes=ps)
+    samples = samples.reshape(*ps) #Ask Alister if this will work in any situation
     return samples  # (...)
 
 

src/libtilt/projection/project_fourier.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 import torch
 import torch.nn.functional as F
 import einops
@@ -33,30 +35,12 @@ def project_fourier(
     projections: torch.Tensor
         `(..., d, d)` array of projection images.
     """
-    # padding
-    if pad is True:
-        pad_length = volume.shape[-1] // 2
-        volume = F.pad(volume, pad=[pad_length] * 6, mode='constant', value=0)
-
-    # premultiply by sinc2
-    grid = fftfreq_grid(
-        image_shape=volume.shape,
-        rfft=False,
-        fftshift=True,
-        norm=True,
-        device=volume.device
-    )
-    volume = volume * torch.sinc(grid) ** 2
-
-    # calculate DFT
-    dft = torch.fft.fftshift(volume, dim=(-3, -2, -1))  # volume center to array origin
-    dft = torch.fft.rfftn(dft, dim=(-3, -2, -1))
-    dft = torch.fft.fftshift(dft, dim=(-3, -2,))  # actual fftshift of rfft
+    dft, vol_shape, pad_length = compute_vol_dtf(volume, pad)
 
     # make projections by taking central slices
     projections = extract_central_slices_rfft(
         dft=dft,
-        image_shape=volume.shape,
+        image_shape=vol_shape,
         rotation_matrices=rotation_matrices,
         rotation_matrix_zyx=rotation_matrix_zyx
     )  # (..., h, w) rfft
@@ -92,7 +76,8 @@ def extract_central_slices_rfft(
 
     # flip coordinates in redundant half transform
     conjugate_mask = grid[..., 2] < 0
-    conjugate_mask = einops.repeat(conjugate_mask, '... -> ... 3')
+    # conjugate_mask = einops.repeat(conjugate_mask, '... -> ... 3')
+    conjugate_mask.unsqueeze(-1).repeat(1, 1, 1, 3)
     grid[conjugate_mask] *= -1
     conjugate_mask = conjugate_mask[..., 0]  # un-repeat
 
@@ -107,3 +92,49 @@ def extract_central_slices_rfft(
     # take complex conjugate of values from redundant half transform
     projections[conjugate_mask] = torch.conj(projections[conjugate_mask])
     return projections
+
+def compute_vol_dtf( #TODO: Is this the best place to have this?
+    volume: torch.Tensor,
+    pad: bool = True,
+    pad_length: int | None = None
+) -> Tuple[torch.Tensor, Tuple[int,int,int], int]:
+    """Project a cubic volume by sampling a central slice through its DFT.
+
+    Parameters
+    ----------
+    volume: torch.Tensor
+        `(d, d, d)` volume.
+    pad: bool
+        Whether to pad the volume with zeros to increase sampling in the DFT.
+    pad_length: bool
+        The lenght used for padding. If None, volume.shape[-1] // 2 is used instead
+
+    Returns
+    -------
+    projections: Tuple[torch.Tensor, torch.Tensor, int]
+        `(..., d, d, d)` dft of the volume. fftshifted rfft
+        Tuple[int,int,int] the shape of the volume after padding
+        int with the padding length
+    """
+    # padding
+    if pad is True:
+        if pad_length is None:
+            pad_length = volume.shape[-1] // 2
+        volume = F.pad(volume, pad=[pad_length] * 6, mode='constant', value=0)
+
+    # premultiply by sinc2
+    grid = fftfreq_grid(
+        image_shape=volume.shape,
+        rfft=False,
+        fftshift=True,
+        norm=True,
+        device=volume.device
+    )
+    volume = volume * torch.sinc(grid) ** 2
+
+    # calculate DFT
+    dft = torch.fft.fftshift(volume, dim=(-3, -2, -1))  # volume center to array origin
+    dft = torch.fft.rfftn(dft, dim=(-3, -2, -1))
+    dft = torch.fft.fftshift(dft, dim=(-3, -2,))  # actual fftshift of rfft
+
+    return dft, volume.shape, pad_length

src/libtilt/projection/project_real.py

@@ -49,7 +49,7 @@ def project_real(volume: torch.Tensor, rotation_matrices: torch.Tensor) -> torch
     torch_padding = einops.rearrange(torch_padding, 'whd pad -> (whd pad)')
     volume = F.pad(volume, pad=tuple(torch_padding), mode='constant', value=0)
     padded_volume_shape = (ps, ps, ps)
-    volume_coordinates = coordinate_grid(image_shape=padded_volume_shape)
+    volume_coordinates = coordinate_grid(image_shape=padded_volume_shape, device=volume.device)
     volume_coordinates -= padded_sidelength // 2  # (d, h, w, zyx)
     volume_coordinates = torch.flip(volume_coordinates, dims=(-1,))  # (d, h, w, zyx)
     volume_coordinates = einops.rearrange(volume_coordinates, 'd h w zyx -> d h w zyx 1')
@@ -73,5 +73,5 @@ def _project_volume(rotation_matrix) -> torch.Tensor:
 
     yl, yh = padding[1, 0], -padding[1, 1]
     xl, xh = padding[2, 0], -padding[2, 1]
-    images = [_project_volume(matrix)[yl:yh, xl:xh] for matrix in rotation_matrices]
+    images = [_project_volume(matrix)[yl:yh, xl:xh] for matrix in rotation_matrices] #TODO: This can probabaly optimized using vmap
     return torch.stack(images, dim=0)

src/libtilt/shapes/soft_edge.py

@@ -8,7 +8,7 @@ def _add_soft_edge_single_binary_image(
 ) -> torch.FloatTensor:
     if smoothing_radius == 0:
         return image.float()
-    distances = ndi.distance_transform_edt(torch.logical_not(image))
+    distances = ndi.distance_transform_edt(torch.logical_not(image)) #TODO: This breaks if the input device is cuda
     distances = torch.as_tensor(distances, device=image.device).float()
     idx = torch.logical_and(distances > 0, distances <= smoothing_radius)
     output = torch.clone(image).float()