change transform_matrix name

Author: jdickerson95

src/leopard_em/backend/core_differentiable_refine.py

@@ -68,7 +68,7 @@ def core_differentiable_refine(
     device: torch.device | list[torch.device],
     batch_size: int = 32,
     num_cuda_streams: int = 1,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> dict[torch.Tensor, torch.Tensor]:
     """Core function to refine orientations and defoci of a set of particles.
 
@@ -113,8 +113,8 @@ def core_differentiable_refine(
         The number of cross-correlations to process in one batch, defaults to 32.
     num_cuda_streams : int, optional
         Number of CUDA streams to use for parallel processing. Defaults to 1.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
 
     Returns
@@ -147,7 +147,7 @@ def core_differentiable_refine(
         batch_size=batch_size,
         devices=device,
         num_cuda_streams=num_cuda_streams,
-        transform_matrix=transform_matrix,
+        mag_matrix=mag_matrix,
     )
 
     results = {}
@@ -235,7 +235,7 @@ def _core_refine_template_single_gpu(
     batch_size: int,
     device: torch.device,
     num_cuda_streams: int = 1,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> None:
     """Run refine template on a subset of particles on a single GPU.
 
@@ -279,8 +279,8 @@ def _core_refine_template_single_gpu(
         Torch device to run this process on.
     num_cuda_streams : int, optional
         Number of CUDA streams to use for parallel processing. Defaults to 1.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
     """
     streams = [torch.cuda.Stream(device=device) for _ in range(num_cuda_streams)]
@@ -302,8 +302,8 @@ def _core_refine_template_single_gpu(
     corr_mean = corr_mean.to(device)
     corr_std = corr_std.to(device)
     projective_filters = projective_filters.to(device)
-    if transform_matrix is not None:
-        transform_matrix = transform_matrix.to(device)
+    if mag_matrix is not None:
+        mag_matrix = mag_matrix.to(device)
 
     ########################################
     ### Setup constants and progress bar ###
@@ -356,7 +356,7 @@ def _core_refine_template_single_gpu(
                 projective_filter=projective_filters[i],
                 batch_size=batch_size,
                 device_id=device_id,
-                transform_matrix=transform_matrix,
+                mag_matrix=mag_matrix,
             )
             refined_statistics.append(refined_stats)
 
@@ -469,7 +469,7 @@ def _core_refine_template_single_thread(
     projective_filter: torch.Tensor,
     batch_size: int = 32,
     device_id: int = 0,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> dict[str, float | int]:
     """Run the single-threaded core refine template function.
 
@@ -510,8 +510,8 @@ def _core_refine_template_single_thread(
         The number of orientations to cross-correlate at once. Default is 32.
     device_id : int, optional
         The ID of the device/process. Default is 0.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
 
     Returns
@@ -616,15 +616,15 @@ def _core_refine_template_single_thread(
                 rotation_matrices=rot_matrix_batch,
                 projective_filters=combined_projective_filter,
                 requires_grad=True,
-                transform_matrix=transform_matrix,
+                mag_matrix=mag_matrix,
             )
         else:
             cross_correlation = do_batched_orientation_cross_correlate_cpu(
                 image_dft=particle_image_dft,
                 template_dft=template_dft,
                 rotation_matrices=rot_matrix_batch,
                 projective_filters=combined_projective_filter,
-                transform_matrix=transform_matrix,
+                mag_matrix=mag_matrix,
             )
 
         cross_correlation = cross_correlation[..., :crop_h, :crop_w]  # valid crop

src/leopard_em/backend/core_match_template.py

@@ -155,7 +155,7 @@ def core_match_template(
     orientation_batch_size: int = 1,
     num_cuda_streams: int = 1,
     backend: str = "streamed",
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> dict[str, torch.Tensor]:
     """Core function for performing the whole-orientation search.
 
@@ -210,8 +210,8 @@ def core_match_template(
     backend : str, optional
         The backend to use for computation. Defaults to 'streamed'.
         Must be 'streamed' or 'batched'.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
 
     Returns
@@ -263,9 +263,9 @@ def core_match_template(
     defocus_values = defocus_values.cpu()
     pixel_values = pixel_values.cpu()
     euler_angles = euler_angles.cpu()
-    # Move transform_matrix to CPU if it's not None
-    if transform_matrix is not None:
-        transform_matrix = transform_matrix.cpu()
+    # Move mag_matrix to CPU if it's not None
+    if mag_matrix is not None:
+        mag_matrix = mag_matrix.cpu()
 
     ##############################################################
     ### Pre-multiply the whitening filter with the CTF filters ###
@@ -315,7 +315,7 @@ def core_match_template(
             "num_cuda_streams": num_cuda_streams,
             "backend": backend,
             "device": d,
-            "transform_matrix": transform_matrix,
+            "mag_matrix": mag_matrix,
         }
 
         kwargs_per_device.append(kwargs)
@@ -379,7 +379,7 @@ def _core_match_template_single_gpu(
     num_cuda_streams: int,
     backend: str,
     device: torch.device,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Single-GPU call for template matching.
 
@@ -419,8 +419,8 @@ def _core_match_template_single_gpu(
         Defaults to 'streamed'. Must be 'streamed' or 'batched'.
     device : torch.device
         Device to run the computation on. All tensors must be allocated on this device.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
 
     Returns
@@ -447,9 +447,9 @@ def _core_match_template_single_gpu(
     template_dft = template_dft.to(device)
     euler_angles = euler_angles.to(device)
     projective_filters = projective_filters.to(device)
-    # Move transform_matrix to device if it's not None
-    if transform_matrix is not None:
-        transform_matrix = transform_matrix.to(device)
+    # Move mag_matrix to device if it's not None
+    if mag_matrix is not None:
+        mag_matrix = mag_matrix.to(device)
 
     num_orientations = euler_angles.shape[0]
     num_defocus = defocus_values.shape[0]
@@ -525,7 +525,7 @@ def _core_match_template_single_gpu(
                         template_dft=template_dft,
                         rotation_matrices=rot_matrix,
                         projective_filters=projective_filters,
-                        transform_matrix=transform_matrix,
+                        mag_matrix=mag_matrix,
                     )
                 else:
                     cross_correlation = do_streamed_orientation_cross_correlate(
@@ -534,7 +534,7 @@ def _core_match_template_single_gpu(
                         rotation_matrices=rot_matrix,
                         projective_filters=projective_filters,
                         streams=streams,
-                        transform_matrix=transform_matrix,
+                        mag_matrix=mag_matrix,
                     )
 
                 # Update the tracked statistics

src/leopard_em/backend/core_refine_template.py

@@ -61,7 +61,7 @@ def core_refine_template(
     device: torch.device | list[torch.device],
     batch_size: int = 32,
     num_cuda_streams: int = 1,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> dict[str, torch.Tensor]:
     """Core function to refine orientations and defoci of a set of particles.
 
@@ -106,8 +106,8 @@ def core_refine_template(
         The number of cross-correlations to process in one batch, defaults to 32.
     num_cuda_streams : int, optional
         Number of CUDA streams to use for parallel processing. Defaults to 1.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
 
     Returns
@@ -140,7 +140,7 @@ def core_refine_template(
         batch_size=batch_size,
         devices=device,
         num_cuda_streams=num_cuda_streams,
-        transform_matrix=transform_matrix,
+        mag_matrix=mag_matrix,
     )
 
     results = run_multiprocess_jobs(
@@ -234,7 +234,7 @@ def construct_multi_gpu_refine_template_kwargs(
     batch_size: int,
     devices: list[torch.device],
     num_cuda_streams: int,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> list[dict]:
     """Split particle stack between requested devices.
 
@@ -272,8 +272,8 @@ def construct_multi_gpu_refine_template_kwargs(
         List of devices to split across.
     num_cuda_streams : int
         Number of CUDA streams to use per device.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
 
     Returns
@@ -330,7 +330,7 @@ def construct_multi_gpu_refine_template_kwargs(
             "batch_size": batch_size,
             "num_cuda_streams": num_cuda_streams,
             "device": device,
-            "transform_matrix": transform_matrix,
+            "mag_matrix": mag_matrix,
         }
 
         kwargs_per_device.append(kwargs)
@@ -360,7 +360,7 @@ def _core_refine_template_single_gpu(
     batch_size: int,
     device: torch.device,
     num_cuda_streams: int = 1,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> None:
     """Run refine template on a subset of particles on a single GPU.
 
@@ -404,8 +404,8 @@ def _core_refine_template_single_gpu(
         Torch device to run this process on.
     num_cuda_streams : int, optional
         Number of CUDA streams to use for parallel processing. Defaults to 1.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
     """
     streams = [torch.cuda.Stream(device=device) for _ in range(num_cuda_streams)]
@@ -478,7 +478,7 @@ def _core_refine_template_single_gpu(
                 corr_std=corr_std[i],
                 projective_filter=projective_filters[i],
                 batch_size=batch_size,
-                transform_matrix=transform_matrix,
+                mag_matrix=mag_matrix,
                 device_id=device_id,
             )
             refined_statistics.append(refined_stats)
@@ -580,7 +580,7 @@ def _core_refine_template_single_thread(
     projective_filter: torch.Tensor,
     batch_size: int = 32,
     device_id: int = 0,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> dict[str, float | int]:
     """Run the single-threaded core refine template function.
 
@@ -621,8 +621,8 @@ def _core_refine_template_single_thread(
         The number of orientations to cross-correlate at once. Default is 32.
     device_id : int, optional
         The ID of the device/process. Default is 0.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
 
     Returns
@@ -712,15 +712,15 @@ def _core_refine_template_single_thread(
                 template_dft=template_dft,
                 rotation_matrices=rot_matrix_batch,
                 projective_filters=combined_projective_filter,
-                transform_matrix=transform_matrix,
+                mag_matrix=mag_matrix,
             )
         else:
             cross_correlation = do_batched_orientation_cross_correlate_cpu(
                 image_dft=particle_image_dft,
                 template_dft=template_dft,
                 rotation_matrices=rot_matrix_batch,
                 projective_filters=combined_projective_filter,
-                transform_matrix=transform_matrix,
+                mag_matrix=mag_matrix,
             )
 
         cross_correlation = cross_correlation[..., :crop_h, :crop_w]  # valid crop
@@ -786,7 +786,7 @@ def cross_correlate_particle_stack(
     projective_filters: torch.Tensor,  # (N, h, w)
     mode: Literal["valid", "same"] = "valid",
     batch_size: int = 1024,
-    transform_matrix: torch.Tensor | None = None,
+    mag_matrix: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Cross-correlate a stack of particle images against a template.
 
@@ -815,8 +815,8 @@ def cross_correlate_particle_stack(
         The number of particle images to cross-correlate at once. Default is 1024.
         Larger sizes will consume more memory. If -1, then the entire stack will be
         cross-correlated at once.
-    transform_matrix : torch.Tensor | None, optional
-        Anisotropic magnification transform matrix of shape (2, 2). If None,
+    mag_matrix : torch.Tensor | None, optional
+        Anisotropic magnification matrix of shape (2, 2). If None,
         no magnification transform is applied. Default is None.
 
     Returns
@@ -868,14 +868,14 @@ def cross_correlate_particle_stack(
         )
         # Apply anisotropic magnification transform if provided
         # pylint: disable=duplicate-code
-        if transform_matrix is not None:
+        if mag_matrix is not None:
             rfft_shape = (template_h, template_w)
             stack_shape = (batch_rotation_matrices.shape[0],)
             fourier_slice = transform_slice_2d(
                 projection_image_dfts=fourier_slice,
                 rfft_shape=rfft_shape,
                 stack_shape=stack_shape,
-                transform_matrix=transform_matrix,
+                transform_matrix=mag_matrix,
             )
         fourier_slice = torch.fft.ifftshift(fourier_slice, dim=(-2,))
         fourier_slice[..., 0, 0] = 0 + 0j  # zero out the DC component (mean zero)