diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index d7a4eb6e..3068b912 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -102,7 +102,7 @@ jobs: - name: 🚢 Publish to PyPI uses: pypa/gh-action-pypi-publish@release/v1 - - uses: softprops/action-gh-release@v2 + - uses: softprops/action-gh-release@v3 with: generate_release_notes: true files: './dist/*' diff --git a/benchmark/benchmark_match_template.py b/benchmark/benchmark_match_template.py index 3d8b73cb..fe620c0d 100644 --- a/benchmark/benchmark_match_template.py +++ b/benchmark/benchmark_match_template.py @@ -76,6 +76,7 @@ def benchmark_match_template_single_run( **core_kwargs, orientation_batch_size=orientation_batch_size, num_cuda_streams=mt_manager.computational_config.num_cpus, + backend=mt_manager.computational_config.backend, ) total_projections = result["total_projections"] # number of CCGs calculated, N @@ -102,13 +103,14 @@ def benchmark_match_template_single_run( # --> r = (N - n) / (T_N - T_n) # --> k = N * (T_N - T_n) / (N - n) - core_kwargs["euler_angles"] = torch.rand(size=(100, 3)) * 180 + core_kwargs["euler_angles"] = torch.rand(size=(300, 3)) * orientation_batch_size start_time = time.perf_counter() result = core_match_template( **core_kwargs, orientation_batch_size=orientation_batch_size, num_cuda_streams=mt_manager.computational_config.num_cpus, + backend=mt_manager.computational_config.backend, ) adjustment_projections = result["total_projections"] # number of CCGs calculated, n @@ -137,7 +139,7 @@ def run_benchmark(orientation_batch_size: int, num_runs: int) -> dict[str, Any]: """Run multiple benchmark iterations and collect statistics.""" # Download example data to use for benchmarking print("Downloading benchmarking data...") - download_comparison_data() + # download_comparison_data() print("Done!") # Get CUDA device properties diff --git a/pyproject.toml b/pyproject.toml index e9e459ad..52e7e950 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,10 +54,13 @@ dependencies = [ "ttsim3d>=v0.4.0", "lmfit", "zenodo-get", + "h5py", + "tensordict", "torch-fourier-shift", "torch-motion-correction>=0.0.4", "torch-grid-utils>=v0.0.9", "torch-ctf", + "h5py", ] [tool.hatch.metadata] diff --git a/src/leopard_em/backend/core_match_template.py b/src/leopard_em/backend/core_match_template.py index 1aa358a4..f72097c4 100644 --- a/src/leopard_em/backend/core_match_template.py +++ b/src/leopard_em/backend/core_match_template.py @@ -4,17 +4,20 @@ # pylint: disable=E1102 import time +import traceback import warnings from functools import partial from multiprocessing import set_start_method from typing import Any, Union import roma +import tensordict import torch import tqdm from leopard_em.backend.cross_correlation import ( do_batched_orientation_cross_correlate, + do_batched_orientation_cross_correlate_zipfft, do_streamed_orientation_cross_correlate, ) from leopard_em.backend.distributed import ( @@ -24,17 +27,20 @@ from leopard_em.backend.process_results import ( aggregate_distributed_results, decode_global_search_index, + process_correlation_table, scale_mip, ) -from leopard_em.backend.utils import do_iteration_statistics_updates_compiled +from leopard_em.backend.utils import do_iteration_and_correlation_table_updates DEFAULT_STATISTIC_DTYPE = torch.float32 +CORRELATION_TABLE_THRESHOLD = 5.5 # Turn off gradient calculations by default torch.set_grad_enabled(False) # Set multiprocessing start method to spawn set_start_method("spawn", force=True) +torch.multiprocessing.set_sharing_strategy("file_system") def monitor_match_template_progress( @@ -77,6 +83,7 @@ def monitor_match_template_progress( time.sleep(poll_interval) except Exception as e: print(f"Error occurred: {e}") + traceback.print_exc() queue.set_error_flag() raise e finally: @@ -156,7 +163,7 @@ def core_match_template( num_cuda_streams: int = 1, backend: str = "streamed", mag_matrix: torch.Tensor | None = None, -) -> dict[str, torch.Tensor]: +) -> dict[str, torch.Tensor | dict | int]: """Core function for performing the whole-orientation search. With the RFFT, the last dimension (fastest dimension) is half the width @@ -216,7 +223,7 @@ def core_match_template( Returns ------- - dict[str, torch.Tensor] + dict[str, torch.Tensor | dict | int] Dictionary containing the following key, value pairs: - "mip": Maximum intensity projection of the cross-correlation values across @@ -226,10 +233,12 @@ def core_match_template( - "best_theta": Best theta angle for each pixel. - "best_psi": Best psi angle for each pixel. - "best_defocus": Best defocus value for each pixel. - - "best_pixel_size": Best pixel size value for each pixel. - - "correlation_sum": Sum of cross-correlation values for each pixel. - - "correlation_squared_sum": Sum of squared cross-correlation values for + - "correlation_mean": Sum of cross-correlation values for each pixel. + - "correlation_variance": Sum of squared cross-correlation values for + - "correlation_table": Processed correlation table with all points in search + space and image positions where correlation value exceeded a threshold. each pixel. + - "total_projections": Total number of cross-correlations computed. - "total_orientations": Total number of orientations searched. - "total_defocus": Total number of defocus values searched. """ @@ -335,7 +344,7 @@ def core_match_template( correlation_squared_sum = aggregated_results["correlation_squared_sum"] # Map from global search index to the best defocus & angles - best_phi, best_theta, best_psi, best_defocus = decode_global_search_index( + best_phi, best_theta, best_psi, best_defocus, _ = decode_global_search_index( best_global_index, pixel_values, defocus_values, euler_angles ) @@ -348,6 +357,14 @@ def core_match_template( total_correlation_positions=total_projections, ) + # Process the correlation table into a more interpretable format + correlation_table = process_correlation_table( + aggregated_results["correlation_table"], + pixel_values, + defocus_values, + euler_angles, + ) + return { "mip": mip, "scaled_mip": mip_scaled, @@ -357,6 +374,7 @@ def core_match_template( "best_defocus": best_defocus, "correlation_mean": correlation_mean, "correlation_variance": correlation_variance, + "correlation_table": correlation_table, "total_projections": total_projections, "total_orientations": euler_angles.shape[0], "total_defocus": defocus_values.shape[0], @@ -380,7 +398,9 @@ def _core_match_template_single_gpu( backend: str, device: torch.device, mag_matrix: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, tensordict.TensorDict +]: """Single-GPU call for template matching. Parameters @@ -433,8 +453,18 @@ def _core_match_template_single_gpu( - correlation_sum: Sum of cross-correlation values for each pixel. - correlation_squared_sum: Sum of squared cross-correlation values for each pixel. + - correlation_table: Table of search indices and image positions where + correlation values exceeded a threshold. """ image_shape_real = (image_dft.shape[0], image_dft.shape[1] * 2 - 2) # adj. for RFFT + projection_shape_real = ( + template_dft.shape[1], + template_dft.shape[2] * 2 - 2, # adj. for RFFT + ) + valid_correlation_shape = ( + image_shape_real[0] - projection_shape_real[0] + 1, + image_shape_real[1] - projection_shape_real[1] + 1, + ) # Create CUDA streams for parallel computation streams = [torch.cuda.Stream(device=device) for _ in range(num_cuda_streams)] @@ -469,21 +499,52 @@ def _core_match_template_single_gpu( ### Initialize the tracked output statistics ### ################################################ + # Correlation table built from 'tensordict' library where any (x, y) positions + # in correlation map which surpass the threshold will be added to the table. + # Keys in table are: + # - "threshold": float threshold value used for the table. + # - "global_idx": int32 global search index. + # - "pos_x": int32 x position in image where corr value surpassed threshold. + # - "pos_y": int32 y position in image where corr value surpassed threshold. + # - "corr_value": float32 correlation value at (pos_x, pos_y) for the given + # global index. + correlation_table = tensordict.TensorDict( + { + "threshold": CORRELATION_TABLE_THRESHOLD, + "global_idx": torch.tensor([], dtype=torch.int32, device=device), + "pos_x": torch.tensor([], dtype=torch.int32, device=device), + "pos_y": torch.tensor([], dtype=torch.int32, device=device), + "corr_value": torch.tensor([], dtype=torch.float32, device=device), + }, + device=device, + ) mip = torch.full( - size=image_shape_real, + size=valid_correlation_shape, fill_value=-float("inf"), dtype=DEFAULT_STATISTIC_DTYPE, device=device, ) best_global_index = torch.full( - image_shape_real, fill_value=-1, dtype=torch.int32, device=device + valid_correlation_shape, + fill_value=-1, + dtype=torch.int32, + device=device, ) correlation_sum = torch.zeros( - size=image_shape_real, dtype=DEFAULT_STATISTIC_DTYPE, device=device + size=valid_correlation_shape, + dtype=DEFAULT_STATISTIC_DTYPE, + device=device, ) correlation_squared_sum = torch.zeros( - size=image_shape_real, dtype=DEFAULT_STATISTIC_DTYPE, device=device + size=valid_correlation_shape, + dtype=DEFAULT_STATISTIC_DTYPE, + device=device, ) + if backend == "zipfft": + # NOTE: zipFFT expects a pre-transformed, pre-transposed input image FFT + # Transpose the 'image_dft' along last two dimensions into contiguous layout + # with shape (..., W // 2 + 1, H) + image_dft = image_dft.transpose(-2, -1).contiguous() ################################## ### Start the orientation loop ### @@ -527,7 +588,7 @@ def _core_match_template_single_gpu( projective_filters=projective_filters, mag_matrix=mag_matrix, ) - else: + elif backend == "streamed": cross_correlation = do_streamed_orientation_cross_correlate( image_dft=image_dft, template_dft=template_dft, @@ -536,18 +597,29 @@ def _core_match_template_single_gpu( streams=streams, mag_matrix=mag_matrix, ) + elif backend == "zipfft": + cross_correlation = do_batched_orientation_cross_correlate_zipfft( + image_dft=image_dft, + template_dft=template_dft, + rotation_matrices=rot_matrix, + projective_filters=projective_filters, + ) - # Update the tracked statistics - do_iteration_statistics_updates_compiled( + # Update tracked statistics and correlation table + do_iteration_and_correlation_table_updates( cross_correlation=cross_correlation, current_indexes=batch_search_indices, + correlation_table=correlation_table, mip=mip, best_global_index=best_global_index, correlation_sum=correlation_sum, correlation_squared_sum=correlation_squared_sum, - img_h=image_shape_real[0], - img_w=image_shape_real[1], + threshold=CORRELATION_TABLE_THRESHOLD, + valid_shape_h=valid_correlation_shape[0], + valid_shape_w=valid_correlation_shape[1], + needs_valid_cropping=(backend != "zipfft"), ) + except Exception as e: index_queue.set_error_flag() print(f"Error occurred in process {rank}: {e}") @@ -559,7 +631,13 @@ def _core_match_template_single_gpu( torch.cuda.synchronize(device) - return mip, best_global_index, correlation_sum, correlation_squared_sum + return ( + mip, + best_global_index, + correlation_sum, + correlation_squared_sum, + correlation_table, + ) def _core_match_template_multiprocess_wrapper( @@ -573,9 +651,13 @@ def _core_match_template_multiprocess_wrapper( See the _core_match_template_single_gpu function for parameter descriptions. """ - mip, best_global_index, correlation_sum, correlation_squared_sum = ( - _core_match_template_single_gpu(rank, **kwargs) # type: ignore[arg-type] - ) + ( + mip, + best_global_index, + correlation_sum, + correlation_squared_sum, + correlation_table, + ) = _core_match_template_single_gpu(rank, **kwargs) # type: ignore[arg-type] # NOTE: Need to send all tensors back to the CPU as numpy arrays for the shared # process dictionary. This is a workaround for now @@ -584,6 +666,7 @@ def _core_match_template_multiprocess_wrapper( "best_global_index": best_global_index.cpu().numpy(), "correlation_sum": correlation_sum.cpu().numpy(), "correlation_squared_sum": correlation_squared_sum.cpu().numpy(), + "correlation_table": correlation_table.cpu(), } # Place the results in the shared multi-process manager dictionary so accessible diff --git a/src/leopard_em/backend/core_match_template_distributed.py b/src/leopard_em/backend/core_match_template_distributed.py index 660cfc4d..b9c64c74 100644 --- a/src/leopard_em/backend/core_match_template_distributed.py +++ b/src/leopard_em/backend/core_match_template_distributed.py @@ -454,21 +454,25 @@ def core_match_template_distributed( ########################################################### dist.barrier() - (mip, best_global_index, correlation_sum, correlation_squared_sum) = ( - _core_match_template_single_gpu( - rank=rank, - index_queue=distributed_queue, # type: ignore - image_dft=image_dft, - template_dft=template_dft, - euler_angles=euler_angles, - projective_filters=projective_filters, - defocus_values=defocus_values, - pixel_values=pixel_values, - orientation_batch_size=orientation_batch_size, - num_cuda_streams=num_cuda_streams, - backend=backend, - device=device, - ) + ( + mip, + best_global_index, + correlation_sum, + correlation_squared_sum, + _, # TODO: include correlation_table in distributed version + ) = _core_match_template_single_gpu( + rank=rank, + index_queue=distributed_queue, # type: ignore + image_dft=image_dft, + template_dft=template_dft, + euler_angles=euler_angles, + projective_filters=projective_filters, + defocus_values=defocus_values, + pixel_values=pixel_values, + orientation_batch_size=orientation_batch_size, + num_cuda_streams=num_cuda_streams, + backend=backend, + device=device, ) dist.barrier() @@ -534,7 +538,7 @@ def core_match_template_distributed( # Map from global search index to the best defocus & angles # pylint: disable=duplicate-code - best_phi, best_theta, best_psi, best_defocus = decode_global_search_index( + best_phi, best_theta, best_psi, best_defocus, _ = decode_global_search_index( best_global_index, pixel_values, defocus_values, euler_angles ) diff --git a/src/leopard_em/backend/core_refine_template.py b/src/leopard_em/backend/core_refine_template.py index 38843c3d..c5cc8efd 100644 --- a/src/leopard_em/backend/core_refine_template.py +++ b/src/leopard_em/backend/core_refine_template.py @@ -297,6 +297,8 @@ def construct_multi_gpu_refine_template_kwargs( device_defocus_v = defocus_v[start_idx:end_idx] device_defocus_angle = defocus_angle[start_idx:end_idx] device_projective_filters = projective_filters[start_idx:end_idx] + device_corr_mean = corr_mean[start_idx:end_idx] + device_corr_std = corr_std[start_idx:end_idx] kwargs = { "particle_stack_dft": device_particle_stack_dft, @@ -309,8 +311,8 @@ def construct_multi_gpu_refine_template_kwargs( "defocus_angle": device_defocus_angle, "defocus_offsets": defocus_offsets, "pixel_size_offsets": pixel_size_offsets, - "corr_mean": corr_mean, - "corr_std": corr_std, + "corr_mean": device_corr_mean, + "corr_std": device_corr_std, "projective_filters": device_projective_filters, "ctf_kwargs": ctf_kwargs, "batch_size": batch_size, diff --git a/src/leopard_em/backend/cross_correlation.py b/src/leopard_em/backend/cross_correlation.py index c078f881..f3acd8d6 100644 --- a/src/leopard_em/backend/cross_correlation.py +++ b/src/leopard_em/backend/cross_correlation.py @@ -8,6 +8,24 @@ normalize_template_projection_compiled, ) +# --- Import handling for zipfft library (which may not be installed) ------------------ +try: + import zipfft + + # Determine which batch sizes are supported by zipFFT for powers of 2 + # pylint: disable=c-extension-no-member + ZIPFFT_SUPPORTED_CONFIGS = zipfft.padded_rconv2d.get_supported_conv_configs() + ZIPFFT_SUPPORTED_BATCH_SIZES = [ + x[-2] + for x in ZIPFFT_SUPPORTED_CONFIGS + if (x[0] == 512 and x[1] == 512 and x[2] == 4096 and x[3] == 4096) + ] + ZIPFFT_SUPPORTED_BATCH_SIZES.sort(reverse=True) # largest to smallest +except ImportError: + zipfft = None + ZIPFFT_SUPPORTED_BATCH_SIZES = [] + ZIPFFT_SUPPORTED_CONFIGS = [] + # pylint: disable=too-many-locals,E1102 def do_streamed_orientation_cross_correlate( @@ -367,3 +385,136 @@ def do_batched_orientation_cross_correlate_cpu( cross_correlation = torch.fft.irfftn(projections_dft, dim=(-2, -1)) return cross_correlation + + +# pylint: disable=E1102 +def do_batched_orientation_cross_correlate_zipfft( + image_dft: torch.Tensor, + template_dft: torch.Tensor, + rotation_matrices: torch.Tensor, + projective_filters: torch.Tensor, +) -> torch.Tensor: + """Batched projection and cross-correlation using zipfft backend. + + This function uses the zipfft library for accelerated 2D cross-correlation + compared to `do_batched_orientation_cross_correlate`. + + NOTE: that this function returns a cross-correlogram with "same" mode (i.e. the + same size as the input image). See numpy correlate docs for more information. + + Parameters + ---------- + image_dft : torch.Tensor + Real-fourier transform (RFFT) of the image with large image filters + already applied. Has shape (H, W // 2 + 1). + template_dft : torch.Tensor + Real-fourier transform (RFFT) of the template volume to take Fourier + slices from. Has shape (l, h, w // 2 + 1) where (l, h, w) is the original + real-space shape of the template volume. + rotation_matrices : torch.Tensor + Rotation matrices to apply to the template volume. Has shape + (num_orientations, 3, 3). + projective_filters : torch.Tensor + Multiplied 'ctf_filters' with 'whitening_filter_template'. Has shape + (num_Cs, num_defocus, h, w // 2 + 1). Is RFFT and not fftshifted. + + Returns + ------- + torch.Tensor + Cross-correlation of the image with the template volume for each + orientation and defocus value. Will have shape + (num_Cs, num_defocus, num_orientations, H, W). + """ + # Accounting for RFFT shape + projection_shape_real = (template_dft.shape[1], template_dft.shape[2] * 2 - 2) + image_shape_real = ( + image_dft.shape[0] * 2 - 2, + image_dft.shape[1], + ) # NOTE: transposed + + num_orientations = rotation_matrices.shape[0] + num_Cs = projective_filters.shape[0] # pylint: disable=invalid-name + num_defocus = projective_filters.shape[1] + + # Output shape for cross-correlation + output_shape = ( + image_shape_real[0] - projection_shape_real[0] + 1, + image_shape_real[1] - projection_shape_real[1] + 1, + ) + + cross_correlation = torch.empty( + size=(num_Cs, num_defocus, num_orientations, *output_shape), + dtype=image_dft.real.dtype, + device=image_dft.device, + ) + + # Extract central slice(s) from the template volume + fourier_slice = extract_central_slices_rfft_3d( + volume_rfft=template_dft, + image_shape=(projection_shape_real[0],) * 3, # NOTE: requires cubic template + rotation_matrices=rotation_matrices, + ) + fourier_slice = torch.fft.ifftshift(fourier_slice, dim=(-2,)) + fourier_slice[..., 0, 0] = 0 + 0j # zero out the DC component (mean zero) + fourier_slice *= -1 # flip contrast + + # Apply the projective filters on a new batch dimension + fourier_slice = fourier_slice[None, None, ...] * projective_filters[:, :, None, ...] + + # Inverse Fourier transform into real space and normalize + projections = torch.fft.irfftn(fourier_slice, dim=(-2, -1)) + projections = torch.fft.ifftshift(projections, dim=(-2, -1)) + projections = normalize_template_projection_compiled( + projections, + projection_shape_real, + image_shape_real, + ) + + # Create workspace for FFT operations + # Shape: (num_orientations, fft_size_y, fft_size_x // 2 + 1) + corr_workspace = torch.empty( + num_orientations, + image_shape_real[0], + image_shape_real[1] // 2 + 1, + dtype=torch.complex64, + device=image_dft.device, + ) + + for j in range(num_defocus): + for k in range(num_Cs): + # Use zipfft for cross-correlation + # projections[k, j, ...] has shape (num_orientations, H_proj, W_proj) + # image_dft has already been pre-transposed into contiguous layout + # with (W // 2 + 1, H) for memory efficiency + # cross_correlation[k, j, ...] has shape (num_orientations, H_out, W_out) + + # NOTE: zipFFT only supports certain batch sizes for optimal performance, + # iterate through ZIPFFT_SUPPORTED_BATCH_SIZES to find the largest supported + # batch size to decompose the projections into. Batch=1 will always be + # supported. + if num_orientations in ZIPFFT_SUPPORTED_BATCH_SIZES: + # pylint: disable=c-extension-no-member + zipfft.padded_rconv2d.corr( + projections[k, j, ...], + corr_workspace, + image_dft, + cross_correlation[k, j, ...], + image_shape_real[0], + image_shape_real[1], + ) + else: + for i in range(num_orientations): + # pylint: disable=c-extension-no-member + zipfft.padded_rconv2d.corr( + projections[k, j, i, ...], + corr_workspace[i, ...], + image_dft, + cross_correlation[k, j, i, ...], + image_shape_real[0], + image_shape_real[1], + ) + + # NOTE: zipFFT internally does not synchronize CUDA, so must do it manually + torch.cuda.synchronize() + + return cross_correlation diff --git a/src/leopard_em/backend/process_results.py b/src/leopard_em/backend/process_results.py index 5bf56647..6b48c462 100644 --- a/src/leopard_em/backend/process_results.py +++ b/src/leopard_em/backend/process_results.py @@ -1,6 +1,9 @@ """Functions related to result processing after backend functions.""" +from typing import Any + import numpy as np +import tensordict import torch @@ -56,20 +59,37 @@ def aggregate_distributed_results( correlation_sum = torch.from_numpy(correlation_sum) correlation_squared_sum = torch.from_numpy(correlation_squared_sum) + # Merge the correlation table dictionaries + # (no key collisions expected after popping "threshold") + full_correlation_table = {} + for result in results: + correlation_table = result["correlation_table"] + correlation_table = ( + correlation_table.cpu().to_dict() + if isinstance(correlation_table, tensordict.TensorDict) + else correlation_table + ) + threshold = correlation_table.pop("threshold") + full_correlation_table.update(correlation_table) + + full_correlation_table["threshold"] = threshold + return { "mip": mip_max, "best_global_index": best_index, "correlation_sum": correlation_sum, "correlation_squared_sum": correlation_squared_sum, + "correlation_table": full_correlation_table, } +# pylint: disable=too-many-locals def decode_global_search_index( global_indices: torch.Tensor, # integer tensor pixel_values: torch.Tensor, # (num_cs,) defocus_values: torch.Tensor, # (num_defocus,) euler_angles: torch.Tensor, # (num_orientations, 3) -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Decode flattened global indices back into (cs, defocus, orientation).""" _ = pixel_values # Unused, but possible to add in future @@ -81,7 +101,7 @@ def decode_global_search_index( stride_defocus = num_orientations # Calculate the indexes for each "best" array - # pixel_idx = global_indices // stride_cs + pixel_idx = global_indices // stride_cs rem = global_indices % stride_cs defocus_idx = rem // stride_defocus orientations_idx = rem % stride_defocus @@ -90,9 +110,101 @@ def decode_global_search_index( theta = euler_angles[orientations_idx, 1] psi = euler_angles[orientations_idx, 2] defocus = defocus_values[defocus_idx] - # pixels = pixel_values[pixel_idx] + pixels = pixel_values[pixel_idx] + + return phi, theta, psi, defocus, pixels + + +# pylint: disable=too-many-locals +def process_correlation_table( + correlation_table: dict[int | str, Any], + pixel_values: torch.Tensor, # (num_cs,) + defocus_values: torch.Tensor, # (num_defocus,) + euler_angles: torch.Tensor, # (num_orientations, 3) +) -> dict[str, list[float | int]]: + """Process the correlation table by applying a threshold. + + Parameters + ---------- + correlation_table : dict[int, torch.Tensor] + Dictionary containing the correlation table. Keys are global search indices, + values are tensors of shape (num_hits, 3) containing (x, y, cc) values. + pixel_values : torch.Tensor + Tensor containing the pixel values used in the search. Shape is (num_cs,). + defocus_values : torch.Tensor + Tensor containing the defocus values used in the search. Shape is + (num_defocus,). + euler_angles : torch.Tensor + Tensor containing the Euler angles used in the search. Shape is + (num_orientations, 3). + + Returns + ------- + dict[str, list[float | int]] + Processed correlation with keys for the unique point in search space and image + position for all cross-correlations which surpassed the threshold. + """ + threshold = correlation_table.pop("threshold") + threshold = threshold.item() if isinstance(threshold, torch.Tensor) else threshold + # processed_table = { + # "threshold": threshold, + # "pixel_size": [], + # "defocus": [], + # "phi": [], + # "theta": [], + # "psi": [], + # "x": [], + # "y": [], + # "correlation": [], + # } + + # Convert string keys to integer tensor for decoding + global_indices = correlation_table["global_idx"] + phi, theta, psi, defocus, pixel_values = decode_global_search_index( + global_indices, pixel_values, defocus_values, euler_angles + ) + + processed_table = { + "threshold": threshold, + "global_idx": global_indices.numpy().tolist(), + "pixel_size": pixel_values.numpy().tolist(), + "defocus": defocus.numpy().tolist(), + "phi": phi.numpy().tolist(), + "theta": theta.numpy().tolist(), + "psi": psi.numpy().tolist(), + "x": correlation_table["pos_x"].numpy().tolist(), + "y": correlation_table["pos_y"].numpy().tolist(), + "correlation": correlation_table["corr_value"].numpy().tolist(), + } + + # # Process each global index + # for i, value in enumerate(correlation_table.values()): + # # Get parameters for this index + # this_pixel_size = pixel_values[i].item() + # this_defocus = defocus[i].item() + # this_phi = phi[i].item() + # this_theta = theta[i].item() # No tuple, just the value + # this_psi = psi[i].item() # No tuple, just the value + + # # Count points in this value + # num_points = value.shape[0] + + # # Extract coordinates and correlation values + # xs = value[:, 0].tolist() + # ys = value[:, 1].tolist() + # ccs = value[:, 2].tolist() + + # # Append all values at once + # processed_table["pixel_size"].extend([this_pixel_size] * num_points) + # processed_table["defocus"].extend([this_defocus] * num_points) + # processed_table["phi"].extend([this_phi] * num_points) + # processed_table["theta"].extend([this_theta] * num_points) + # processed_table["psi"].extend([this_psi] * num_points) + # processed_table["x"].extend([int(x) for x in xs]) + # processed_table["y"].extend([int(y) for y in ys]) + # processed_table["correlation"].extend(ccs) - return phi, theta, psi, defocus + return processed_table def correlation_sum_and_squared_sum_to_mean_and_variance( diff --git a/src/leopard_em/backend/utils.py b/src/leopard_em/backend/utils.py index 006fe0c4..93d2af8b 100644 --- a/src/leopard_em/backend/utils.py +++ b/src/leopard_em/backend/utils.py @@ -5,6 +5,7 @@ import warnings from typing import Any, Callable, TypeVar +import tensordict import roma import torch @@ -188,25 +189,125 @@ def normalize_template_projection( return projections -def do_iteration_statistics_updates( +@torch.compile # type: ignore[misc] +# pylint: disable=too-many-locals +def _stats_and_table_core( cross_correlation: torch.Tensor, current_indexes: torch.Tensor, mip: torch.Tensor, best_global_index: torch.Tensor, + threshold: float, + valid_shape_h: int, + valid_shape_w: int, + needs_valid_cropping: bool = True, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Compiled function to find new maxima and do correlation table updates. + + Parameters + ---------- + cross_correlation : torch.Tensor + Cross-correlation values for the current iteration. Has shape + (num_cs, num_defocus, num_orientations, H, W) where 'num_cs' are the number of + different pixel sizes (controlled by spherical aberration Cs) in the + cross-correlation batch, 'num_defocus' are the number of different defocus + values in the cross-correlation batch, and 'num_orientations' are the number of + different orientations in the cross-correlation batch. H and W can either be the + full image heigh/width or the valid cropped height/width. + current_indexes : torch.Tensor + The global search indexes for the *current* batch of pixel sizes, defocus + values, and orientations. Has shape `num_cs * num_defocus * num_orientations` + to uniquely identify the set of pixel sizes, defocus values, and orientations + associated with the batch from the global search space. + mip : torch.Tensor + Maximum intensity projection of the cross-correlation values. + best_global_index : torch.Tensor + Previous best global search indexes. Has shape (H, W) and is int32 type. + threshold : float + The threshold value for adding entries to the correlation table. + valid_shape_h : int + Height of the valid region of the cross-correlation values. + valid_shape_w : int + Width of the valid region of the cross-correlation values. + needs_valid_cropping : bool, optional + Whether the cross-correlation tensor should be cropped (via a view operation). + If False, the cross-correlation tensor is assumed to already be in the valid + shape. + """ + # create cropped view as in existing functions + if needs_valid_cropping: + cc_reshaped = cross_correlation.view( + -1, cross_correlation.shape[-2], cross_correlation.shape[-1] + ) + cc_reshaped = cc_reshaped.as_strided( + size=(cc_reshaped.shape[0], valid_shape_h, valid_shape_w), + stride=( + cc_reshaped.stride(0), + cc_reshaped.stride(1), + cc_reshaped.stride(2), + ), + ) + else: + cc_reshaped = cross_correlation.view( + -1, cross_correlation.shape[-2], cross_correlation.shape[-1] + ) + + # per-pixel maxima across the unraveled batch dimension + max_values, max_indices = torch.max(cc_reshaped, dim=0) + + # masked mip / index updates (do not modify originals here; return updated tensors) + update_mask = max_values > mip + new_mip = torch.where(update_mask, max_values, mip) + new_best_global_index = torch.where( + update_mask, current_indexes[max_indices], best_global_index + ) + + # sums used for statistics + corr_sum = cc_reshaped.sum(dim=0) + corr_sq_sum = (cc_reshaped * cc_reshaped).sum(dim=0) + + # find threshold exceedances (for correlation table) + batch_idxs, y_idxs, x_idxs = torch.where(cc_reshaped > threshold) + values = cc_reshaped[batch_idxs, y_idxs, x_idxs] + global_idxs = current_indexes[batch_idxs] + + return ( + new_mip, + new_best_global_index, + corr_sum, + corr_sq_sum, + global_idxs, + y_idxs, + x_idxs, + values, + ) + + +# pylint: disable=too-many-arguments +# pylint: disable=too-many-positional-arguments +# pylint: disable=too-many-locals +def do_iteration_and_correlation_table_updates( + cross_correlation: torch.Tensor, + current_indexes: torch.Tensor, + correlation_table: tensordict.TensorDict, + mip: torch.Tensor, + best_global_index: torch.Tensor, correlation_sum: torch.Tensor, correlation_squared_sum: torch.Tensor, - img_h: int, - img_w: int, + threshold: float, + valid_shape_h: int, + valid_shape_w: int, + needs_valid_cropping: bool = True, ) -> None: - """Helper function for updating maxima and tracked statistics. - - NOTE: The batch dimensions are effectively unraveled since taking the - maximum over a single batch dimensions is much faster than - multi-dimensional maxima. - - NOTE: Updating the maxima was found to be fastest and least memory - impactful when using torch.where directly. Other methods tested were - boolean masking and torch.where with tuples of tensor indexes. + """Helper function for updating maxima, tracked statistics, and correlation table. Parameters ---------- @@ -222,6 +323,9 @@ def do_iteration_statistics_updates( values, and orientations. Has shape `num_cs * num_defocus * num_orientations` to uniquely identify the set of pixel sizes, defocus values, and orientations associated with the batch from the global search space. + correlation_table : tensordict.TensorDict + The correlation table to update. Has keys + ["threshold", "pos_x", "pos_y", "corr_value"] each of which are tensors. mip : torch.Tensor Maximum intensity projection of the cross-correlation values. best_global_index : torch.Tensor @@ -230,37 +334,58 @@ def do_iteration_statistics_updates( Sum of cross-correlation values for each pixel. correlation_squared_sum : torch.Tensor Sum of squared cross-correlation values for each pixel. - img_h : int - Height of the cross-correlation values. - img_w : int - Width of the cross-correlation values. + threshold : float + The threshold value for adding entries to the correlation table. + valid_shape_h : int + Height of the valid region of the cross-correlation values. + valid_shape_w : int + Width of the valid region of the cross-correlation values. + needs_valid_cropping : bool, optional + Whether the cross-correlation tensor should be cropped (via a view operation) + to the valid dimensions (defined by `img_h` and `img_w`). If False, the + cross-correlation tensor is assumed to already be in the valid shape. """ - cc_reshaped = cross_correlation.view(-1, img_h, img_w) - - # Need two passes for maxima operator for memory efficiency - # and to distinguish between batch position which would both update - max_values, max_indices = torch.max(cc_reshaped, dim=0) - - # Do masked updates with torch.where directly (in-place) - update_mask = max_values > mip - torch.where(update_mask, max_values, mip, out=mip) - torch.where( - update_mask, - current_indexes[max_indices], + # call compiled core + ( + new_mip, + new_best_global_index, + corr_sum, + corr_sq_sum, + global_idxs, + y_idxs, + x_idxs, + values, + ) = _stats_and_table_core( + cross_correlation, + current_indexes, + mip, best_global_index, - out=best_global_index, + threshold, + valid_shape_h, + valid_shape_w, + needs_valid_cropping=needs_valid_cropping, ) - correlation_sum += cc_reshaped.sum(dim=0) - correlation_squared_sum += (cc_reshaped**2).sum(dim=0) + # update inplace the statistics tensors + mip.copy_(new_mip) + best_global_index.copy_(new_best_global_index) + + correlation_sum += corr_sum + correlation_squared_sum += corr_sq_sum + + # update correlation_table (tensordict operations not compiled) + if global_idxs.numel() > 0: + correlation_table["global_idx"] = torch.cat( + [correlation_table["global_idx"], global_idxs] + ) + correlation_table["pos_x"] = torch.cat([correlation_table["pos_x"], x_idxs]) + correlation_table["pos_y"] = torch.cat([correlation_table["pos_y"], y_idxs]) + correlation_table["corr_value"] = torch.cat( + [correlation_table["corr_value"], values] + ) # These are compiled normalization and stat update functions normalize_template_projection_compiled = attempt_torch_compilation( normalize_template_projection, backend="inductor", mode="default" ) -do_iteration_statistics_updates_compiled = attempt_torch_compilation( - do_iteration_statistics_updates, - backend="inductor", - mode="max-autotune-no-cudagraphs", -) diff --git a/src/leopard_em/pydantic_models/config/computational_config.py b/src/leopard_em/pydantic_models/config/computational_config.py index e795c8ff..b7cc1b63 100644 --- a/src/leopard_em/pydantic_models/config/computational_config.py +++ b/src/leopard_em/pydantic_models/config/computational_config.py @@ -6,7 +6,7 @@ from pydantic import BaseModel, Field # Type alias for non-negative integer -NonNegativeInt = Annotated[int, Field(ge=0)] +NonNegativeInt = Annotated[int, Field(ge=0)] # pylint: disable=invalid-name class BaseComputationalConfig(BaseModel): @@ -73,10 +73,13 @@ class ComputationalConfigMatch(BaseComputationalConfig): - The specific string "cpu" which means to use CPU. num_cpus : int Total number of CPUs to use, defaults to 1. - backend : Optional[str] - The backend to use for match template. - Must be "streamed" or "batched". - Defaults to "streamed". + backend : Literal["streamed", "batched", "zipfft"], optional + The cross-correlation backend to use for match template. Must be one of + "streamed", "batched", or "zipfft". When "streamed", individual 2D + cross-correlations are computed across multiple streams using PyTorch while + with "batched", all the 2D cross-correlations are computed in a single batched + call also with PyTorch. When "zipfft", the zipFFT library is used to compute the + cross-correlations. Defaults to "streamed". """ # Type-hinting here is ensuring non-negative integers, and list of at least one @@ -89,7 +92,7 @@ class ComputationalConfigMatch(BaseComputationalConfig): ] ] = [0] num_cpus: Annotated[int, Field(ge=1)] = 1 - backend: Literal["streamed", "batched"] = "streamed" + backend: Literal["streamed", "batched", "zipfft"] = "streamed" class ComputationalConfigRefine(BaseComputationalConfig): diff --git a/src/leopard_em/pydantic_models/data_structures/__init__.py b/src/leopard_em/pydantic_models/data_structures/__init__.py index 6538ed01..4bf76437 100644 --- a/src/leopard_em/pydantic_models/data_structures/__init__.py +++ b/src/leopard_em/pydantic_models/data_structures/__init__.py @@ -1,9 +1,11 @@ """Pydantic models for reused data structures across Leopard-EM programs.""" from .optics_group import OpticsGroup -from .particle_stack import ParticleStack +from .particle_stack import ParticleStack, ParticleStackCSV, ParticleStackHDF5 __all__ = [ "OpticsGroup", "ParticleStack", + "ParticleStackCSV", + "ParticleStackHDF5", ] diff --git a/src/leopard_em/pydantic_models/data_structures/particle_stack.py b/src/leopard_em/pydantic_models/data_structures/particle_stack.py index 3feab0a3..5d7c0044 100644 --- a/src/leopard_em/pydantic_models/data_structures/particle_stack.py +++ b/src/leopard_em/pydantic_models/data_structures/particle_stack.py @@ -1,14 +1,32 @@ -"""Particle stack Pydantic model for dealing with extracted particle data.""" +"""Particle stack Pydantic model for dealing with extracted particle data. + +Two public classes are provided for different storage back-ends: + +* ``ParticleStackCSV`` - the original behavior, loading particle data from a + CSV file and micrograph images from referenced paths on disk. + ``ParticleStack`` is an alias for this class for backward compatibility. +* ``ParticleStackHDF5`` - stores the particle table, optional image stack, and + optional per-particle local correlation statistics in a single HDF5 file. + +The base class ``_ParticleStackBase`` holds all shared computation methods and +tensor fields. It is not intended to be used directly. +""" + +# TODO: Move these into two separate files (long file) # pylint: disable=too-many-lines +import json import warnings +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path from typing import Any, ClassVar, Literal +import h5py import numpy as np import pandas as pd import torch -from pydantic import ConfigDict +from pydantic import ConfigDict, Field, model_validator from torch.utils.checkpoint import checkpoint from torch_cubic_spline_grids import CubicCatmullRomGrid3d from torch_fourier_shift import fourier_shift_dft_2d @@ -17,6 +35,7 @@ from torch_motion_correction.deformation_field_utils import ( evaluate_deformation_field_at_t, ) +from typing_extensions import Self from leopard_em.pydantic_models.config import PreprocessingFilters from leopard_em.pydantic_models.custom_types import ( @@ -33,21 +52,146 @@ "replicate": "edge", } +_HDF5_PARTICLES_GROUP = "particles" +_HDF5_LOCAL_STATS_GROUP = "local_stats" +_HDF5_IMAGE_STACK_DATASET = "image_stack" +_HDF5_STRING_DTYPE = h5py.string_dtype() + + +# TODO: Make this a shared utility function across the package somehow +def _leopard_em_version() -> str: + try: + return version("leopard_em") + except PackageNotFoundError: + return "uninstalled" + def _any_nan_or_inf(s: pd.Series) -> bool: - """Helper function to check if any value in the Series is NaN or infinite. + """Helper function to check if any value in the Series is NaN or infinite.""" + return bool(s.isna().any() or s.isin([float("inf"), float("-inf")]).any()) - Parameters - ---------- - s : pd.Series - The Series to check. - Returns - ------- - bool - True if any value in the Series is NaN or infinite, False otherwise. +def _generate_particle_ids(df: pd.DataFrame) -> list[str]: + """Generate particle IDs of the form ``{mic_stem}_{local_idx:05d}``.""" + ids: pd.Series = pd.Series("", index=df.index, dtype=object) + for mic_path, group in df.groupby("micrograph_path", sort=False): + stem = Path(str(mic_path)).stem + for local_idx, row_label in enumerate(group.index): + ids.at[row_label] = f"{stem}_{local_idx:05d}" + + res: list[str] = ids.tolist() + return res + + +# TODO: Better management of Zernikie coefficient columns/arrays in the HDF5 format... +# This is a lot of boilerplate code, and probably a better schema would eliminate +# these parsing needs. +def _value_to_str(v: Any) -> str: + """Serialize a value to a string for HDF5 storage.""" + if v is None: + return "" + if isinstance(v, str): + return v + return json.dumps(v) + + +def _str_to_value(s: str) -> Any: + """Deserialize a string back to a Python value after HDF5 load.""" + if s == "": + return None + try: + parsed = json.loads(s) + if isinstance(parsed, (list, dict)): + return parsed + # Plain JSON scalars (numbers) that were originally strings stay as strings + return s + except (json.JSONDecodeError, ValueError): + return s + + +# NOTE: How are the internals of the hdf5 particle stack being handled? Are they just a +# pass through for the DataFrame type backed class (don't want this). Need to +# implement things at the base class level somehow. +def _write_df_to_hdf5_group(f: h5py.File, df: pd.DataFrame) -> None: + """Write a DataFrame's columns to ``f[_HDF5_PARTICLES_GROUP]``. + + Numeric columns are stored as float64 datasets. String / object columns + (including path columns, Zernike coefficient arrays, etc.) are serialized + to variable-length UTF-8 strings via ``_value_to_str``. + + The dataset names match the DataFrame column names. ``particle_id`` + (which may be the DataFrame index) is always written as an explicit + dataset and listed first in ``attrs["columns"]``. """ - return bool(s.isna().any() or s.isin([float("inf"), float("-inf")]).any()) + grp = f.create_group(_HDF5_PARTICLES_GROUP) + + # Build the list of columns to write, ensuring particle_id comes first. + if df.index.name == "particle_id": + particle_ids = df.index.tolist() + col_names = ["particle_id", *list(df.columns)] + else: + # particle_id may be an ordinary column + particle_ids = df["particle_id"].tolist() if "particle_id" in df.columns else [] + col_names = list(df.columns) + + grp.attrs["columns"] = col_names + + # Write particle_id dataset + if particle_ids: + grp.create_dataset( + "particle_id", + data=np.array([str(v) for v in particle_ids], dtype=object), + dtype=_HDF5_STRING_DTYPE, + ) + + for col in df.columns: + if col == "particle_id": + # Already written above (or will be skipped if index) + if df.index.name != "particle_id": + continue + series = df[col] + if pd.api.types.is_float_dtype(series) or pd.api.types.is_integer_dtype(series): + grp.create_dataset(col, data=series.to_numpy(dtype=np.float64)) + else: + str_data = [_value_to_str(v) for v in series] + grp.create_dataset( + col, + data=np.array(str_data, dtype=object), + dtype=_HDF5_STRING_DTYPE, + ) + + +def _read_df_from_hdf5_group(f: h5py.File) -> pd.DataFrame: + """Reconstruct a DataFrame from ``f[_HDF5_PARTICLES_GROUP]``. + + ``particle_id`` is restored as the pandas ``Index``. + """ + grp = f[_HDF5_PARTICLES_GROUP] + columns: list[str] = list(grp.attrs["columns"]) + + data: dict[str, Any] = {} + for col in columns: + if col not in grp: + continue + raw = grp[col][:] + if raw.dtype.kind in ("O", "S", "U"): + decoded = [s.decode() if isinstance(s, bytes) else s for s in raw] + data[col] = [_str_to_value(s) for s in decoded] + else: + data[col] = raw + + df = pd.DataFrame(data) + + if "particle_id" in df.columns: + df = df.set_index("particle_id") + df.index.name = "particle_id" + + return df + + +# --------------------------------------------------------------------------- +# Stand-alone image-extraction helpers (unchanged from original module) +# --------------------------------------------------------------------------- def get_cropped_image_regions( @@ -130,8 +274,6 @@ def get_cropped_image_regions( if isinstance(box_size, int): box_size = (box_size, box_size) - # The underlying numpy/torch functions only operate on the top-left corner - # reference, so shift the position half a box height/width if using center. if pos_reference == "center": pos_y = pos_y - box_size[0] // 2 pos_x = pos_x - box_size[1] // 2 @@ -198,7 +340,6 @@ def _get_cropped_image_regions_numpy( regions = [] for y, x in zip(pos_y, pos_x): - # Check bounds and raise error if out of bounds if ( y < 0 or x < 0 @@ -251,12 +392,10 @@ def _get_cropped_image_regions_torch( regions = [] for y, x in zip(pos_y, pos_x): - # Convert to Python ints for comparison y = int(y.item() if hasattr(y, "item") else y) x = int(x.item() if hasattr(x, "item") else x) original_y, original_x = y, x - # Check bounds if ( y < 0 or x < 0 @@ -269,7 +408,6 @@ def _get_cropped_image_regions_torch( f"{original_x}:{original_x + box_size[1]}] exceed " f"image dimensions {image.shape}" ) - # For "pad" mode, warn and clamp coordinates warnings.warn( f"Region bounds [{original_y}:{original_y + box_size[0]}, " f"{original_x}:{original_x + box_size[1]}] exceed " @@ -277,131 +415,301 @@ def _get_cropped_image_regions_torch( UserWarning, stacklevel=2, ) - # Clamp coordinates to keep region within image bounds y = max(0, min(y, image.shape[0] - box_size[0])) x = max(0, min(x, image.shape[1] - box_size[1])) regions.append(image[y : y + box_size[0], x : x + box_size[1]]) - # Stack all regions cropped_images = torch.stack(regions) return cropped_images -class ParticleStack(BaseModel2DTM): - """Pydantic model for dealing with particle stack data. +# --------------------------------------------------------------------------- +# Base class +# --------------------------------------------------------------------------- + + +# pylint: disable=too-many-instance-attributes +class _ParticleStackBase(BaseModel2DTM): + """Base class holding particle stack data, preprocessing state, and compute methods. + + Not intended to be instantiated directly — use ``ParticleStackCSV`` or + ``ParticleStackHDF5`` depending on the desired storage back-end. Attributes ---------- - df_path : str - Path to the DataFrame containing the particle data. The DataFrame must have - the following columns (see the documentation for further information): - - - mip - - scaled_mip - - correlation_mean - - correlation_variance - - total_correlations - - pos_x - - pos_y - - pos_x_img - - pos_y_img - - pos_x_img_angstrom - - pos_y_img_angstrom - - psi - - theta - - phi - - relative_defocus - - refined_relative_defocus - - defocus_u - - defocus_v - - astigmatism_angle - - pixel_size - - refined_pixel_size - - voltage - - spherical_aberration - - amplitude_contrast_ratio - - phase_shift - - ctf_B_factor - - micrograph_path - - template_path - - mip_path - - scaled_mip_path - - psi_path - - theta_path - - phi_path - - defocus_path - - correlation_average_path - - correlation_variance_path - + leopard_em_version : str + Version of Leopard-EM that created this particle stack. Auto-populated + from installed package metadata; preserved as-recorded when loading from + a file. extracted_box_size : tuple[int, int] - The size of the extracted particle boxes in pixels in units of pixels. + Size of extracted particle boxes in pixels (height, width). original_template_size : tuple[int, int] - The original size of the template used during the matching process. Should be - smaller than the extracted box size. + Size of the template used during template matching (height, width). + Must be smaller than or equal to ``extracted_box_size``. + global_whitening_applied : bool + True if whitening was computed from and applied to the full micrograph + before particle extraction. + local_whitening_applied : bool + True if whitening was computed from and applied to each individual + extracted particle box. + global_normalization_applied : bool + True if normalization was computed from the full micrograph before + extraction. + local_normalization_applied : bool + True if normalization was computed from and applied to each extracted + particle box. image_stack : ExcludedTensor - The stack of images extracted from the micrographs. Is effectively a pytorch - Tensor with shape (N, H, W) where N is the number of particles and (H, W) is - the extracted box size. + Stack of extracted particle images, shape ``(N, box_h, box_w)``. + Not serialized to YAML/JSON. + local_stats_correlation_average : ExcludedTensor + Per-particle local mean of the cross-correlation map, extracted from + the valid cross-correlation region around each particle center. + Shape ``(N, valid_h, valid_w)`` where + ``valid_h = extracted_box_size[0] - original_template_size[0] + 1`` and + ``valid_w = extracted_box_size[1] - original_template_size[1] + 1``. + Not serialized to YAML/JSON. + local_stats_correlation_variance : ExcludedTensor + Per-particle local variance of the cross-correlation map. Same shape + as ``local_stats_correlation_average``. Not serialized to YAML/JSON. """ model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True) - # Serialized fields - df_path: str + leopard_em_version: str = Field(default_factory=_leopard_em_version) extracted_box_size: tuple[int, int] original_template_size: tuple[int, int] - # Imported tabular data (not serialized) + # Pre-processing state flags + global_whitening_applied: bool = False + local_whitening_applied: bool = False + global_normalization_applied: bool = False + local_normalization_applied: bool = False + + # Private: tabular data (not part of Pydantic schema) + # TODO: Move away from having a df-backed implementation in favor of either + # getter/setter methods OR private fields for the relevant data. _df: pd.DataFrame - # Cropped out view of the particles from images + # Image and statistics tensors (excluded from YAML/JSON serialization) image_stack: ExcludedTensor + local_stats_correlation_average: ExcludedTensor + local_stats_correlation_variance: ExcludedTensor - def __init__(self, skip_df_load: bool = False, **data: dict[str, Any]): - """Initialize the ParticleStack object. + def __init__(self, skip_df_load: bool = False, **data: Any): + """Initialize the particle stack. Parameters ---------- skip_df_load : bool, optional - Whether to skip loading the DataFrame, by default False and the dataframe - is loaded automatically. + When True the subclass ``load_df`` is not called automatically. + Use this when constructing an empty instance before populating + ``_df`` manually (e.g., during ``from_hdf5``). data : dict[str, Any] - The data to initialize the object with. + Fields forwarded to the Pydantic constructor. """ super().__init__(**data) - if not skip_df_load: self.load_df() def load_df(self) -> None: - """Load the DataFrame from the specified path. + """Load the particle DataFrame from the backing store. - Raises - ------ - ValueError - If the DataFrame is missing required columns. + Subclasses must override this method. """ - tmp_df = pd.read_csv(self.df_path) - - # Validate the DataFrame columns - missing_columns = [ - col for col in MATCH_TEMPLATE_DF_COLUMN_ORDER if col not in tmp_df.columns - ] - if missing_columns: - raise ValueError( - f"Missing the following columns in DataFrame: {missing_columns}" - ) + raise NotImplementedError("Subclasses must implement load_df()") - self._df = tmp_df + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ def _get_position_reference_columns(self) -> tuple[str, str]: - """Get the position reference columns based on the DataFrame.""" + """Return the y/x position column names to use (refined preferred).""" y_col = "refined_pos_y" if "refined_pos_y" in self._df.columns else "pos_y" x_col = "refined_pos_x" if "refined_pos_x" in self._df.columns else "pos_x" return y_col, x_col + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def df_columns(self) -> list[str]: + """Column names of the underlying DataFrame.""" + return list(self._df.columns.tolist()) + + @property + def num_particles(self) -> int: + """Number of particles in the stack.""" + return len(self._df) + + # ------------------------------------------------------------------ + # DataFrame accessor / mutator helpers + # ------------------------------------------------------------------ + + def __getitem__(self, key: str) -> Any: + """Get a column from the underlying DataFrame.""" + try: + return self._df[key] + except KeyError as err: + raise KeyError(f"Key '{key}' not found in underlying DataFrame.") from err + + def set_column(self, column_name: str, value: Any) -> None: + """Set a column in the underlying DataFrame. + + Parameters + ---------- + column_name : str + The column name to set. + value : Any + The value(s) to assign. + """ + self._df.loc[:, column_name] = value + + def get_dataframe_copy(self) -> pd.DataFrame: + """Return a copy of the underlying DataFrame. + + Returns + ------- + pd.DataFrame + """ + return self._df.copy() + + # ------------------------------------------------------------------ + # CTF / orientation accessors + # ------------------------------------------------------------------ + + def get_relative_defocus( + self, + prefer_refined_defocus: bool = True, + ) -> torch.Tensor: + """Get the relative defocus values for each particle. + + Parameters + ---------- + prefer_refined_defocus : bool, optional + Whether to use the refined defocus values, by default True. + + Returns + ------- + torch.Tensor + """ + rel_defocus_col = "relative_defocus" + if prefer_refined_defocus: + if "refined_relative_defocus" not in self._df.columns: + warnings.warn( + "Refined defocus values not found in DataFrame, using original " + "defocus values...", + stacklevel=2, + ) + elif _any_nan_or_inf(self._df["refined_relative_defocus"]): + warnings.warn( + "Refined defocus values contain NaN or inf values, using original " + "defocus values...", + stacklevel=2, + ) + else: + rel_defocus_col = "refined_relative_defocus" + + return torch.tensor(self._df[rel_defocus_col].to_numpy().copy()) + + def get_absolute_defocus( + self, prefer_refined_defocus: bool = True + ) -> tuple[torch.Tensor, torch.Tensor]: + """Get the absolute defocus (u, v) values for each particle. + + Parameters + ---------- + prefer_refined_defocus : bool, optional + Whether to use refined defocus, by default True. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + ``(defocus_u, defocus_v)`` tensors in Angstroms. + """ + particle_defocus = self.get_relative_defocus(prefer_refined_defocus) + defocus_u = torch.tensor(self._df["defocus_u"].to_numpy().copy()) + defocus_v = torch.tensor(self._df["defocus_v"].to_numpy().copy()) + defocus_u = defocus_u + particle_defocus + defocus_v = defocus_v + particle_defocus + return defocus_u, defocus_v + + def get_pixel_size( + self, + prefer_refined_pixel_size: bool = True, + ) -> torch.Tensor: + """Get the pixel size for each particle. + + Parameters + ---------- + prefer_refined_pixel_size : bool, optional + Whether to use the refined pixel size, by default True. + + Returns + ------- + torch.Tensor + """ + pixel_size_col = "pixel_size" + if prefer_refined_pixel_size: + if "refined_pixel_size" not in self._df.columns: + warnings.warn( + "Refined pixel size not found in DataFrame, using original" + " pixel size values...", + stacklevel=2, + ) + elif _any_nan_or_inf(self._df["refined_pixel_size"]): + warnings.warn( + "Refined pixel size contain NaN or inf values, using original" + " pixel size values...", + stacklevel=2, + ) + else: + pixel_size_col = "refined_pixel_size" + + return torch.tensor(self._df[pixel_size_col].to_numpy().copy()) + + def get_euler_angles(self, prefer_refined_angles: bool = True) -> torch.Tensor: + """Return the Euler angles (phi, theta, psi) of all particles as a tensor. + + Parameters + ---------- + prefer_refined_angles : bool, optional + When true, refined angles are used if present, by default True. + + Returns + ------- + torch.Tensor + Shape ``(N, 3)`` — columns correspond to (phi, theta, psi) in ZYZ. + """ + phi_col = "phi" + theta_col = "theta" + psi_col = "psi" + if prefer_refined_angles: + if not all( + x in self._df.columns + for x in ["refined_phi", "refined_theta", "refined_psi"] + ): + warnings.warn( + "Refined angles not found in DataFrame, using original angles...", + stacklevel=2, + ) + else: + phi_col = "refined_phi" + theta_col = "refined_theta" + psi_col = "refined_psi" + + phi = torch.tensor(self._df[phi_col].to_numpy().copy()) + theta = torch.tensor(self._df[theta_col].to_numpy().copy()) + psi = torch.tensor(self._df[psi_col].to_numpy().copy()) + + return torch.stack((phi, theta, psi), dim=-1) + + # ------------------------------------------------------------------ + # Image-stack construction + # ------------------------------------------------------------------ + def load_images_grouped_by_column( self, column_name: str ) -> tuple[torch.Tensor, list[pd.Index]]: @@ -424,7 +732,6 @@ def load_images_grouped_by_column( if column_name not in self._df.columns: raise ValueError(f"Column '{column_name}' not found in the DataFrame.") - # Find the indexes in the DataFrame that correspond to each unique image image_index_groups = self._df.groupby(column_name).groups images_list = [] indices = [] @@ -433,7 +740,6 @@ def load_images_grouped_by_column( images_list.append(img) indices.append(indexes) - # Stack images into a tensor (N, H, W) images_tensor = torch.stack(images_list, dim=0) return images_tensor, indices @@ -493,86 +799,53 @@ def construct_image_stack( Parameters ---------- images : torch.Tensor - A tensor of loaded images with shape (N, H, W) where N is the number of - images and (H, W) is the image size. + A tensor of loaded images with shape (N, H, W). indices : list[pd.Index] - A list of pandas Index objects containing the row indexes for particles - from each corresponding image. Should be the same length as the first - dimension of `images`. + Row indexes for particles from each corresponding image. extraction_size : tuple[int, int] - The size of the extracted boxes in pixels (height, width). + Size of the extracted boxes in pixels (height, width). pos_reference : Literal["center", "top-left"], optional - The reference point for the positions, by default "top-left". If "center", - the boxes extracted will be - image[y - box_size // 2 : y + box_size // 2, ...]. - Columns in the dataframe which are used as position references are always - pos_x and pos_y, or refined_pos_x and refined_pos_y if available. - If "top-left", the boxes will be image[y : y + box_size, ...]. - Leopard-EM uses the "top-left" reference position, and unless you know data - was processed in a different way you should not change this value. - handle_bounds : Literal["pad", "clip", "error"], optional - How to handle the bounds of the image, by default "pad". If "pad", the image - will be padded with the padding value based on the padding mode. If "error", - an error will be raised if any region exceeds the image bounds. NOTE: - clipping is not supported since returned stack may have inhomogeneous sizes. + Reference point for the positions, by default "top-left". + handle_bounds : Literal["pad", "error"], optional + How to handle out-of-bounds regions, by default "pad". padding_mode : Literal["constant", "reflect", "replicate"], optional - The padding mode to use when padding the image, by default "constant". - "constant" pads with the value `padding_value`, "reflect" pads with the - reflection of the image at the edge, and "replicate" pads with the last - pixel of the image. These match the modes available in - `torch.nn.functional.pad`. + Padding mode when ``handle_bounds="pad"``, by default "constant". padding_value : float, optional - The value to use for padding when `padding_mode` is "constant", by default - 0.0. + Constant padding value, by default 0.0. Returns ------- torch.Tensor - The stack of images, this is the internal 'image_stack' attribute. + Stack of extracted images ``(N, extraction_h, extraction_w)``. """ - # Determine which position columns to use (refined if available) y_col, x_col = self._get_position_reference_columns() - # Create an empty tensor to store the image stack on the same device as images h, w = self.original_template_size box_h, box_w = self.extracted_box_size device = images.device image_stack = torch.zeros((self.num_particles, *extraction_size), device=device) - # Verify that the number of images matches the number of indices if images.shape[0] != len(indices): raise ValueError( f"Number of images ({images.shape[0]}) does not match the number of " f"indices ({len(indices)})." ) - # Loop over each image and its corresponding indexes for i, indexes in enumerate(indices): img = images[i] - # Get the positions as numpy arrays for indexing pos_y = self._df.loc[indexes, y_col].to_numpy().copy() pos_x = self._df.loc[indexes, x_col].to_numpy().copy() - # If the position reference is "center", shift (x, y) by half the original - # template width/height so reference is now the top-left corner if pos_reference == "center": pos_y = pos_y - h // 2 pos_x = pos_x - w // 2 - # Our reference is now a top-left corner of a box of the original template - # shape, BUT we want a slightly larger box of extraction_size AND this - # box to be centered around the particle. Therefore, need to shift the - # position half the difference between the original template size and - # the extraction size. pos_y = pos_y - (box_h - h) // 2 pos_x = pos_x - (box_w - w) // 2 pos_y = torch.tensor(pos_y, device=img.device) pos_x = torch.tensor(pos_x, device=img.device) - # Code logic is simplified by only using the top-left reference position - # in the `get_cropped_image_regions` function. Relative referencing handled - # by the ParticleStack class. cropped_images = get_cropped_image_regions( img, pos_y, @@ -621,7 +894,6 @@ def construct_image_filters( num_images = images_dft.shape[0] filter_stack = torch.zeros((num_images, *output_shape), device=device) - # Loop over each image and compute the filter for i in range(num_images): img_dft = images_dft[i] cumulative_filter = preprocess_filters.get_combined_filter( @@ -653,30 +925,23 @@ def construct_projective_filters( output_shape : tuple[int, int] What shape along the last two dimensions the filters should be. images_dft : torch.Tensor - A tensor of micrograph images in Fourier space with shape (N, H, W) where N - is the number of unique micrographs and (H, W) is the Fourier space size. + A tensor of micrograph images in Fourier space with shape (N, H, W). indices : list[pd.Index] - A list of pandas Index objects containing the row indexes for particles - from each corresponding micrograph. Should be the same length as the first - dimension of `images_dft`. + Row indexes for particles from each corresponding micrograph. Returns ------- torch.Tensor - The stack of filters with shape (M, h, w) where M is the number of particles - and (h, w) is the output shape. + Filter stack of shape ``(M, h, w)`` where M is the number of particles. """ - # Create an empty tensor to store the filter stack device = images_dft.device filter_stack = torch.zeros((self.num_particles, *output_shape), device=device) - # Verify that the number of images matches the number of indices if images_dft.shape[0] != len(indices): raise ValueError( f"Number of images ({images_dft.shape[0]}) does not match " f"the number of indices ({len(indices)})." ) - # Loop over each micrograph and its corresponding indexes for i, indexes in enumerate(indices): img_dft = images_dft[i] cumulative_filter = preprocess_filters.get_combined_filter( @@ -688,199 +953,6 @@ def construct_projective_filters( return filter_stack - @property - def df_columns(self) -> list[str]: - """Get the columns of the DataFrame.""" - return list(self._df.columns.tolist()) - - @property - def num_particles(self) -> int: - """Get the number of particles in the stack.""" - return len(self._df) - - def get_relative_defocus( - self, - prefer_refined_defocus: bool = True, - ) -> torch.Tensor: - """Get the relative defocus values for each particle. - - Parameters - ---------- - prefer_refined_defocus : bool, optional - Whether to use the refined defocus values (columns prefixed with 'refined_') - or not, by default True. - - Returns - ------- - torch.Tensor - The relative defocus values for each particle. - - Warnings - -------- - Warns if NaN values or no column present for either - 'refined_relative_defocus' or 'relative_defocus'. - Falls back to the unrefined values. - """ - rel_defocus_col = "relative_defocus" - # Both refined columns must be present AND no values can be NaN or inf - if prefer_refined_defocus: - if "refined_relative_defocus" not in self._df.columns: - warnings.warn( - "Refined defocus values not found in DataFrame, using original " - "defocus values...", - stacklevel=2, - ) - elif _any_nan_or_inf(self._df["refined_relative_defocus"]): - warnings.warn( - "Refined defocus values contain NaN or inf values, using original " - "defocus values...", - stacklevel=2, - ) - else: - rel_defocus_col = "refined_relative_defocus" - - return torch.tensor(self._df[rel_defocus_col].to_numpy().copy()) - - def get_absolute_defocus( - self, prefer_refined_defocus: bool = True - ) -> tuple[torch.Tensor, torch.Tensor]: - """Get the absolute defocus values for each particle. - - NOTE: If the refined defocus values are requested but not present in the - DataFrame (either no column or any NaN values), a user warning is raised - and the original defocus values are returned instead. - - Parameters - ---------- - prefer_refined_defocus : bool, optional - Whether to use the refined defocus values - (columns prefixed with 'refined_') or not, by default True. - - Returns - ------- - tuple[torch.Tensor, torch.Tensor] - A tuple of two tensors containing the absolute defocus values along the - major (defocus_u) and minor axes (defocus_v), respectively in units of - Angstroms. - """ - particle_defocus = self.get_relative_defocus(prefer_refined_defocus) - defocus_u = torch.tensor(self._df["defocus_u"].to_numpy().copy()) - defocus_v = torch.tensor(self._df["defocus_v"].to_numpy().copy()) - defocus_u = defocus_u + particle_defocus - defocus_v = defocus_v + particle_defocus - - return defocus_u, defocus_v - - def get_pixel_size( - self, - prefer_refined_pixel_size: bool = True, - ) -> torch.Tensor: - """Get the relative pixel size values for each particle. - - Parameters - ---------- - prefer_refined_pixel_size : bool, optional - Whether to use the refined pixel size values - (columns prefixed with 'refined_') or not, by default True. - - Returns - ------- - torch.Tensor - The relative pixel size values for each particle. - - Warnings - -------- - Warns if NaN values or no column present for either 'refined_pixel_size' - or 'pixel_size'. Falls back to the unrefined values. - """ - pixel_size_col = "pixel_size" - if prefer_refined_pixel_size: - if "refined_pixel_size" not in self._df.columns: - warnings.warn( - "Refined pixel size not found in DataFrame, using original" - " pixel size values...", - stacklevel=2, - ) - elif _any_nan_or_inf(self._df["refined_pixel_size"]): - warnings.warn( - "Refined pixel size contain NaN or inf values, using original" - " pixel size values...", - stacklevel=2, - ) - else: - pixel_size_col = "refined_pixel_size" - - return torch.tensor(self._df[pixel_size_col].to_numpy().copy()) - - def get_euler_angles(self, prefer_refined_angles: bool = True) -> torch.Tensor: - """Return the Euler angles (phi, theta, psi) of all particles as a tensor. - - Parameters - ---------- - prefer_refined_angles : bool, optional - When true, the refined Euler angles are used (columns prefixed with - 'refined_'), otherwise the original angles are used, by default True. - - Returns - ------- - torch.Tensor - A tensor of shape (N, 3) where N is the number of particles and the columns - correspond to (phi, theta, psi) in ZYZ format. - """ - # Ensure all three refined columns are present, warning if not - phi_col = "phi" - theta_col = "theta" - psi_col = "psi" - if prefer_refined_angles: - if not all( - x in self._df.columns - for x in ["refined_phi", "refined_theta", "refined_psi"] - ): - warnings.warn( - "Refined angles not found in DataFrame, using original angles...", - stacklevel=2, - ) - else: - phi_col = "refined_phi" - theta_col = "refined_theta" - psi_col = "refined_psi" - - # Get the angles from the DataFrame - phi = torch.tensor(self._df[phi_col].to_numpy().copy()) - theta = torch.tensor(self._df[theta_col].to_numpy().copy()) - psi = torch.tensor(self._df[psi_col].to_numpy().copy()) - - return torch.stack((phi, theta, psi), dim=-1) - - def __getitem__(self, key: str) -> Any: - """Get an item from the DataFrame.""" - try: - return self._df[key] - except KeyError as err: - raise KeyError(f"Key '{key}' not found in underlying DataFrame.") from err - - def set_column(self, column_name: str, value: Any) -> None: - """Set a column in the underlying DataFrame. - - Parameters - ---------- - column_name : str - The name of the column to set - value : Any - The value to set the column to - """ - self._df.loc[:, column_name] = value - - def get_dataframe_copy(self) -> pd.DataFrame: - """Return a copy of the underlying DataFrame. - - Returns - ------- - pd.DataFrame - A copy of the underlying DataFrame - """ - return self._df.copy() - @staticmethod # pylint: disable=too-many-arguments # pylint: disable=too-many-positional-arguments @@ -894,33 +966,30 @@ def _process_single_frame_with_shifts_checkpoint( padding_mode: Literal["constant", "reflect", "replicate"], padding_value: float, ) -> torch.Tensor: - """ - Process a single frame using *precomputed particle shifts*. + """Process a single frame using precomputed particle shifts. - This function is safe for gradient checkpointing and contains no - deformation-field evaluation. + Safe for gradient checkpointing; contains no deformation-field evaluation. Parameters ---------- movie_frame : torch.Tensor - Single movie frame (H, W) + Single movie frame (H, W). shifts : torch.Tensor - Per-particle shifts with shape (N, 2) as (dy, dx) + Per-particle shifts with shape (N, 2) as (dy, dx). pos_y, pos_x : torch.Tensor - Top-left extraction positions + Top-left extraction positions. extracted_box_size : tuple[int, int] - (box_h, box_w) + ``(box_h, box_w)``. handle_bounds, padding_mode, padding_value - Passed through to cropping + Passed through to cropping. Returns ------- torch.Tensor - Shifted FFTs with shape (N, box_h, box_w//2 + 1) + Shifted FFTs with shape ``(N, box_h, box_w//2 + 1)``. """ box_h, box_w = extracted_box_size - # Extract particle images cropped_images = get_cropped_image_regions( movie_frame, pos_y, @@ -932,12 +1001,10 @@ def _process_single_frame_with_shifts_checkpoint( padding_value=padding_value, ) - # FFT cropped_images_dft = torch.fft.rfftn( # pylint: disable=not-callable cropped_images, dim=(-2, -1) ) - # Fourier shift shifted_fft = fourier_shift_dft_2d( dft=cropped_images_dft, image_shape=(box_h, box_w), @@ -960,34 +1027,33 @@ def compute_frame_particle_shifts_from_deformation( gh: int, gw: int, ) -> torch.Tensor: - """ - Compute per-particle shifts for a single frame from a deformation field. + """Compute per-particle shifts for a single frame from a deformation field. Parameters ---------- movie_frame : torch.Tensor - Single movie frame (H, W) + Single movie frame (H, W). deformation_field : CubicCatmullRomGrid3d The deformation field grid. normalized_t_value : torch.Tensor - The normalized time value for the frame. + Normalized time value for the frame. pixel_grid : torch.Tensor The pixel grid tensor. pixel_spacing : float The pixel spacing. pos_y_center : torch.Tensor - The center y position. + Center y positions. pos_x_center : torch.Tensor - The center x position. + Center x positions. gh : int - The height of the deformation field grid. + Height of the deformation field grid. gw : int - The width of the deformation field grid. + Width of the deformation field grid. Returns ------- torch.Tensor - Shifts with shape (N, 2) as (dy, dx) + Shifts with shape ``(N, 2)`` as (dy, dx). """ frame_deformation_field = evaluate_deformation_field_at_t( deformation_field=deformation_field, @@ -1034,55 +1100,36 @@ def construct_image_stack_from_movie( deformation_field : CubicCatmullRomGrid3d | None, optional The deformation field grid. particle_shifts : torch.Tensor | None, optional - The particle shifts to apply to the movie. If None, the particle shifts - are computed from the deformation field. If provided, the particle shifts - are used to shift the movie. One must be provided. - Shape is (T, N, 2) where T = number of frames, N = number of particles, + Per-particle shifts, shape ``(T, N, 2)``. Exactly one of + ``deformation_field`` and ``particle_shifts`` must be provided. pos_reference : Literal["center", "top-left"], optional - The reference point for the positions, by default "top-left". If "center", - the boxes extracted are image[y - box_size // 2 : y + box_size // 2, ...]. - If "top-left", the boxes will be image[y : y + box_size, ...]. + Position reference for extraction, by default "top-left". handle_bounds : Literal["pad", "error"], optional - How to handle the bounds of the image, by default "pad". If "pad", the image - will be padded with the padding value based on the padding mode. - If "error", an error will be raised if any region exceeds the image bounds. - Note clipping is not supported - since returned stack may have inhomogeneous sizes. + How to handle out-of-bounds regions, by default "pad". padding_mode : Literal["constant", "reflect", "replicate"], optional - The padding mode to use when padding the image, by default "constant". - "constant" pads with the value `padding_value`, "reflect" pads with the - reflection of the image, and "replicate" pads with the last pixel - of the image. These match the modes available in `torch.nn.functional.pad`. + Padding mode, by default "constant". padding_value : float, optional - The value to use for padding when `padding_mode` is "constant", - by default 0.0. + Constant padding value, by default 0.0. pre_exposure : float, optional - The pre-exposure time in seconds, by default 0.0. + Pre-exposure in electrons per pixel, by default 0.0. fluence_per_frame : float, optional - The dose per frame in electrons per pixel, by default 0.0. + Dose per frame in electrons per pixel, by default 0.0. use_gradient_checkpointing : bool, optional - Whether to use gradient checkpointing to save memory during frame - processing. Checkpointing trades compute time for memory by not - storing intermediate activations. Defaults to True. + Trade compute for memory during frame processing, by default True. particle_indices : list[int] | None, optional - Indices of particles to process from the dataframe. If None, - processes all particles. Use this to batch particles for memory - efficiency during gradient-based optimization. Defaults to None. + Subset of particles to process. If None, all particles are used. Returns ------- torch.Tensor - The stack of images with shape (N, H, W) where N is the number of particles - and (H, W) is the extracted box size. + Image stack of shape ``(N, box_h, box_w)``. """ if (deformation_field is None) == (particle_shifts is None): raise ValueError( "One of `deformation_field` or `particle_shifts` must be provided." ) pixel_sizes = self.get_pixel_size() - # Determine which position columns to use (refined if available) y_col, x_col = self._get_position_reference_columns() - # Create an empty tensor to store the image stack h, w = self.original_template_size box_h, box_w = self.extracted_box_size t, img_h, img_w = movie.shape @@ -1095,20 +1142,15 @@ def construct_image_stack_from_movie( image_shape=(img_h, img_w), device=movie.device, ) - # Find the indexes in the DataFrame that correspond to each unique image if particle_indices is not None: - # Use provided subset of particles paticle_indexes = [self._df.index[i] for i in particle_indices] num_particles_to_process = len(particle_indices) else: - # Use all particles paticle_indexes = self._df.index.tolist() num_particles_to_process = self.num_particles pos_y = self._df.loc[paticle_indexes, y_col].to_numpy() pos_x = self._df.loc[paticle_indexes, x_col].to_numpy() - # If the position reference is "top-left", shift (x, y) by half the original - # template width/height so reference is now in the center if pos_reference == "center": pos_y = pos_y - h // 2 pos_x = pos_x - w // 2 @@ -1127,13 +1169,9 @@ def construct_image_stack_from_movie( dtype=torch.complex64, device=movie.device, ) - # set frames mean zero movie = movie - torch.mean(movie, dim=(-2, -1), keepdim=True) for frame_index, movie_frame in enumerate(movie): - # ------------------------------------------------------------ - # Obtain shifts (dy, dx) for this frame - # ------------------------------------------------------------ if particle_shifts is not None: frame_shifts = particle_shifts[frame_index] # (N, 2) else: @@ -1149,9 +1187,6 @@ def construct_image_stack_from_movie( gw=gw, ) - # ------------------------------------------------------------ - # Apply shifts + FFT (checkpointed) - # ------------------------------------------------------------ if use_gradient_checkpointing: shifted_fft = checkpoint( self._process_single_frame_with_shifts_checkpoint, @@ -1177,13 +1212,11 @@ def construct_image_stack_from_movie( padding_value=padding_value, ) - # Store the shifted FFTs aligned_particle_movies_rfft[:, frame_index] = shifted_fft - # Clear cache periodically to help with memory if frame_index % 10 == 0 and frame_index > 0: torch.cuda.empty_cache() - # Dose weight the aligned particle images + aligned_particle_images = torch.zeros( (num_particles_to_process, box_h, box_w), device=movie.device, @@ -1191,7 +1224,6 @@ def construct_image_stack_from_movie( for particle_index in range(num_particles_to_process): particle_dft = aligned_particle_movies_rfft[particle_index] - # Get the actual dataframe index for this particle df_idx = paticle_indexes[particle_index] df_loc = self._df.index.get_loc(df_idx) @@ -1201,10 +1233,393 @@ def construct_image_stack_from_movie( pre_exposure=pre_exposure, fluence_per_frame=fluence_per_frame, voltage=self._df["voltage"].to_numpy()[df_loc], - ) # (box_h, box_w) + ) aligned_particle_images[particle_index] = dw_sum - # Only update self.image_stack if processing all particles if particle_indices is None: self.image_stack = aligned_particle_images return aligned_particle_images + + +# --------------------------------------------------------------------------- +# CSV-backed subclass +# --------------------------------------------------------------------------- + + +class ParticleStackCSV(_ParticleStackBase): + """Particle stack whose tabular data is loaded from a CSV file. + + Particle images are extracted from the micrograph paths referenced in the + CSV at run time. This is the original ``ParticleStack`` behavior. + + Attributes + ---------- + df_path : str + Path to the CSV file containing the particle data. + """ + + df_path: str + + def load_df(self) -> None: + """Load and validate the particle DataFrame from ``df_path``. + + Raises + ------ + ValueError + If required columns are missing from the CSV. + """ + tmp_df = pd.read_csv(self.df_path) + + missing_columns = [ + col for col in MATCH_TEMPLATE_DF_COLUMN_ORDER if col not in tmp_df.columns + ] + if missing_columns: + raise ValueError( + f"Missing the following columns in DataFrame: {missing_columns}" + ) + + self._df = tmp_df + + def to_hdf5( + self, + hdf5_path: str, + allow_file_overwrite: bool = False, + include_image_stack: bool = False, + include_local_stats: bool = False, + ) -> "ParticleStackHDF5": + """Convert this CSV-backed stack to an HDF5-backed stack and write to disk. + + Parameters + ---------- + hdf5_path : str + Destination path for the HDF5 file. + allow_file_overwrite : bool, optional + Whether to overwrite an existing file, by default False. + include_image_stack : bool, optional + Write ``image_stack`` to the HDF5 file, by default False. + Raises ``ValueError`` if the image stack has not been loaded. + include_local_stats : bool, optional + Write per-particle local stats to the HDF5 file, by default False. + Raises ``ValueError`` if the local stats have not been set. + + Returns + ------- + ParticleStackHDF5 + The new HDF5-backed stack instance pointing at ``hdf5_path``. + """ + # Generate particle_id for each row and add to the copied DataFrame + df = self._df.copy() + particle_ids = _generate_particle_ids(df) + df.insert(0, "particle_id", particle_ids) + df = df.set_index("particle_id") + df.index.name = "particle_id" + + hdf5_stack = ParticleStackHDF5( + hdf5_path=hdf5_path, + allow_file_overwrite=allow_file_overwrite, + extracted_box_size=self.extracted_box_size, + original_template_size=self.original_template_size, + leopard_em_version=self.leopard_em_version, + global_whitening_applied=self.global_whitening_applied, + local_whitening_applied=self.local_whitening_applied, + global_normalization_applied=self.global_normalization_applied, + local_normalization_applied=self.local_normalization_applied, + image_stack=self.image_stack if include_image_stack else None, + local_stats_correlation_average=( + self.local_stats_correlation_average if include_local_stats else None + ), + local_stats_correlation_variance=( + self.local_stats_correlation_variance if include_local_stats else None + ), + skip_df_load=True, + ) + hdf5_stack._df = df + hdf5_stack.to_hdf5( + include_image_stack=include_image_stack, + include_local_stats=include_local_stats, + ) + return hdf5_stack + + +# --------------------------------------------------------------------------- +# HDF5-backed subclass +# --------------------------------------------------------------------------- + + +class ParticleStackHDF5(_ParticleStackBase): + """Particle stack stored entirely within a single HDF5 file. + + The particle table, optional image stack, and optional per-particle local + correlation statistics are all held in one ``.h5`` file. Two loading + modes are supported — choose one; mixing them raises errors: + + * **Load from referenced files**: ``image_stack`` and ``local_stats`` are + computed from the paths stored in the particle table. The HDF5 file + stores only the particle table (``image_stack_stored=False``). + * **Load from HDF5**: ``image_stack`` and ``local_stats`` are read + directly from the HDF5 datasets (``image_stack_stored=True`` and/or + ``local_stats_stored=True``). + + HDF5 file layout + ---------------- + + :: + + / (root) + │ attrs: leopard_em_version, extracted_box_size, original_template_size, + │ image_stack_stored, local_stats_stored, + │ global_whitening_applied, local_whitening_applied, + │ global_normalization_applied, local_normalization_applied + ├─ particles/ + │ particle_id (N,) variable-length str "{mic_stem}_{idx:05d}" + │ (N,) float64 or variable-length str + │ ... + ├─ image_stack (N, box_h, box_w) float32 [optional] + └─ local_stats/ [optional] + correlation_average (N, valid_h, valid_w) float32 + correlation_variance (N, valid_h, valid_w) float32 + + where ``valid_h = extracted_box_size[0] - original_template_size[0] + 1`` + and ``valid_w = extracted_box_size[1] - original_template_size[1] + 1``. + + Attributes + ---------- + hdf5_path : str + Path to the HDF5 file. + allow_file_overwrite : bool + Whether to permit overwriting an existing file, by default False. + image_stack_stored : bool + True when ``/image_stack`` is present in the HDF5 file. + local_stats_stored : bool + True when ``/local_stats`` group is present in the HDF5 file. + """ + + hdf5_path: str + allow_file_overwrite: bool = False + image_stack_stored: bool = False + local_stats_stored: bool = False + + ########################### + ### Pydantic Validators ### + ########################### + + @model_validator(mode="after") # type: ignore + def _validate_hdf5_path(self) -> Self: + """Validate that the HDF5 path is writable and the overwrite policy is met. + + Returns + ------- + Self + + Raises + ------ + ValueError + If the path is not writable or the file exists and overwrite is + disabled. + """ + import os + + directory = str(Path(self.hdf5_path).parent) + if directory and not os.path.exists(directory): + os.makedirs(directory, exist_ok=True) + if directory and not os.access(directory, os.W_OK): + raise ValueError( + f"Directory '{directory}' does not permit writing to " + f"'{self.hdf5_path}'." + ) + if not self.allow_file_overwrite and os.path.exists(self.hdf5_path): + raise ValueError( + f"File '{self.hdf5_path}' already exists but " + "'allow_file_overwrite' is False." + ) + return self + + ########################### + ### Data loading ### + ########################### + + def load_df(self) -> None: + """Load the particle DataFrame from the HDF5 file at ``hdf5_path``. + + Raises + ------ + FileNotFoundError + If ``hdf5_path`` does not exist. + """ + import os + + if not os.path.exists(self.hdf5_path): + raise FileNotFoundError( + f"HDF5 file '{self.hdf5_path}' does not exist. " + "Pass skip_df_load=True if you intend to write a new file." + ) + with h5py.File(self.hdf5_path, "r") as f: + self._df = _read_df_from_hdf5_group(f) + + ########################### + ### I/O methods ### + ########################### + + def to_hdf5( + self, + include_image_stack: bool = False, + include_local_stats: bool = False, + ) -> None: + """Write the particle table and optional tensors to ``hdf5_path``. + + Parameters + ---------- + include_image_stack : bool, optional + Write ``image_stack`` to ``/image_stack``, by default False. + Raises ``ValueError`` if ``image_stack`` is None. + include_local_stats : bool, optional + Write per-particle correlation stats to ``/local_stats``, by + default False. Raises ``ValueError`` if either stats tensor is + None. + """ + with h5py.File(self.hdf5_path, "w") as f: + # Root attributes — metadata + f.attrs["leopard_em_version"] = self.leopard_em_version + f.attrs["extracted_box_size"] = list(self.extracted_box_size) + f.attrs["original_template_size"] = list(self.original_template_size) + f.attrs["global_whitening_applied"] = self.global_whitening_applied + f.attrs["local_whitening_applied"] = self.local_whitening_applied + f.attrs["global_normalization_applied"] = self.global_normalization_applied + f.attrs["local_normalization_applied"] = self.local_normalization_applied + + # Particle table + _write_df_to_hdf5_group(f, self._df) + + # Optional image stack + if include_image_stack: + if self.image_stack is None: + raise ValueError( + "image_stack is None; cannot write to HDF5. " + "Call construct_image_stack() first." + ) + f.create_dataset( + _HDF5_IMAGE_STACK_DATASET, + data=self.image_stack.cpu().to(torch.float32).numpy(), + ) + self.image_stack_stored = True + + f.attrs["image_stack_stored"] = self.image_stack_stored + + # Optional per-particle local stats + if include_local_stats: + if ( + self.local_stats_correlation_average is None + or self.local_stats_correlation_variance is None + ): + raise ValueError( + "local_stats tensors are None; cannot write to HDF5." + ) + local_grp = f.create_group(_HDF5_LOCAL_STATS_GROUP) + local_grp.create_dataset( + "correlation_average", + data=self.local_stats_correlation_average.cpu() + .to(torch.float32) + .numpy(), + ) + local_grp.create_dataset( + "correlation_variance", + data=self.local_stats_correlation_variance.cpu() + .to(torch.float32) + .numpy(), + ) + self.local_stats_stored = True + + f.attrs["local_stats_stored"] = self.local_stats_stored + + @classmethod + def from_hdf5( + cls, + path: str, + allow_file_overwrite: bool = True, + ) -> "ParticleStackHDF5": + """Load a ``ParticleStackHDF5`` from an existing HDF5 file. + + Parameters + ---------- + path : str + Path to the HDF5 file written by ``to_hdf5``. + allow_file_overwrite : bool, optional + Passed to the constructor so that the model validator does not + reject the path of the file being loaded, by default True. + + Returns + ------- + ParticleStackHDF5 + """ + with h5py.File(path, "r") as f: + leopard_em_version = str(f.attrs.get("leopard_em_version", "unknown")) + extracted_box_size = tuple(int(v) for v in f.attrs["extracted_box_size"]) + original_template_size = tuple( + int(v) for v in f.attrs["original_template_size"] + ) + global_whitening_applied = bool( + f.attrs.get("global_whitening_applied", False) + ) + local_whitening_applied = bool( + f.attrs.get("local_whitening_applied", False) + ) + global_normalization_applied = bool( + f.attrs.get("global_normalization_applied", False) + ) + local_normalization_applied = bool( + f.attrs.get("local_normalization_applied", False) + ) + image_stack_stored = bool(f.attrs.get("image_stack_stored", False)) + local_stats_stored = bool(f.attrs.get("local_stats_stored", False)) + + df = _read_df_from_hdf5_group(f) + + image_stack: torch.Tensor | None = None + if image_stack_stored: + if _HDF5_IMAGE_STACK_DATASET not in f: + raise ValueError( + f"'image_stack_stored' is True but dataset " + f"'{_HDF5_IMAGE_STACK_DATASET}' is absent in '{path}'." + ) + image_stack = torch.from_numpy(f[_HDF5_IMAGE_STACK_DATASET][:]) + + local_avg: torch.Tensor | None = None + local_var: torch.Tensor | None = None + if local_stats_stored: + if _HDF5_LOCAL_STATS_GROUP not in f: + raise ValueError( + f"'local_stats_stored' is True but group " + f"'{_HDF5_LOCAL_STATS_GROUP}' is absent in '{path}'." + ) + local_grp = f[_HDF5_LOCAL_STATS_GROUP] + local_avg = torch.from_numpy(local_grp["correlation_average"][:]) + local_var = torch.from_numpy(local_grp["correlation_variance"][:]) + + instance = cls( + hdf5_path=str(path), + allow_file_overwrite=allow_file_overwrite, + extracted_box_size=extracted_box_size, + original_template_size=original_template_size, + leopard_em_version=leopard_em_version, + global_whitening_applied=global_whitening_applied, + local_whitening_applied=local_whitening_applied, + global_normalization_applied=global_normalization_applied, + local_normalization_applied=local_normalization_applied, + image_stack_stored=image_stack_stored, + local_stats_stored=local_stats_stored, + image_stack=image_stack, + local_stats_correlation_average=local_avg, + local_stats_correlation_variance=local_var, + skip_df_load=True, + ) + instance._df = df + return instance + + +# --------------------------------------------------------------------------- +# Backward-compatibility alias +# --------------------------------------------------------------------------- + +# Existing code that imports `ParticleStack` continues to receive +# `ParticleStackCSV` unchanged. +ParticleStack = ParticleStackCSV diff --git a/src/leopard_em/pydantic_models/managers/match_template_manager.py b/src/leopard_em/pydantic_models/managers/match_template_manager.py index 7c239a57..0bce9820 100644 --- a/src/leopard_em/pydantic_models/managers/match_template_manager.py +++ b/src/leopard_em/pydantic_models/managers/match_template_manager.py @@ -23,7 +23,11 @@ from leopard_em.pydantic_models.custom_types import BaseModel2DTM, ExcludedTensor from leopard_em.pydantic_models.data_structures import OpticsGroup from leopard_em.pydantic_models.formats import MATCH_TEMPLATE_DF_COLUMN_ORDER -from leopard_em.pydantic_models.results import MatchTemplateResult +from leopard_em.pydantic_models.results import ( + MatchTemplateResultHDF5, + MatchTemplateResultMRC, +) +from leopard_em.pydantic_models.results.correlation_table import CorrelationTable from leopard_em.utils.ctf_utils import calculate_ctf_filter_stack from leopard_em.utils.data_io import load_mrc_image, load_mrc_volume from leopard_em.utils.image_processing import ( @@ -55,9 +59,10 @@ class MatchTemplateManager(BaseModel2DTM): preprocessing_filters : PreprocessingFilters Configurations for the preprocessing filters to apply during correlation. - match_template_result : MatchTemplateResult - Result of the match template program stored as an instance of the - `MatchTemplateResult` class. + match_template_result : MatchTemplateResultMRC | MatchTemplateResultHDF5 + Result of the match template program. Use ``MatchTemplateResultMRC`` + to write individual MRC files or ``MatchTemplateResultHDF5`` to bundle + all tensors into a single HDF5 file. computational_config : ComputationalConfigMatch Parameters for controlling computational resources. @@ -98,7 +103,7 @@ class MatchTemplateManager(BaseModel2DTM): defocus_search_config: DefocusSearchConfig orientation_search_config: OrientationSearchConfig | MultipleOrientationConfig preprocessing_filters: PreprocessingFilters - match_template_result: MatchTemplateResult + match_template_result: MatchTemplateResultMRC | MatchTemplateResultHDF5 computational_config: ComputationalConfigMatch # Non-serialized large array-like attributes @@ -221,7 +226,7 @@ def run_match_template( self, orientation_batch_size: int = 16, do_result_export: bool = True, - do_valid_cropping: bool = True, + do_valid_cropping: bool = False, ) -> None: """Runs the base match template in pytorch. @@ -233,7 +238,10 @@ def run_match_template( If True, call the `MatchTemplateResult.export_results` method to save the results to disk directly after running the match template. Default is True. do_valid_cropping : bool - If True, apply the valid cropping mode to the results. Default is True. + If True, then apply valid cropping to the result maps based on the relative + size of the image and template (N-n+1 along each axis). The backend of + Leopard-EM will automatically do this, so generally set this to False. The + default is False. Returns ------- @@ -250,6 +258,8 @@ def run_match_template( # Populate the MatchTemplateResult via a private helper self._populate_match_template_result( results, + defocus_values=core_kwargs["defocus_values"], + euler_angles=core_kwargs["euler_angles"], do_result_export=do_result_export, do_valid_cropping=do_valid_cropping, ) @@ -261,7 +271,7 @@ def run_match_template_distributed( local_rank: int, orientation_batch_size: int = 16, do_result_export: bool = True, - do_valid_cropping: bool = True, + do_valid_cropping: bool = False, ) -> None: """Runs the base match template in a distributed, multi-node environment. @@ -279,7 +289,10 @@ def run_match_template_distributed( If True, call the `MatchTemplateResult.export_results` method to save the results to disk directly after running the match template. Default is True. do_valid_cropping : bool - If True, apply the valid cropping mode to the results. Default is True. + If True, then apply valid cropping to the result maps based on the relative + size of the image and template (N-n+1 along each axis). The backend of + Leopard-EM will automatically do this, so generally set this to False. The + default is False. Raises ------ @@ -320,6 +333,8 @@ def run_match_template_distributed( if torch.distributed.get_rank() == 0: self._populate_match_template_result( results, + defocus_values=core_kwargs["defocus_values"], + euler_angles=core_kwargs["euler_angles"], do_result_export=do_result_export, do_valid_cropping=do_valid_cropping, ) @@ -327,8 +342,10 @@ def run_match_template_distributed( def _populate_match_template_result( self, results: dict[str, Any], + defocus_values: torch.Tensor, + euler_angles: torch.Tensor, do_result_export: bool = True, - do_valid_cropping: bool = True, + do_valid_cropping: bool = False, ) -> None: """Helper function to populate the MatchTemplateResult object post-core call.""" # Place results into the `MatchTemplateResult` object @@ -348,7 +365,20 @@ def _populate_match_template_result( self.match_template_result.total_orientations = results["total_orientations"] self.match_template_result.total_defocus = results["total_defocus"] + # Build a typed CorrelationTable from the processed backend output, looking up + # per-detection mean/variance from the statistics tensors independently. + self.match_template_result.correlation_table = ( + CorrelationTable.from_match_template_results( + processed_correlation_table=results["correlation_table"], + defocus_values=defocus_values, + euler_angles=euler_angles, + correlation_average=results["correlation_mean"], + correlation_variance_map=results["correlation_variance"], + ) + ) + # Apply the valid cropping mode to the results + # NOTE: zipFFT already applies valid cropping internally if do_valid_cropping: nx = self.template_volume.shape[-1] self.match_template_result.apply_valid_cropping((nx, nx)) @@ -453,19 +483,30 @@ def results_to_dataframe( df["micrograph_path"] = self.micrograph_path df["template_path"] = self.template_volume_path - # Add paths to the output statistic files - df["mip_path"] = self.match_template_result.mip_path - df["scaled_mip_path"] = self.match_template_result.scaled_mip_path - df["psi_path"] = self.match_template_result.orientation_psi_path - df["theta_path"] = self.match_template_result.orientation_theta_path - df["phi_path"] = self.match_template_result.orientation_phi_path - df["defocus_path"] = self.match_template_result.relative_defocus_path - df["correlation_average_path"] = ( - self.match_template_result.correlation_average_path - ) - df["correlation_variance_path"] = ( - self.match_template_result.correlation_variance_path - ) + # Add paths to the output statistic files, branching on storage back-end + if isinstance(self.match_template_result, MatchTemplateResultMRC): + df["mip_path"] = self.match_template_result.mip_path + df["scaled_mip_path"] = self.match_template_result.scaled_mip_path + df["psi_path"] = self.match_template_result.orientation_psi_path + df["theta_path"] = self.match_template_result.orientation_theta_path + df["phi_path"] = self.match_template_result.orientation_phi_path + df["defocus_path"] = self.match_template_result.relative_defocus_path + df["correlation_average_path"] = ( + self.match_template_result.correlation_average_path + ) + df["correlation_variance_path"] = ( + self.match_template_result.correlation_variance_path + ) + else: + # HDF5: all tensors are in one file; individual MRC paths are not applicable + df["mip_path"] = self.match_template_result.hdf5_path + df["scaled_mip_path"] = self.match_template_result.hdf5_path + df["psi_path"] = self.match_template_result.hdf5_path + df["theta_path"] = self.match_template_result.hdf5_path + df["phi_path"] = self.match_template_result.hdf5_path + df["defocus_path"] = self.match_template_result.hdf5_path + df["correlation_average_path"] = self.match_template_result.hdf5_path + df["correlation_variance_path"] = self.match_template_result.hdf5_path # Add particle index df["particle_index"] = df.index diff --git a/src/leopard_em/pydantic_models/results/__init__.py b/src/leopard_em/pydantic_models/results/__init__.py index 9f2714d0..66a88b05 100644 --- a/src/leopard_em/pydantic_models/results/__init__.py +++ b/src/leopard_em/pydantic_models/results/__init__.py @@ -1,7 +1,15 @@ """Pydantic models for Leopard-EM program results.""" -from .match_template_result import MatchTemplateResult +from .correlation_table import CorrelationTable +from .match_template_result import ( + MatchTemplateResult, + MatchTemplateResultHDF5, + MatchTemplateResultMRC, +) __all__ = [ + "CorrelationTable", "MatchTemplateResult", + "MatchTemplateResultHDF5", + "MatchTemplateResultMRC", ] diff --git a/src/leopard_em/pydantic_models/results/correlation_table.py b/src/leopard_em/pydantic_models/results/correlation_table.py new file mode 100644 index 00000000..1d2e69a0 --- /dev/null +++ b/src/leopard_em/pydantic_models/results/correlation_table.py @@ -0,0 +1,341 @@ +"""Storage of sparse particle detections in a 2DTM search.""" + +import h5py +import numpy as np +import pandas as pd +import torch + +from leopard_em.pydantic_models.custom_types import BaseModel2DTM + + +def derive_orientation_grid_from_full_angles( + euler_angles: torch.Tensor, +) -> tuple[list[tuple[float, float]], list[float]]: + """Extract unique (phi, theta) pairs and psi values from a grid angles tensor. + + Assumes ``euler_angles`` is ordered as a Cartesian product: all psi values for + the first (phi, theta) pair, then all psi values for the second pair, etc. + + Parameters + ---------- + euler_angles : torch.Tensor + All Euler angles used in the search, shape (num_orientations, 3), in ZYZ + convention (degrees). + + Returns + ------- + tuple[list[tuple[float, float]], list[float]] + - ``phi_theta_angles``: list of unique (phi, theta) pairs, one per out-of-plane + orientation, in the order they appear in the search. + - ``psi_angles``: list of unique psi values, in the order they cycle within each + (phi, theta) group. + """ + n_orientations = euler_angles.shape[0] + n_psi = int(torch.unique(euler_angles[:, 2]).shape[0]) + n_phi_theta = n_orientations // n_psi + + phi_theta_angles = [ + (float(euler_angles[i * n_psi, 0]), float(euler_angles[i * n_psi, 1])) + for i in range(n_phi_theta) + ] + psi_angles = euler_angles[:n_psi, 2].tolist() + + return phi_theta_angles, psi_angles + + +class CorrelationTable(BaseModel2DTM): + """Correlation table data structure storing possible detections along a 2DTM search. + + Attributes + ---------- + correlation_threshold : float + Pre-defined threshold a cross-correlation value must surpass to be included + in the correlation table. + num_observations : int + Total number of detections in the correlation table (number of search indices + which surpassed the correlation threshold). + defocus_offsets : list[float] + List of defocus offsets (in Angstroms) used in the search. + phi_theta_angles : list[tuple[float, float]] + List out-of-plane rotation angles (in degrees, Euler angles phi and theta, in + ZYZ convention) used in the search. + psi_angles : list[float] + List of in-plane rotation angles (in degrees, Euler angle psi, in ZYZ + convention) used in the search. + search_index : list[int] + Global search index defining defocus offset, phi/theta angles, and psi angle for + each detection. Calculated as `i * (n_j * n_k) + j * n_k + k`, where `i` is the + index of the defocus offset, `j` is the index of the phi/theta angles, and `k` + is the index of the psi angle. Length will be equal to `num_observations`. + x : list[int] + List of x-coordinates (in pixels) of the detections in the micrograph. + y : list[int] + List of y-coordinates (in pixels) of the detections in the micrograph. + correlation_value : list[float] + List of cross-correlation values for each detection. + correlation_mean : list[float] + List of mean cross-correlation values for each detection, calculated across all + search indices for the same x/y coordinates. + correlation_variance : list[float] + List of variance of cross-correlation values for each detection, calculated + across all search indices for the same x/y coordinates. + + Methods + ------- + to_dataframe() -> pd.DataFrame + from_dataframe(df: pd.DataFrame) -> CorrelationTable + to_hdf5(file_path: str) + from_hdf5(file_path: str) -> CorrelationTable + from_match_template_results(...) -> CorrelationTable + """ + + correlation_threshold: float + num_observations: int + + # Defining and indexing search space + defocus_offsets: list[float] # index 'i' + phi_theta_angles: list[tuple[float, float]] # index 'j', out-of-plane rotations + psi_angles: list[float] # index 'k', in-plane rotations + search_index: list[int] # i * (n_j * n_k) + j * n_k + k, length == num_observations + + # Other detection attributes + x: list[int] + y: list[int] + correlation_value: list[float] + correlation_mean: list[float] + correlation_variance: list[float] + + def to_dataframe(self) -> pd.DataFrame: + """Convert per-detection data to a DataFrame. + + Search-space metadata is stored in ``df.attrs`` so that + ``from_dataframe`` can reconstruct the full object. + + Returns + ------- + pd.DataFrame + One row per detection with columns: search_index, x, y, + correlation_value, correlation_mean, correlation_variance. + """ + df = pd.DataFrame( + { + "search_index": self.search_index, + "x": self.x, + "y": self.y, + "correlation_value": self.correlation_value, + "correlation_mean": self.correlation_mean, + "correlation_variance": self.correlation_variance, + } + ) + df.attrs["correlation_threshold"] = self.correlation_threshold + df.attrs["num_observations"] = self.num_observations + df.attrs["defocus_offsets"] = self.defocus_offsets + df.attrs["phi_theta_angles"] = self.phi_theta_angles + df.attrs["psi_angles"] = self.psi_angles + return df + + @classmethod + def from_dataframe(cls, df: pd.DataFrame) -> "CorrelationTable": + """Reconstruct a CorrelationTable from a DataFrame produced by ``to_dataframe``. + + Parameters + ---------- + df : pd.DataFrame + DataFrame with detection columns and search-space metadata in + ``df.attrs``. + + Returns + ------- + CorrelationTable + """ + return cls( + correlation_threshold=float(df.attrs["correlation_threshold"]), + num_observations=int(df.attrs["num_observations"]), + defocus_offsets=list(df.attrs["defocus_offsets"]), + phi_theta_angles=[tuple(pair) for pair in df.attrs["phi_theta_angles"]], + psi_angles=list(df.attrs["psi_angles"]), + search_index=df["search_index"].tolist(), + x=df["x"].tolist(), + y=df["y"].tolist(), + correlation_value=df["correlation_value"].tolist(), + correlation_mean=df["correlation_mean"].tolist(), + correlation_variance=df["correlation_variance"].tolist(), + ) + + def to_hdf5(self, file_path: str) -> None: + """Write this CorrelationTable to an HDF5 file. + + Layout:: + + /metadata (attrs: correlation_threshold, num_observations) + /search_space/ + defocus_offsets float32 1-D + phi_theta_angles float32 (n, 2) + psi_angles float32 1-D + /detections/ + search_index int32 1-D + x int32 1-D + y int32 1-D + correlation_value float32 1-D + correlation_mean float32 1-D + correlation_variance float32 1-D + + Parameters + ---------- + file_path : str + Destination HDF5 file path. + """ + with h5py.File(file_path, "w") as f: + meta = f.create_group("metadata") + meta.attrs["correlation_threshold"] = self.correlation_threshold + meta.attrs["num_observations"] = self.num_observations + + ss = f.create_group("search_space") + ss.create_dataset( + "defocus_offsets", + data=np.array(self.defocus_offsets, dtype=np.float32), + ) + ss.create_dataset( + "phi_theta_angles", + data=np.array(self.phi_theta_angles, dtype=np.float32), + ) + ss.create_dataset( + "psi_angles", + data=np.array(self.psi_angles, dtype=np.float32), + ) + + det = f.create_group("detections") + det.create_dataset( + "search_index", + data=np.array(self.search_index, dtype=np.int32), + ) + det.create_dataset("x", data=np.array(self.x, dtype=np.int32)) + det.create_dataset("y", data=np.array(self.y, dtype=np.int32)) + det.create_dataset( + "correlation_value", + data=np.array(self.correlation_value, dtype=np.float32), + ) + det.create_dataset( + "correlation_mean", + data=np.array(self.correlation_mean, dtype=np.float32), + ) + det.create_dataset( + "correlation_variance", + data=np.array(self.correlation_variance, dtype=np.float32), + ) + + @classmethod + def from_hdf5(cls, file_path: str) -> "CorrelationTable": + """Load a CorrelationTable from an HDF5 file written by ``to_hdf5``. + + Parameters + ---------- + file_path : str + Path to the HDF5 file. + + Returns + ------- + CorrelationTable + """ + with h5py.File(file_path, "r") as f: + correlation_threshold = float(f["metadata"].attrs["correlation_threshold"]) + num_observations = int(f["metadata"].attrs["num_observations"]) + + defocus_offsets = f["search_space/defocus_offsets"][:].tolist() + phi_theta_raw = f["search_space/phi_theta_angles"][:] + phi_theta_angles = [(float(row[0]), float(row[1])) for row in phi_theta_raw] + psi_angles = f["search_space/psi_angles"][:].tolist() + + search_index = f["detections/search_index"][:].tolist() + x = f["detections/x"][:].tolist() + y = f["detections/y"][:].tolist() + correlation_value = f["detections/correlation_value"][:].tolist() + correlation_mean = f["detections/correlation_mean"][:].tolist() + correlation_variance = f["detections/correlation_variance"][:].tolist() + + return cls( + correlation_threshold=correlation_threshold, + num_observations=num_observations, + defocus_offsets=defocus_offsets, + phi_theta_angles=phi_theta_angles, + psi_angles=psi_angles, + search_index=search_index, + x=x, + y=y, + correlation_value=correlation_value, + correlation_mean=correlation_mean, + correlation_variance=correlation_variance, + ) + + @classmethod + def from_match_template_results( + cls, + processed_correlation_table: dict, + defocus_values: torch.Tensor, + euler_angles: torch.Tensor, + correlation_average: torch.Tensor, + correlation_variance_map: torch.Tensor, + ) -> "CorrelationTable": + """Construct a CorrelationTable from backend outputs. + + Parameters + ---------- + processed_correlation_table : dict + Output of ``process_correlation_table`` with an additional ``global_idx`` + key (list[int]). Expected keys: ``threshold``, ``global_idx``, ``x``, + ``y``, ``correlation``. + defocus_values : torch.Tensor + Defocus offsets used in the search. Shape (num_defocus,). + euler_angles : torch.Tensor + All Euler angles used in the search, shape (num_orientations, 3), in ZYZ + convention (degrees). Must be ordered as a grid: all psi values for the + first (phi, theta) pair, then all psi values for the second pair, etc. + correlation_average : torch.Tensor + Per-pixel mean cross-correlation, shape (H, W). + correlation_variance_map : torch.Tensor + Per-pixel standard deviation of cross-correlation, shape (H, W). + + Returns + ------- + CorrelationTable + """ + threshold = processed_correlation_table["threshold"] + global_idx = processed_correlation_table["global_idx"] # list[int] + pos_x = processed_correlation_table["x"] # list[int] + pos_y = processed_correlation_table["y"] # list[int] + corr_values = processed_correlation_table["correlation"] # list[float] + + defocus_offsets = defocus_values.tolist() + phi_theta_angles, psi_angles = derive_orientation_grid_from_full_angles( + euler_angles + ) + + # For match_template (num_cs == 1), the backend global_idx equals the + search_index = ( + list(global_idx) if isinstance(global_idx, list) else global_idx.tolist() + ) + + # Look up per-detection statistics from the pre-computed statistics tensors + num_observations = len(pos_x) + if num_observations > 0: + x_tensor = torch.tensor(pos_x, dtype=torch.long) + y_tensor = torch.tensor(pos_y, dtype=torch.long) + det_mean = correlation_average[y_tensor, x_tensor].tolist() + det_variance = correlation_variance_map[y_tensor, x_tensor].tolist() + else: + det_mean = [] + det_variance = [] + + return cls( + correlation_threshold=float(threshold), + num_observations=num_observations, + defocus_offsets=defocus_offsets, + phi_theta_angles=phi_theta_angles, + psi_angles=psi_angles, + search_index=search_index, + x=list(pos_x), + y=list(pos_y), + correlation_value=list(corr_values), + correlation_mean=det_mean, + correlation_variance=det_variance, + ) diff --git a/src/leopard_em/pydantic_models/results/match_template_result.py b/src/leopard_em/pydantic_models/results/match_template_result.py index 2c028fc3..6f83c784 100644 --- a/src/leopard_em/pydantic_models/results/match_template_result.py +++ b/src/leopard_em/pydantic_models/results/match_template_result.py @@ -1,13 +1,28 @@ -"""Reading, storing, and exporting results from the match_template program.""" +"""Reading, storing, and exporting results from the match_template program. + +Two public classes are provided for different storage back-ends: + +* ``MatchTemplateResultMRC`` - stores each result tensor in a separate MRC + file (the original behavior). ``MatchTemplateResult`` is an alias for + this class for backward compatibility. +* ``MatchTemplateResultHDF5`` - bundles every tensor and all scalar metadata + into a single HDF5 file. + +The base class ``_MatchTemplateResultBase`` holds the tensors, scalar +metadata, and all analysis methods and is not intended to be used directly. +""" # NOTE: Disabling pylint for too-many-instance-attributes since this class holds a # number of result attributes that are independent and should not be grouped further. # pylint: disable=too-many-instance-attributes import os +from importlib.metadata import PackageNotFoundError, version from typing import ClassVar +import h5py import pandas as pd +import torch from pydantic import ConfigDict, Field, model_validator from typing_extensions import Self @@ -18,9 +33,31 @@ match_template_peaks_to_dict, ) from leopard_em.pydantic_models.custom_types import BaseModel2DTM, ExcludedTensor +from leopard_em.pydantic_models.results.correlation_table import CorrelationTable from leopard_em.utils.data_io import load_mrc_image, write_mrc_from_tensor +def _leopard_em_version() -> str: + try: + return version("leopard_em") + except PackageNotFoundError: + return "uninstalled" + + +_TENSOR_NAMES = ( + "mip", + "scaled_mip", + "correlation_average", + "correlation_variance", + "orientation_psi", + "orientation_theta", + "orientation_phi", + "relative_defocus", +) + +_HDF5_TENSORS_GROUP = "tensors" + + def check_file_path_and_permissions(path: str, allow_overwrite: bool) -> None: """Ensures path is writable and it does not exist, if `allow_overwrite` is False.""" # 1. Create path to file, if it does not exist @@ -45,88 +82,40 @@ def check_file_path_and_permissions(path: str, allow_overwrite: bool) -> None: ) -class MatchTemplateResult(BaseModel2DTM): - """Class to hold and export results from the match_template program. +class _MatchTemplateResultBase(BaseModel2DTM): + """Base class holding result tensors, scalar metadata, and analysis methods. - TODO: Implement tracking of how far along the template matching is - (e.g. orientations up to what index have been searched). - TODO: Implement method for exporting intermediary results in case of error - or program interruption. - TODO: Implement functionality for restarting template matching from a - saved state (e.g. after a program interruption). + Not intended to be instantiated directly — use ``MatchTemplateResultMRC`` + or ``MatchTemplateResultHDF5`` depending on the desired storage back-end. Attributes ---------- - allow_file_overwrite : bool = False - Weather to allow overwriting of existing files. Default is False. - WARNING: Setting to True can overwrite existing files! - mip_path : str - Path to the output maximum intensity projection (MIP) file. - scaled_mip_path : str - Path to the output scaled MIP file. - correlation_average_path : str - Path to the output correlation average file. - correlation_variance_path : str - Path to the output correlation variance file. - orientation_psi_path : str - Path to the output orientation psi file. - orientation_theta_path : str - Path to the output orientation theta file. - orientation_phi_path : str - Path to the output orientation phi file. - relative_defocus_path : str - Path to the output relative defocus file. + leopard_em_version : str + Version of Leopard-EM that produced this result. Auto-populated from + the installed package metadata on construction; preserved as-recorded + when loading from a file. + total_projections : int + Total cross-correlations computed (orientations x defocus steps). + total_orientations : int + Total orientations searched. + total_defocus : int + Total defocus values searched. mip : ExcludedTensor - Maximum intensity projection (MIP). + Maximum intensity projection. scaled_mip : ExcludedTensor - Scaled MIP. + Scaled MIP (z-score normalized). correlation_average : ExcludedTensor - Correlation average. + Running mean of the correlation over the search space, per pixel. correlation_variance : ExcludedTensor - Correlation variance. + Running variance of the correlation over the search space, per pixel. orientation_psi : ExcludedTensor - Best orientation angle psi. + Per-pixel best psi angle (degrees, ZYZ convention). orientation_theta : ExcludedTensor - Best orientation angle theta. + Per-pixel best theta angle (degrees, ZYZ convention). orientation_phi : ExcludedTensor - Best orientation angle phi. + Per-pixel best phi angle (degrees, ZYZ convention). relative_defocus : ExcludedTensor - Best relative defocus. - total_projections : int, optional - Total number of cross-correlograms of projections computed. Should be - 'total_orientations x total_defocus' Default is 0, and this field is updated - automatically after a match_template run. - total_orientations : int, optional - Total number of orientations searched. Default is 0, and this field is updated - automatically after a match_template run. - total_defocus : int, optional - Total number of defocus values searched. Default is 0, and this field is updated - automatically after a match_template run. - match_template_peaks : MatchTemplatePeaks - Named tuple object containing the peak locations, heights, and pose statistics. - See the 'analysis.pick_match_template_peaks' module for more information. - - Methods - ------- - validate_paths() - Validates the output paths for write permissions and overwriting. - - load_tensors_from_paths() - Load tensors from the specified (held) paths into memory. - - locate_peaks(**kwargs) - Updates the 'match_template_peaks' attribute with info from held tensors. - Additional keyword arguments can be passed to the 'extract_peaks_and_statistics' - function. - - peaks_to_dict() - Convert the 'match_template_peaks' attribute to a dictionary. - - peaks_to_dataframe() - Convert the 'match_template_peaks' attribute to a pandas DataFrame. - - export_results() - Export the torch.Tensor results to the specified mrc files. + Per-pixel best relative defocus offset (Angstroms). """ model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True) @@ -136,23 +125,17 @@ class MatchTemplateResult(BaseModel2DTM): # it will lead to headaches when attempting to load a result, this is set # to True, and the result files already exist. allow_file_overwrite: bool = False - mip_path: str - scaled_mip_path: str - correlation_average_path: str - correlation_variance_path: str - orientation_psi_path: str - orientation_theta_path: str - orientation_phi_path: str - relative_defocus_path: str + correlation_table_path: str | None = Field(default=None) # Scalar (non-tensor) attributes + leopard_em_version: str = Field(default_factory=_leopard_em_version) total_projections: int = 0 total_orientations: int = 0 total_defocus: int = 0 match_template_peaks: MatchTemplatePeaks = Field(default=None, exclude=True) + correlation_table: CorrelationTable | None = Field(default=None, exclude=True) - # Large array-like attributes saved to individual files (not in JSON) mip: ExcludedTensor scaled_mip: ExcludedTensor correlation_average: ExcludedTensor @@ -162,43 +145,6 @@ class MatchTemplateResult(BaseModel2DTM): orientation_phi: ExcludedTensor relative_defocus: ExcludedTensor - ########################### - ### Pydantic Validators ### - ########################### - - @model_validator(mode="after") # type: ignore - def validate_paths(self) -> Self: - """Validate output paths for write permissions and overwriting. - - Note: This method runs after instantiation, so attributes are already - set. We can safely access them with `self`. - - Returns - ------- - Self - The validated instance. - - Raises - ------ - ValueError - If the output paths are not writable or do not permit overwriting. - """ - # 1. Check write permissions and overwriting for each path - paths = [ - self.mip_path, - self.scaled_mip_path, - self.correlation_average_path, - self.correlation_variance_path, - self.orientation_psi_path, - self.orientation_theta_path, - self.orientation_phi_path, - self.relative_defocus_path, - ] - for path in paths: - check_file_path_and_permissions(path, self.allow_file_overwrite) - - return self - ############################################ ### Functional (data processing) methods ### ############################################ @@ -221,7 +167,9 @@ def apply_valid_cropping(self, template_shape: tuple[int, int]) -> None: ------- None """ - # Assuming all statistic files have the same shape (which should be true!) + # NOTE: Assuming all statistic files have the same shape (which should be true) + # NOTE: Assuming we the correlation maps have not already been cropped, which + # is not true for the zipFFT backend. img_h, img_w = self.mip.shape h, w = template_shape slice_obj = (slice(img_h - h + 1), slice(img_w - w + 1)) @@ -235,42 +183,17 @@ def apply_valid_cropping(self, template_shape: tuple[int, int]) -> None: self.orientation_phi = self.orientation_phi[slice_obj] self.relative_defocus = self.relative_defocus[slice_obj] - def load_tensors_from_paths(self) -> None: - """Use the held paths to load tensors into memory. - - NOTE: Currently only supports .mrc files. - """ - self.mip = load_mrc_image(self.mip_path) - self.scaled_mip = load_mrc_image(self.scaled_mip_path) - self.correlation_average = load_mrc_image(self.correlation_average_path) - self.correlation_variance = load_mrc_image(self.correlation_variance_path) - self.orientation_psi = load_mrc_image(self.orientation_psi_path) - self.orientation_theta = load_mrc_image(self.orientation_theta_path) - self.orientation_phi = load_mrc_image(self.orientation_phi_path) - self.relative_defocus = load_mrc_image(self.relative_defocus_path) - def locate_peaks(self, **kwargs) -> MatchTemplatePeaks: # type: ignore - """Updates the 'match_template_peaks' attribute with info from held tensors. - - This method calls the `extract_peaks_and_statistics` function to first locate - particles based on the z-scores of the correlation results, then finds the - best orientations and defocus values at those locations. Returned named tuple - object is stored in the 'match_template_peaks' attribute. - - NOTE: Method intended to be called after running match_template or loading - the tensors from disk. + """Locate peaks and store results in ``match_template_peaks``. Parameters ---------- **kwargs - Additional keyword arguments to pass to the 'extract_peaks_and_statistics' - function. + Forwarded to ``extract_peaks_and_statistics_zscore``. Returns ------- MatchTemplatePeaks - Named tuple object containing the peak locations, heights, and pose - statistics. """ self.match_template_peaks = extract_peaks_and_statistics_zscore( mip=self.mip, @@ -284,29 +207,83 @@ def locate_peaks(self, **kwargs) -> MatchTemplatePeaks: # type: ignore total_correlation_positions=self.total_projections, **kwargs, ) - return self.match_template_peaks def peaks_to_dict(self) -> dict: - """Convert the 'match_template_peaks' attribute to a dictionary.""" + """Convert ``match_template_peaks`` to a dictionary.""" if self.match_template_peaks is None: self.locate_peaks() - return match_template_peaks_to_dict(self.match_template_peaks) def peaks_to_dataframe(self) -> pd.DataFrame: - """Convert the 'match_template_peaks' attribute to a pandas DataFrame.""" + """Convert ``match_template_peaks`` to a pandas DataFrame.""" if self.match_template_peaks is None: self.locate_peaks() - return match_template_peaks_to_dataframe(self.match_template_peaks) - ###################### - ### Export methods ### - ###################### - def export_results(self) -> None: - """Export the torch.Tensor results to the specified mrc files.""" +class MatchTemplateResultMRC(_MatchTemplateResultBase): + """Stores each result tensor in a separate MRC file. + + Attributes + ---------- + allow_file_overwrite : bool + Whether to allow overwriting of existing files. Default is False. + mip_path : str + Output path for the maximum intensity projection MRC file. + scaled_mip_path : str + Output path for the scaled MIP MRC file. + correlation_average_path : str + Output path for the correlation average MRC file. + correlation_variance_path : str + Output path for the correlation variance MRC file. + orientation_psi_path : str + Output path for the orientation psi MRC file. + orientation_theta_path : str + Output path for the orientation theta MRC file. + orientation_phi_path : str + Output path for the orientation phi MRC file. + relative_defocus_path : str + Output path for the relative defocus MRC file. + + Methods + ------- + validate_paths() + Validates write permissions and overwrite policy for all eight paths. + load_tensors_from_paths() + Reads MRC files from the held paths into memory. + export_results() + Writes the held tensors to their respective MRC paths. + """ + + allow_file_overwrite: bool = False + mip_path: str + scaled_mip_path: str + correlation_average_path: str + correlation_variance_path: str + orientation_psi_path: str + orientation_theta_path: str + orientation_phi_path: str + relative_defocus_path: str + + ########################### + ### Pydantic Validators ### + ########################### + + @model_validator(mode="after") # type: ignore + def validate_paths(self) -> Self: + """Validate output paths for write permissions and overwriting. + + Returns + ------- + Self + + Raises + ------ + ValueError + If any path is not writable or already exists and overwriting is + disabled. + """ paths = [ self.mip_path, self.scaled_mip_path, @@ -317,21 +294,217 @@ def export_results(self) -> None: self.orientation_phi_path, self.relative_defocus_path, ] - tensors = [ - self.mip, - self.scaled_mip, - self.correlation_average, - self.correlation_variance, - self.orientation_psi, - self.orientation_theta, - self.orientation_phi, - self.relative_defocus, - ] + for path in paths: + check_file_path_and_permissions(path, self.allow_file_overwrite) + return self - for path, tensor in zip(paths, tensors): + ###################### + ### I/O methods ### + ###################### + + def load_tensors_from_paths(self) -> None: + """Read MRC files from the held paths into the tensor attributes.""" + self.mip = load_mrc_image(self.mip_path) + self.scaled_mip = load_mrc_image(self.scaled_mip_path) + self.correlation_average = load_mrc_image(self.correlation_average_path) + self.correlation_variance = load_mrc_image(self.correlation_variance_path) + self.orientation_psi = load_mrc_image(self.orientation_psi_path) + self.orientation_theta = load_mrc_image(self.orientation_theta_path) + self.orientation_phi = load_mrc_image(self.orientation_phi_path) + self.relative_defocus = load_mrc_image(self.relative_defocus_path) + + def export_results(self) -> None: + """Write the held tensors to their respective MRC paths.""" + pairs = [ + (self.mip_path, self.mip), + (self.scaled_mip_path, self.scaled_mip), + (self.correlation_average_path, self.correlation_average), + (self.correlation_variance_path, self.correlation_variance), + (self.orientation_psi_path, self.orientation_psi), + (self.orientation_theta_path, self.orientation_theta), + (self.orientation_phi_path, self.orientation_phi), + (self.relative_defocus_path, self.relative_defocus), + ] + for path, tensor in pairs: write_mrc_from_tensor( data=tensor, mrc_path=path, mrc_header=None, overwrite=self.allow_file_overwrite, ) + + self.export_correlation_table() + + def export_correlation_table(self) -> None: + """Write the held CorrelationTable to ``self.correlation_table_path``.""" + if self.correlation_table is None: + raise ValueError("No correlation_table to export.") + if self.correlation_table_path is None: + raise ValueError("No correlation_table_path specified to export to.") + self.correlation_table.to_hdf5(self.correlation_table_path) + + def load_correlation_table_from_path(self) -> None: + """Load CorrelationTable from HDF5 file at ``self.correlation_table_path``.""" + if self.correlation_table_path is None: + raise ValueError("No correlation_table_path specified to load from.") + self.correlation_table = CorrelationTable.from_hdf5(self.correlation_table_path) + + +class MatchTemplateResultHDF5(_MatchTemplateResultBase): + """Bundles all result tensors and metadata into a single HDF5 file. + + HDF5 file layout + ---------------- + All eight 2-D result tensors are stored as float32 datasets inside a + ``/tensors`` group. When ``compress`` is ``True`` (the default) each + dataset is compressed with gzip at level 4. Scalar metadata + (``total_projections``, ``total_orientations``, ``total_defocus``) are + stored as attributes on the HDF5 root group. No MRC paths are written to + the file; the path to the HDF5 file itself is the only path required at + load time. + + / (root) + │ attrs: leopard_em_version, total_projections, + │ total_orientations, total_defocus + └─ tensors/ + mip float32, shape (H, W), gzip-4 (if compress=True) + scaled_mip float32, shape (H, W), gzip-4 + correlation_average float32, shape (H, W), gzip-4 + correlation_variance float32, shape (H, W), gzip-4 + orientation_psi float32, shape (H, W), gzip-4 + orientation_theta float32, shape (H, W), gzip-4 + orientation_phi float32, shape (H, W), gzip-4 + relative_defocus float32, shape (H, W), gzip-4 + + Attributes + ---------- + hdf5_path : str + Path to the HDF5 output file. + allow_file_overwrite : bool + Whether to allow overwriting an existing file. Default is False. + compress : bool + Whether to apply gzip-4 compression to tensor datasets. Default is + True. Disable for faster writes at the cost of larger files. + + Methods + ------- + validate_hdf5_path() + Validates write permissions and overwrite policy for ``hdf5_path``. + to_hdf5() + Writes tensors and metadata to ``hdf5_path``. + from_hdf5(path, allow_file_overwrite) + Class method that loads an instance from an existing HDF5 file. + """ + + hdf5_path: str + allow_file_overwrite: bool = False + compress: bool = True + + ########################### + ### Pydantic Validators ### + ########################### + + @model_validator(mode="after") # type: ignore + def validate_hdf5_path(self) -> Self: + """Validate ``hdf5_path`` for write permissions and overwriting. + + Returns + ------- + Self + + Raises + ------ + ValueError + If the path is not writable or the file already exists and + overwriting is disabled. + """ + check_file_path_and_permissions(self.hdf5_path, self.allow_file_overwrite) + return self + + ###################### + ### I/O methods ### + ###################### + + def export_results(self) -> None: + """Write tensors and metadata to ``hdf5_path``. Alias for ``to_hdf5``.""" + self.to_hdf5() + + def to_hdf5(self) -> None: + """Write tensors and scalar metadata to ``hdf5_path``. + + Tensors are cast to float32 before writing. When ``self.compress`` is + ``True``, each dataset is compressed with gzip at level 4. + """ + compression_kwargs: dict = ( + {"compression": "gzip", "compression_opts": 4} if self.compress else {} + ) + + with h5py.File(self.hdf5_path, "w") as f: + f.attrs["leopard_em_version"] = self.leopard_em_version + f.attrs["total_projections"] = self.total_projections + f.attrs["total_orientations"] = self.total_orientations + f.attrs["total_defocus"] = self.total_defocus + + tensors_group = f.create_group(_HDF5_TENSORS_GROUP) + for name in _TENSOR_NAMES: + tensor: torch.Tensor | None = getattr(self, name) + if tensor is not None: + tensors_group.create_dataset( + name, + data=tensor.cpu().to(torch.float32).numpy(), + **compression_kwargs, + ) + + @classmethod + def from_hdf5( + cls, + path: str | os.PathLike, + allow_file_overwrite: bool = True, + compress: bool = True, + ) -> "MatchTemplateResultHDF5": + """Load a ``MatchTemplateResultHDF5`` from an existing HDF5 file. + + Parameters + ---------- + path : str | os.PathLike + Path to the HDF5 file written by ``to_hdf5``. + allow_file_overwrite : bool + Passed to the constructor. Defaults to ``True`` so that the + model validator does not reject the path of the file being loaded. + compress : bool + Passed to the constructor. Controls compression on any subsequent + ``to_hdf5`` call made on the returned instance. Default is ``True``. + + Returns + ------- + MatchTemplateResultHDF5 + """ + tensors: dict[str, torch.Tensor] = {} + + with h5py.File(path, "r") as f: + leopard_em_version = str(f.attrs.get("leopard_em_version", "unknown")) + total_projections = int(f.attrs["total_projections"]) + total_orientations = int(f.attrs["total_orientations"]) + total_defocus = int(f.attrs["total_defocus"]) + + if _HDF5_TENSORS_GROUP in f: + grp = f[_HDF5_TENSORS_GROUP] + for name in _TENSOR_NAMES: + if name in grp: + tensors[name] = torch.from_numpy(grp[name][:]) + + return cls( + hdf5_path=str(path), + allow_file_overwrite=allow_file_overwrite, + compress=compress, + leopard_em_version=leopard_em_version, + total_projections=total_projections, + total_orientations=total_orientations, + total_defocus=total_defocus, + **tensors, + ) + + +# Backward-compatibility alias — existing code importing MatchTemplateResult +# continues to receive MatchTemplateResultMRC unchanged. +MatchTemplateResult = MatchTemplateResultMRC diff --git a/tests/backend/test_core_match_template.py b/tests/backend/test_core_match_template.py index efe1b8b2..7db081c2 100644 --- a/tests/backend/test_core_match_template.py +++ b/tests/backend/test_core_match_template.py @@ -56,15 +56,21 @@ def mrcfile_allclose(path_a: str, path_b: str, **kwargs) -> bool: ) @pytest.mark.slow def test_core_match_template(): - download_comparison_data() + # download_comparison_data() mt_manager = setup_match_template_manager() # Run the match template program mt_manager.run_match_template( orientation_batch_size=ORIENTATION_BATCH_SIZE, do_result_export=True, # Saves the statistics immediately upon completion + do_valid_cropping=False, # testing backend doing valid cropping in-place ) + corr_table = mt_manager.match_template_result.correlation_table + import pandas as pd + + pd.DataFrame(corr_table).to_csv("tests/tmp/test_corr_table.csv") + # Ensure the MIPs are the same, if they are not then there's an issue... assert mrcfile_allclose( "tests/tmp/test_match_template_xenon_216_000_0_output_mip.mrc", diff --git a/tests/pydantic_models/test_basic_imports.py b/tests/pydantic_models/test_basic_imports.py index c9441302..b19ea571 100644 --- a/tests/pydantic_models/test_basic_imports.py +++ b/tests/pydantic_models/test_basic_imports.py @@ -56,7 +56,12 @@ def test_data_structure_imports(): def test_results_import(): """Test for results imports.""" try: - from leopard_em.pydantic_models.results import MatchTemplateResult + from leopard_em.pydantic_models.results import ( + CorrelationTable, + MatchTemplateResult, + MatchTemplateResultHDF5, + MatchTemplateResultMRC, + ) except ImportError as e: raise ImportError( "Failed to import one or more results classes from " diff --git a/tests/pydantic_models/test_correlation_table.py b/tests/pydantic_models/test_correlation_table.py new file mode 100644 index 00000000..f3d7d0e4 --- /dev/null +++ b/tests/pydantic_models/test_correlation_table.py @@ -0,0 +1,322 @@ +"""Unit tests for CorrelationTable and derive_orientation_grid_from_full_angles.""" + +import os +import tempfile + +import pytest +import torch + +from leopard_em.pydantic_models.results.correlation_table import ( + CorrelationTable, + derive_orientation_grid_from_full_angles, +) + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def grid_euler_angles() -> torch.Tensor: + """2 (phi, theta) pairs x 3 psi values → 6 orientations.""" + return torch.tensor( + [ + [0.0, 0.0, 0.0], + [0.0, 0.0, 90.0], + [0.0, 0.0, 180.0], + [45.0, 30.0, 0.0], + [45.0, 30.0, 90.0], + [45.0, 30.0, 180.0], + ] + ) + + +@pytest.fixture() +def minimal_table() -> CorrelationTable: + """A small, hand-crafted CorrelationTable for roundtrip tests.""" + return CorrelationTable( + correlation_threshold=5.5, + num_observations=3, + defocus_offsets=[-500.0, 0.0, 500.0], + phi_theta_angles=[(0.0, 0.0), (45.0, 30.0)], + psi_angles=[0.0, 90.0, 180.0], + search_index=[0, 5, 11], + x=[10, 20, 30], + y=[15, 25, 35], + correlation_value=[6.1, 7.2, 5.8], + correlation_mean=[0.1, 0.2, 0.3], + correlation_variance=[0.5, 0.6, 0.7], + ) + + +@pytest.fixture() +def empty_table() -> CorrelationTable: + """A CorrelationTable with no detections.""" + return CorrelationTable( + correlation_threshold=5.5, + num_observations=0, + defocus_offsets=[-500.0, 0.0], + phi_theta_angles=[(0.0, 0.0)], + psi_angles=[0.0, 90.0], + search_index=[], + x=[], + y=[], + correlation_value=[], + correlation_mean=[], + correlation_variance=[], + ) + + +# --------------------------------------------------------------------------- +# derive_orientation_grid_from_full_angles +# --------------------------------------------------------------------------- + + +class TestDeriveOrientationGrid: + def test_basic_grid(self, grid_euler_angles): + phi_theta, psi = derive_orientation_grid_from_full_angles(grid_euler_angles) + assert phi_theta == [(0.0, 0.0), (45.0, 30.0)] + assert psi == [0.0, 90.0, 180.0] + + def test_single_phi_theta(self): + angles = torch.tensor([[10.0, 20.0, 0.0], [10.0, 20.0, 45.0]]) + phi_theta, psi = derive_orientation_grid_from_full_angles(angles) + assert phi_theta == [(10.0, 20.0)] + assert psi == [0.0, 45.0] + + def test_single_psi(self): + angles = torch.tensor([[0.0, 0.0, 0.0], [45.0, 30.0, 0.0]]) + phi_theta, psi = derive_orientation_grid_from_full_angles(angles) + assert phi_theta == [(0.0, 0.0), (45.0, 30.0)] + assert psi == [0.0] + + def test_return_lengths_match_grid(self, grid_euler_angles): + phi_theta, psi = derive_orientation_grid_from_full_angles(grid_euler_angles) + assert len(phi_theta) * len(psi) == grid_euler_angles.shape[0] + + +# --------------------------------------------------------------------------- +# CorrelationTable construction +# --------------------------------------------------------------------------- + + +class TestCorrelationTableConstruction: + def test_basic_construction(self, minimal_table): + assert minimal_table.num_observations == 3 + assert minimal_table.correlation_threshold == 5.5 + assert len(minimal_table.search_index) == 3 + assert len(minimal_table.x) == 3 + + def test_empty_construction(self, empty_table): + assert empty_table.num_observations == 0 + assert empty_table.search_index == [] + assert empty_table.x == [] + + +# --------------------------------------------------------------------------- +# DataFrame roundtrip +# --------------------------------------------------------------------------- + + +class TestDataFrameRoundtrip: + def test_columns_present(self, minimal_table): + df = minimal_table.to_dataframe() + expected = { + "search_index", + "x", + "y", + "correlation_value", + "correlation_mean", + "correlation_variance", + } + assert expected == set(df.columns) + + def test_metadata_in_attrs(self, minimal_table): + df = minimal_table.to_dataframe() + assert df.attrs["correlation_threshold"] == minimal_table.correlation_threshold + assert df.attrs["num_observations"] == minimal_table.num_observations + assert df.attrs["defocus_offsets"] == minimal_table.defocus_offsets + assert df.attrs["psi_angles"] == minimal_table.psi_angles + + def test_roundtrip_detection_data(self, minimal_table): + recovered = CorrelationTable.from_dataframe(minimal_table.to_dataframe()) + assert recovered.search_index == minimal_table.search_index + assert recovered.x == minimal_table.x + assert recovered.y == minimal_table.y + assert recovered.correlation_value == pytest.approx( + minimal_table.correlation_value + ) + assert recovered.correlation_mean == pytest.approx( + minimal_table.correlation_mean + ) + + def test_roundtrip_search_space(self, minimal_table): + recovered = CorrelationTable.from_dataframe(minimal_table.to_dataframe()) + assert recovered.defocus_offsets == minimal_table.defocus_offsets + assert recovered.phi_theta_angles == minimal_table.phi_theta_angles + assert recovered.psi_angles == minimal_table.psi_angles + + def test_row_count(self, minimal_table): + df = minimal_table.to_dataframe() + assert len(df) == minimal_table.num_observations + + def test_empty_table_roundtrip(self, empty_table): + recovered = CorrelationTable.from_dataframe(empty_table.to_dataframe()) + assert recovered.num_observations == 0 + assert recovered.x == [] + + +# --------------------------------------------------------------------------- +# HDF5 roundtrip +# --------------------------------------------------------------------------- + + +class TestHDF5Roundtrip: + def test_roundtrip_detection_data(self, minimal_table): + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: + path = f.name + try: + minimal_table.to_hdf5(path) + recovered = CorrelationTable.from_hdf5(path) + assert recovered.search_index == minimal_table.search_index + assert recovered.x == minimal_table.x + assert recovered.y == minimal_table.y + assert recovered.correlation_value == pytest.approx( + minimal_table.correlation_value, abs=1e-5 + ) + finally: + os.unlink(path) + + def test_roundtrip_search_space(self, minimal_table): + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: + path = f.name + try: + minimal_table.to_hdf5(path) + recovered = CorrelationTable.from_hdf5(path) + assert recovered.defocus_offsets == pytest.approx( + minimal_table.defocus_offsets, abs=1e-5 + ) + assert recovered.phi_theta_angles == pytest.approx( + minimal_table.phi_theta_angles, abs=1e-5 + ) + assert recovered.psi_angles == pytest.approx( + minimal_table.psi_angles, abs=1e-5 + ) + finally: + os.unlink(path) + + def test_roundtrip_metadata(self, minimal_table): + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: + path = f.name + try: + minimal_table.to_hdf5(path) + recovered = CorrelationTable.from_hdf5(path) + assert ( + recovered.correlation_threshold == minimal_table.correlation_threshold + ) + assert recovered.num_observations == minimal_table.num_observations + finally: + os.unlink(path) + + def test_empty_table_roundtrip(self, empty_table): + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: + path = f.name + try: + empty_table.to_hdf5(path) + recovered = CorrelationTable.from_hdf5(path) + assert recovered.num_observations == 0 + assert recovered.search_index == [] + finally: + os.unlink(path) + + def test_file_is_created(self, minimal_table): + with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as f: + path = f.name + os.unlink(path) + try: + minimal_table.to_hdf5(path) + assert os.path.isfile(path) + finally: + if os.path.exists(path): + os.unlink(path) + + +# --------------------------------------------------------------------------- +# from_match_template_results factory +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def factory_inputs(grid_euler_angles): + """Common inputs for from_match_template_results tests.""" + H, W = 64, 80 + defocus_values = torch.tensor([-500.0, 0.0, 500.0]) + corr_avg = torch.rand(H, W) + corr_var = torch.rand(H, W) + proc_table = { + "threshold": 5.5, + "global_idx": [0, 5, 11], + "x": [10, 20, 30], + "y": [15, 25, 35], + "correlation": [6.1, 7.2, 5.8], + } + return { + "processed_correlation_table": proc_table, + "defocus_values": defocus_values, + "euler_angles": grid_euler_angles, + "correlation_average": corr_avg, + "correlation_variance_map": corr_var, + } + + +class TestFromMatchTemplateResults: + def test_search_space_derivation(self, factory_inputs): + ct = CorrelationTable.from_match_template_results(**factory_inputs) + assert ct.phi_theta_angles == [(0.0, 0.0), (45.0, 30.0)] + assert ct.psi_angles == [0.0, 90.0, 180.0] + assert ct.defocus_offsets == pytest.approx([-500.0, 0.0, 500.0]) + + def test_num_observations(self, factory_inputs): + ct = CorrelationTable.from_match_template_results(**factory_inputs) + assert ct.num_observations == 3 + + def test_search_index_passthrough(self, factory_inputs): + ct = CorrelationTable.from_match_template_results(**factory_inputs) + assert ct.search_index == [0, 5, 11] + + def test_xy_positions(self, factory_inputs): + ct = CorrelationTable.from_match_template_results(**factory_inputs) + assert ct.x == [10, 20, 30] + assert ct.y == [15, 25, 35] + + def test_mean_variance_looked_up_from_tensors(self, factory_inputs): + corr_avg = factory_inputs["correlation_average"] + corr_var = factory_inputs["correlation_variance_map"] + ct = CorrelationTable.from_match_template_results(**factory_inputs) + + xs = factory_inputs["processed_correlation_table"]["x"] + ys = factory_inputs["processed_correlation_table"]["y"] + expected_mean = [corr_avg[y, x].item() for x, y in zip(xs, ys)] + expected_var = [corr_var[y, x].item() for x, y in zip(xs, ys)] + + assert ct.correlation_mean == pytest.approx(expected_mean) + assert ct.correlation_variance == pytest.approx(expected_var) + + def test_empty_detections(self, factory_inputs, grid_euler_angles): + empty_proc = { + "threshold": 5.5, + "global_idx": [], + "x": [], + "y": [], + "correlation": [], + } + ct = CorrelationTable.from_match_template_results( + processed_correlation_table=empty_proc, + defocus_values=factory_inputs["defocus_values"], + euler_angles=grid_euler_angles, + correlation_average=factory_inputs["correlation_average"], + correlation_variance_map=factory_inputs["correlation_variance_map"], + ) + assert ct.num_observations == 0 + assert ct.correlation_mean == [] + assert ct.correlation_variance == [] diff --git a/tests/pydantic_models/test_match_template_result.py b/tests/pydantic_models/test_match_template_result.py new file mode 100644 index 00000000..813f5172 --- /dev/null +++ b/tests/pydantic_models/test_match_template_result.py @@ -0,0 +1,368 @@ +"""Unit tests for MatchTemplateResultMRC and MatchTemplateResultHDF5.""" + +import importlib.metadata + +import h5py +import pytest +import torch + +from leopard_em.pydantic_models.results import ( + MatchTemplateResult, + MatchTemplateResultHDF5, + MatchTemplateResultMRC, +) +from leopard_em.pydantic_models.results.match_template_result import _TENSOR_NAMES + +# --------------------------------------------------------------------------- +# Shared helpers / fixtures +# --------------------------------------------------------------------------- + +_TENSOR_SHAPE = (16, 16) + + +def _make_tensors(shape=_TENSOR_SHAPE) -> dict[str, torch.Tensor]: + """Return a dict of distinct float32 tensors for all eight result fields.""" + return { + name: torch.full(shape, float(i), dtype=torch.float32) + for i, name in enumerate(_TENSOR_NAMES) + } + + +def _mrc_paths(tmp_path) -> dict[str, str]: + """Return a dict of non-existent MRC output paths under *tmp_path*.""" + names = [ + "mip_path", + "scaled_mip_path", + "correlation_average_path", + "correlation_variance_path", + "orientation_psi_path", + "orientation_theta_path", + "orientation_phi_path", + "relative_defocus_path", + ] + return {name: str(tmp_path / f"{name}.mrc") for name in names} + + +@pytest.fixture +def tensors(): + return _make_tensors() + + +@pytest.fixture +def mrc_result(tmp_path, tensors): + """Construct a fully-populated MatchTemplateResultMRC.""" + return MatchTemplateResultMRC( + **_mrc_paths(tmp_path), + total_projections=100, + total_orientations=50, + total_defocus=2, + **tensors, + ) + + +@pytest.fixture +def hdf5_result(tmp_path, tensors): + """Construct a fully-populated MatchTemplateResultHDF5.""" + return MatchTemplateResultHDF5( + hdf5_path=str(tmp_path / "result.h5"), + allow_file_overwrite=True, + total_projections=100, + total_orientations=50, + total_defocus=2, + **tensors, + ) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def test_public_imports(): + """All expected public names are importable from the results package.""" + from leopard_em.pydantic_models.results import ( # noqa: F401 + MatchTemplateResult, + MatchTemplateResultHDF5, + MatchTemplateResultMRC, + ) + + +def test_base_class_not_public(): + """The private base class is not exported from the public API.""" + with pytest.raises(ImportError): + from leopard_em.pydantic_models.results import ( + _MatchTemplateResultBase, # noqa: F401 + ) + + +def test_backward_compat_alias(): + """MatchTemplateResult is an alias for MatchTemplateResultMRC.""" + assert MatchTemplateResult is MatchTemplateResultMRC + + +# --------------------------------------------------------------------------- +# MatchTemplateResultMRC — construction and validation +# --------------------------------------------------------------------------- + + +def test_mrc_construction(tmp_path): + """MatchTemplateResultMRC is instantiated with the required MRC paths.""" + result = MatchTemplateResultMRC(**_mrc_paths(tmp_path)) + assert result.allow_file_overwrite is False + assert result.total_projections == 0 + assert result.total_orientations == 0 + assert result.total_defocus == 0 + + +def test_mrc_rejects_existing_file(tmp_path): + """validate_paths raises when a path already exists and overwrite is off.""" + paths = _mrc_paths(tmp_path) + # Pre-create one of the output files + existing = paths["mip_path"] + open(existing, "w").close() + + with pytest.raises(ValueError, match="already exists"): + MatchTemplateResultMRC(**paths, allow_file_overwrite=False) + + +def test_mrc_allows_existing_file_with_overwrite_flag(tmp_path): + """validate_paths passes when allow_file_overwrite=True, even if file exists.""" + paths = _mrc_paths(tmp_path) + open(paths["mip_path"], "w").close() + # Should not raise + MatchTemplateResultMRC(**paths, allow_file_overwrite=True) + + +def test_mrc_version_auto_populated(tmp_path): + """leopard_em_version is set to the installed package version on construction.""" + result = MatchTemplateResultMRC(**_mrc_paths(tmp_path)) + expected = importlib.metadata.version("leopard_em") + assert result.leopard_em_version == expected + + +# --------------------------------------------------------------------------- +# MatchTemplateResultHDF5 — construction and validation +# --------------------------------------------------------------------------- + + +def test_hdf5_construction(tmp_path): + """MatchTemplateResultHDF5 is instantiated with only hdf5_path.""" + result = MatchTemplateResultHDF5(hdf5_path=str(tmp_path / "out.h5")) + assert result.compress is True + assert result.allow_file_overwrite is False + assert result.total_projections == 0 + + +def test_hdf5_rejects_existing_file(tmp_path): + """validate_hdf5_path raises when the file exists and overwrite is off.""" + path = tmp_path / "out.h5" + path.touch() + with pytest.raises(ValueError, match="already exists"): + MatchTemplateResultHDF5(hdf5_path=str(path), allow_file_overwrite=False) + + +def test_hdf5_allows_existing_file_with_overwrite_flag(tmp_path): + """validate_hdf5_path passes when allow_file_overwrite=True.""" + path = tmp_path / "out.h5" + path.touch() + MatchTemplateResultHDF5(hdf5_path=str(path), allow_file_overwrite=True) + + +def test_hdf5_version_auto_populated(tmp_path): + """leopard_em_version is set to the installed package version on construction.""" + result = MatchTemplateResultHDF5(hdf5_path=str(tmp_path / "out.h5")) + expected = importlib.metadata.version("leopard_em") + assert result.leopard_em_version == expected + + +# --------------------------------------------------------------------------- +# MatchTemplateResultHDF5 — to_hdf5 +# --------------------------------------------------------------------------- + + +def test_to_hdf5_creates_file(hdf5_result): + """to_hdf5 creates the file at hdf5_path.""" + hdf5_result.to_hdf5() + import os + + assert os.path.exists(hdf5_result.hdf5_path) + + +def test_to_hdf5_root_attributes(hdf5_result): + """to_hdf5 writes scalar metadata as HDF5 root attributes.""" + hdf5_result.to_hdf5() + with h5py.File(hdf5_result.hdf5_path, "r") as f: + assert f.attrs["total_projections"] == 100 + assert f.attrs["total_orientations"] == 50 + assert f.attrs["total_defocus"] == 2 + assert f.attrs["leopard_em_version"] == hdf5_result.leopard_em_version + + +def test_to_hdf5_tensor_group_exists(hdf5_result): + """to_hdf5 creates a 'tensors' group containing all eight datasets.""" + hdf5_result.to_hdf5() + with h5py.File(hdf5_result.hdf5_path, "r") as f: + assert "tensors" in f + for name in _TENSOR_NAMES: + assert name in f["tensors"], f"missing dataset: {name}" + + +def test_to_hdf5_tensor_shape_and_dtype(hdf5_result): + """Each tensor dataset has the expected shape and float32 dtype.""" + hdf5_result.to_hdf5() + with h5py.File(hdf5_result.hdf5_path, "r") as f: + for name in _TENSOR_NAMES: + ds = f["tensors"][name] + assert ds.shape == _TENSOR_SHAPE, f"{name}: shape mismatch" + assert ds.dtype == "float32", f"{name}: dtype mismatch" + + +def test_to_hdf5_tensor_values(hdf5_result): + """Tensor values written to HDF5 match those held in the object.""" + hdf5_result.to_hdf5() + with h5py.File(hdf5_result.hdf5_path, "r") as f: + for name in _TENSOR_NAMES: + stored = torch.from_numpy(f["tensors"][name][:]) + original = getattr(hdf5_result, name).cpu().to(torch.float32) + assert torch.allclose(stored, original), f"{name}: value mismatch" + + +def test_to_hdf5_compression_enabled(hdf5_result): + """When compress=True (default), datasets are gzip-compressed.""" + hdf5_result.to_hdf5() + with h5py.File(hdf5_result.hdf5_path, "r") as f: + for name in _TENSOR_NAMES: + assert f["tensors"][name].compression == "gzip", f"{name}: expected gzip" + + +def test_to_hdf5_compression_disabled(tmp_path, tensors): + """When compress=False, datasets are written without compression.""" + result = MatchTemplateResultHDF5( + hdf5_path=str(tmp_path / "result_uncompressed.h5"), + allow_file_overwrite=True, + compress=False, + **tensors, + ) + result.to_hdf5() + with h5py.File(result.hdf5_path, "r") as f: + for name in _TENSOR_NAMES: + assert ( + f["tensors"][name].compression is None + ), f"{name}: expected no compression" + + +# --------------------------------------------------------------------------- +# MatchTemplateResultHDF5 — from_hdf5 round-trip +# --------------------------------------------------------------------------- + + +def test_from_hdf5_roundtrip_metadata(hdf5_result): + """Scalar metadata survives a to_hdf5 / from_hdf5 round-trip.""" + hdf5_result.to_hdf5() + loaded = MatchTemplateResultHDF5.from_hdf5(hdf5_result.hdf5_path) + + assert loaded.total_projections == hdf5_result.total_projections + assert loaded.total_orientations == hdf5_result.total_orientations + assert loaded.total_defocus == hdf5_result.total_defocus + assert loaded.leopard_em_version == hdf5_result.leopard_em_version + + +def test_from_hdf5_roundtrip_tensors(hdf5_result): + """All eight tensors survive a to_hdf5 / from_hdf5 round-trip.""" + hdf5_result.to_hdf5() + loaded = MatchTemplateResultHDF5.from_hdf5(hdf5_result.hdf5_path) + + for name in _TENSOR_NAMES: + original = getattr(hdf5_result, name).cpu().to(torch.float32) + restored = getattr(loaded, name) + assert restored is not None, f"{name} is None after loading" + assert torch.allclose( + original, restored + ), f"{name}: value mismatch after round-trip" + + +def test_from_hdf5_preserves_hdf5_path(hdf5_result): + """The loaded instance's hdf5_path matches the path it was loaded from.""" + hdf5_result.to_hdf5() + loaded = MatchTemplateResultHDF5.from_hdf5(hdf5_result.hdf5_path) + assert loaded.hdf5_path == hdf5_result.hdf5_path + + +def test_from_hdf5_missing_version_attribute(tmp_path): + """Files without leopard_em_version degrade gracefully to 'unknown'.""" + path = tmp_path / "old_result.h5" + with h5py.File(str(path), "w") as f: + f.attrs["total_projections"] = 0 + f.attrs["total_orientations"] = 0 + f.attrs["total_defocus"] = 0 + + loaded = MatchTemplateResultHDF5.from_hdf5(str(path)) + assert loaded.leopard_em_version == "unknown" + + +def test_from_hdf5_partial_tensors(tmp_path): + """from_hdf5 succeeds when only a subset of tensors are stored.""" + path = tmp_path / "partial.h5" + with h5py.File(str(path), "w") as f: + f.attrs["leopard_em_version"] = "test" + f.attrs["total_projections"] = 0 + f.attrs["total_orientations"] = 0 + f.attrs["total_defocus"] = 0 + grp = f.create_group("tensors") + grp.create_dataset("mip", data=torch.zeros(_TENSOR_SHAPE).numpy()) + + loaded = MatchTemplateResultHDF5.from_hdf5(str(path)) + assert loaded.mip is not None + assert loaded.scaled_mip is None + + +# --------------------------------------------------------------------------- +# export_results alias +# --------------------------------------------------------------------------- + + +def test_export_results_alias(hdf5_result): + """export_results() on MatchTemplateResultHDF5 is equivalent to to_hdf5().""" + import os + + hdf5_result.export_results() + assert os.path.exists(hdf5_result.hdf5_path) + with h5py.File(hdf5_result.hdf5_path, "r") as f: + assert "tensors" in f + + +# --------------------------------------------------------------------------- +# apply_valid_cropping (inherited from base) +# --------------------------------------------------------------------------- + + +def test_apply_valid_cropping_mrc(mrc_result): + """apply_valid_cropping reduces tensor dimensions correctly for MRC result.""" + template_shape = (4, 4) + expected_h = _TENSOR_SHAPE[0] - template_shape[0] + 1 # 13 + expected_w = _TENSOR_SHAPE[1] - template_shape[1] + 1 # 13 + + mrc_result.apply_valid_cropping(template_shape) + + for name in _TENSOR_NAMES: + t = getattr(mrc_result, name) + assert t.shape == ( + expected_h, + expected_w, + ), f"{name}: wrong shape after cropping" + + +def test_apply_valid_cropping_hdf5(hdf5_result): + """apply_valid_cropping reduces tensor dimensions correctly for HDF5 result.""" + template_shape = (4, 4) + expected_h = _TENSOR_SHAPE[0] - template_shape[0] + 1 + expected_w = _TENSOR_SHAPE[1] - template_shape[1] + 1 + + hdf5_result.apply_valid_cropping(template_shape) + + for name in _TENSOR_NAMES: + t = getattr(hdf5_result, name) + assert t.shape == ( + expected_h, + expected_w, + ), f"{name}: wrong shape after cropping" diff --git a/tests/self_consistency/test_backend_cross_correlate.py b/tests/self_consistency/test_backend_cross_correlate.py index e1f0fbae..216a3530 100644 --- a/tests/self_consistency/test_backend_cross_correlate.py +++ b/tests/self_consistency/test_backend_cross_correlate.py @@ -26,15 +26,16 @@ from leopard_em.backend.cross_correlation import ( do_batched_orientation_cross_correlate, + do_batched_orientation_cross_correlate_zipfft, do_streamed_orientation_cross_correlate, ) from leopard_em.utils import get_cs_range -IMAGE_SHAPE = (1024, 1024) -TEMPLATE_SHAPE = (128, 128, 128) -NUM_ORIENTATIONS = 10 +IMAGE_SHAPE = (4096, 4096) +TEMPLATE_SHAPE = (512, 512, 512) +NUM_ORIENTATIONS = 4 NUM_DEFOCUS_VALUES = 5 -NUM_PIXEL_SIZES = 4 +NUM_PIXEL_SIZES = 2 NUM_STREAMS = 1 @@ -133,3 +134,55 @@ def test_stream_and_batch_cross_correlate_consistency(sample_input_data): f"Streamed and batched cross-correlation results not within tolerance.\n" f"Max absolute difference: {max_abs_diff}\n" ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device not available") +def test_batched_zipfft_cross_correlate_consistency(sample_input_data): + """Test that the batched zip-FFT cross-correlation method is consistent.""" + cross_correlate_kwargs = sample_input_data + + batched_result = do_batched_orientation_cross_correlate(**cross_correlate_kwargs) + + # NOTE: zipFFT does valid cropping internally, so need to adjust the bached result + batched_result = batched_result[ + ..., + : IMAGE_SHAPE[-2] - TEMPLATE_SHAPE[-2] + 1, + : IMAGE_SHAPE[-1] - TEMPLATE_SHAPE[-1] + 1, + ] + + # NOTE: Need to transpose the input 'image_dft' for the zip-FFT method + image_dft = cross_correlate_kwargs["image_dft"].transpose(-1, -2).clone() + cross_correlate_kwargs["image_dft"] = image_dft + + zipfft_result = do_batched_orientation_cross_correlate_zipfft( + **cross_correlate_kwargs + ) + + assert zipfft_result.shape == batched_result.shape + + max_abs_diff = (zipfft_result - batched_result).abs().max().item() + max_rel_diff = ( + ((zipfft_result - batched_result).abs() / batched_result.abs().clamp_min(1e-8)) + .max() + .item() + ) + + # NOTE: Taking the L2 norm of the difference since FFT plans execute differently + # and are not guaranteed to be bitwise identical and likely include small numerical + # differences. + l2_norm_diff = torch.norm(zipfft_result - batched_result).item() + l2_norm_diff /= torch.prod(torch.tensor(zipfft_result.shape)).item() + + # ### DEBUGGING: Save each result to file + # import numpy as np + + # np.save("batched_result.npy", batched_result.cpu().numpy()) + # np.save("zipfft_result.npy", zipfft_result.cpu().numpy()) + # ### END DEBUGGING + + assert l2_norm_diff < 1e-6, ( + f"Batched and zip-FFT cross-correlation results differ too much.\n" + f"L2 norm of difference: {l2_norm_diff}\n" + f"Max absolute difference: {max_abs_diff}\n" + f"Max relative difference: {max_rel_diff}\n" + )