diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 00000000..01a10ab7 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,55 @@ +name: Documentation + +on: + push: + branches: + - main + tags: + - "v*" + workflow_dispatch: + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: "pages" + cancel-in-progress: false + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.x" + cache: "pip" + cache-dependency-path: "pyproject.toml" + + - name: Install docs dependencies + run: pip install ".[docs]" + + - name: Generate API stubs + run: python scripts/gen_api_docs.py + + - name: Build docs + run: zensical build --clean + + - name: Upload Pages artifact + uses: actions/upload-pages-artifact@v3 + with: + path: ./site + + deploy: + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + runs-on: ubuntu-latest + needs: build + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/pylint.yml b/.github/workflows/pylint.yml index ab3f27e5..69c77f0c 100644 --- a/.github/workflows/pylint.yml +++ b/.github/workflows/pylint.yml @@ -7,7 +7,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.11", "3.12", "3.13"] steps: - uses: actions/checkout@v6 - name: Set up Python ${{ matrix.python-version }} diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 5ea00187..fbfd4183 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -32,7 +32,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.11", "3.12", "3.13"] platform: [ubuntu-latest, macos-latest] steps: diff --git a/.gitignore b/.gitignore index 825c1ac8..e3856631 100644 --- a/.gitignore +++ b/.gitignore @@ -99,9 +99,10 @@ ENV/ # Rope project settings .ropeproject -# mkdocs documentation +# mkdocs/zensical documentation /site /docs/generated/ +/docs/api/ /docs/examples/60S_*.mrc /docs/examples/xenon_*.mrc /docs/examples/xenon_*.csv diff --git a/.pylintrc b/.pylintrc index 69ac17d3..d1476f0f 100644 --- a/.pylintrc +++ b/.pylintrc @@ -6,4 +6,16 @@ generated-members=numpy.*,torch.* # See https://github.com/pytorch/pytorch/issu # NOTE: We've adopted a more data/functional programming style, so some data related # data is not packed/unpacked as objects and rather passed as data arrays/options. max-args=10 -max-positional-arguments=10 \ No newline at end of file +max-positional-arguments=10 + +[DESIGN] +# Result/manager Pydantic models and the batched inspection/correlation backends +# legitimately carry more fields and locals than pylint's defaults allow; raise the +# ceilings rather than sprinkling per-class/per-function disables. +max-attributes=12 +max-locals=20 +max-statements=60 + +[MESSAGES CONTROL] +# TODOs are tracked intentionally in-code; don't fail CI on them. +disable=fixme \ No newline at end of file diff --git a/README.md b/README.md index e6bc9a50..aeffb16d 100644 --- a/README.md +++ b/README.md @@ -88,18 +88,18 @@ Further information on development and contributing to the repo can be found in If you use Leopard-EM in your research, please include the following citations in your biblography: -- [Leopard-EM pre-print](https://doi.org/10.1101/2025.08.26.672452) +- [Leopard-EM Manuscript](https://doi.org/10.1107/S2059798325009982) > ``` -> @article {Giammar2025.08.26.672452, -> author = {Giammar, Matthew David and Dickerson, Joshua Luke and Hall, Laina Nicole and Lucas, Bronwyn Ayla}, -> title = {Leopard-EM: An extensible 2DTM package to accelerate in situ structural biology}, -> elocation-id = {2025.08.26.672452}, -> year = {2025}, -> doi = {10.1101/2025.08.26.672452}, -> publisher = {Cold Spring Harbor Laboratory}, -> URL = {https://www.biorxiv.org/content/early/2025/08/29/2025.08.26.672452}, -> eprint = {https://www.biorxiv.org/content/early/2025/08/29/2025.08.26.672452.full.pdf}, -> journal = {bioRxiv} +> @article{Giammar:rr5256, +> author = "Giammar, Matthew D. and Dickerson, Joshua L. and Hall, Laina N. and Lucas, Bronwyn A.", +> title = "{{\it Leopard-EM}: an extensible 2D template-matching package to accelerate {\it in situ} structural biology}", +> journal = "Acta Crystallographica Section D", +> year = "2026", +> volume = "82", +> number = "1", +> pages = "30--42", +> month = "Jan", +> doi = {10.1107/S2059798325009982}, +> url = {https://doi.org/10.1107/S2059798325009982}, > } > ``` - diff --git a/benchmark/benchmark_match_template.py b/benchmark/benchmark_match_template.py index 3d8b73cb..18566c63 100644 --- a/benchmark/benchmark_match_template.py +++ b/benchmark/benchmark_match_template.py @@ -28,8 +28,17 @@ ZENODO_URL = "https://zenodo.org/records/17069838" -def download_comparison_data() -> None: - """Downloads the example data from Zenodo.""" +def download_comparison_data(force_download: bool = False) -> None: + """Downloads the example data from Zenodo, skipping if already present. + + Parameters + ---------- + force_download : bool + If True, re-download the data even if it appears to already be present. + """ + if YAML_PATH.exists() and not force_download: + return + subprocess.run( ["zenodo_get", f"--output-dir={DOWNLOAD_DIR}", ZENODO_URL], check=True ) @@ -76,6 +85,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 +112,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 @@ -133,11 +144,13 @@ def benchmark_match_template_single_run( } -def run_benchmark(orientation_batch_size: int, num_runs: int) -> dict[str, Any]: +def run_benchmark( + orientation_batch_size: int, num_runs: int, force_download: bool = False +) -> 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(force_download=force_download) print("Done!") # Get CUDA device properties @@ -228,7 +241,19 @@ def save_benchmark_results(result: dict, output_file: str) -> None: type=str, help="Output file for benchmark results (default: benchmark_results.json)", ) -def main(orientation_batch_size: int, num_runs: int, output_file: str): +@click.option( + "--force-download", + is_flag=True, + default=False, + help="Re-download benchmarking data even if it already appears to be present " + "(e.g. to restore files that were manually modified for local testing).", +) +def main( + orientation_batch_size: int, + num_runs: int, + output_file: str, + force_download: bool, +): """Main benchmarking function with Click CLI interface.""" if not torch.cuda.is_available(): print("CUDA not available exiting...") @@ -239,7 +264,9 @@ def main(orientation_batch_size: int, num_runs: int, output_file: str): print(f" Number of runs: {num_runs}") print(f" Output file: {output_file}") - result = run_benchmark(orientation_batch_size, num_runs) + result = run_benchmark( + orientation_batch_size, num_runs, force_download=force_download + ) # pprint(result) save_benchmark_results(result, output_file) diff --git a/docs/data_formats.md b/docs/data_formats.md index b97b5fd5..730f9a04 100644 --- a/docs/data_formats.md +++ b/docs/data_formats.md @@ -1,24 +1,26 @@ # Description of Data Formats -To increase interoperability between external packages, we herein describe the different data formats used as input to and export from programs. +To increase interoperability between external packages, we describe the different data formats used as input to and export from programs. Orientations on a per-particle bases are currently stored as [Euler angles](https://en.wikipedia.org/wiki/Euler_angles#) in ZYZ format with angles ordered as \( \phi, \theta, \psi \). !!! note - + Leopard-Em is still undergoing development, and the exact way the data is represented might change in the future. We will document any data format changes between versions. ## Data from the match template program The match template program collates statistics from a large number of cross-correlograms taken over an orientation and defocus search space. -See the API on the `MatchTemplateResult` object for further information on how these data are stored in memory, but here we provide an overview of what files get written to disk. +See the API on the `MatchTemplateResult` object for further information on how these data are stored in memory. +Here we provide an overview of what files get written to disk. ### Best statistic maps -Each of the "best" statistics (i.e. orientation, defocus) is stored on a per-position basis in what we dub "statistics maps" saved as `.mrc` files. + +Each of the "best" statistics (i.e. orientation, defocus) is stored on a per-position basis in what we dub "statistics maps". We have the following tracked statistics for each valid (x, y) position: -- Maximum Intensity Projection (MIP): Maximum attained cross-correlation value over the entire search space. -- Scaled MIP (z-score or 2DTM SNR): The MIP value normalized by the mean and variance of the cross-correlation over the entire search space. +- Maximum Intensity Projection (MIP): Maximum attained cross-correlation value over the entire search space, per-pixel. +- Scaled MIP (z-score or 2DTM SNR): The MIP value normalized by the mean and variance of the cross-correlation over the entire search space, per-pixel. - Correlation Mean: The mean of the cross-correlation values over the entire search space. Used to calculate the scaled MIP. - Correlation Variance: The variance of the cross-correlation values over the entire search space. Used to calculate the scaled MIP. - Phi: The \( \phi \) angle (in degrees) which produced the MIP value. @@ -26,7 +28,9 @@ We have the following tracked statistics for each valid (x, y) position: - Psi: The \( \psi \) angle (in degrees) which produced the MIP value. - Defocus: The relative defocus value (in Angstroms, relative to CTF defocus of micrograph) which produced the MIP value. -Each of these statistics maps are saved to disk in the [MRC format](https://www.ccpem.ac.uk/mrc-format/mrc2014/) based on paths provided in the `MatchTemplateResult` object. +Leopard-EM supports two storage back-ends for these statistics maps: `MatchTemplateResultMRC` (the default, saving each statistic as its own [MRC format](https://www.ccpem.ac.uk/mrc-format/mrc2014/) file, based on paths provided in the object) and `MatchTemplateResultHDF5` (bundling all statistics into a single HDF5 file). +`match_template` also optionally records a sparse `CorrelationTable` of every detection which crossed a correlation threshold, in addition to the per-pixel best-statistic maps above. +See [Match Template Result Formats](data_formats/match_template_results.md) for a full comparison of the two back-ends, the `CorrelationTable` format, and code examples for reading/writing each. #### A note on correlation modes and output shapes @@ -60,43 +64,54 @@ We take a verbose approach to constructing this DataFrame where some columns sto Additional columns besides locations and orientations are included in the DataFrame to increase the utility of the data, namely the construction of `ParticleStack` objects. The columns and corresponding descriptions are as follows: -| Column Name | Type | Description | -|-------------------------------|-------|-------------| -| `particle_index` | int | An integer descriptor for ordering picked particles. Useful when operating on multiple dataframes. -| `mip` | float | Maximum cross-correlation value over all search orientations and relative defocus values. -| `scaled_mip` | float | Scaled MIP value (z-score) normalized by cross-correlation mean and variance. -| `correlation_mean` | float | Mean of the cross-correlation values over the entire search space. -| `correlation_variance` | float | Variance of the cross-correlation values over the entire search space. -| `total_correlations` | int | Total number of cross-correlations performs in the search space (number of defocus planes times number of orientations). -| `pos_x` | int | Particle x position (units of pixels) in the statistics maps. Corresponds to the top-left corner of the template. -| `pos_y` | int | Particle y position (units of pixels) in the statistics maps. Corresponds to the top-left corner of the template. -| `pos_x_img` | int | Center of of the particle (x position, units of pixels) in the micrograph. -| `pos_y_img` | int | Center of of the particle (y position, units of pixels) in the micrograph. -| `pos_x_img_angstrom` | float | Center of the particle (x position, in *Angstroms*) in the micrograph. -| `pos_y_img_angstrom` | float | Center of the particle (y position, in *Angstroms*) in the micrograph. -| `phi` | float | The \( \phi \) angle which (in degrees) produced the MIP value. Orientation angles are in the ZYZ format. -| `theta` | float | The \( \theta \) angle (in degrees) which produced the MIP value. Orientation angles are in the ZYZ format. -| `psi` | float | The \( \psi \) angle (in degrees) which produced the MIP value. Orientation angles are in the ZYZ format. -| `relative_defocus` | float | The relative defocus value (in Angstroms) which produced the MIP value. Relative to `defocus_u` and `defocus_v`. -| `defocus_u` | float | Defocus value along the major axis for the **micrograph** (in Angstroms). -| `defocus_v` | float | Defocus value along the minor axis for the **micrograph** (in Angstroms). -| `astigmatism_angle` | float | Angle of the astigmatism (in degrees) for defocus. -| `pixel_size` | float | Pixel size of the micrograph (in Angstroms). -| `voltage` | float | Voltage of the microscope (in kV). -| `spherical_aberration` | float | Spherical aberration of the microscope (in mm). -| `amplitude_contrast_ratio` | float | Amplitude contrast ratio of the microscope. -| `phase_shift` | float | Phase shift of the microscope (in degrees). -| `ctf_B_factor` | float | B-factor of the CTF, in Angstroms^2. -| `micrograph_path` | str | Path to the micrograph searched over in the match template program. Can be used to identify where particles came from when data frames from multiple match template runs are merged. -| `template_path` | str | Path to the template used for the search. -| `mip_path` | str | Path to the saved MIP map. -| `scaled_mip_path` | str | Path to the saved scaled MIP map. -| `psi_path` | str | Path to the saved psi map. -| `theta_path` | str | Path to the saved theta map. -| `phi_path` | str | Path to the saved phi map. -| `defocus_path` | str | Path to the saved defocus map. -| `correlation_average_path` | str | Path to the saved correlation mean map. -| `correlation_variance_path` | str | Path to the saved correlation variance map. +| Column Name | Type | Description | +|-------------------------------|-------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `particle_index` | int | An integer descriptor for ordering picked particles. Useful when operating on multiple dataframes. | +| `mip` | float | Maximum cross-correlation value over all search orientations and relative defocus values. | +| `scaled_mip` | float | Scaled MIP value (z-score) normalized by cross-correlation mean and variance. | +| `correlation_mean` | float | Mean of the cross-correlation values over the entire search space. | +| `correlation_variance` | float | Variance of the cross-correlation values over the entire search space. | +| `total_correlations` | int | Total number of cross-correlations performs in the search space (number of defocus planes times number of orientations). | +| `pos_x` | int | Particle x position (units of pixels) in the statistics maps. Corresponds to the top-left corner of the template. | +| `pos_y` | int | Particle y position (units of pixels) in the statistics maps. Corresponds to the top-left corner of the template. | +| `pos_x_img` | int | Center of of the particle (x position, units of pixels) in the micrograph. | +| `pos_y_img` | int | Center of of the particle (y position, units of pixels) in the micrograph. | +| `pos_x_img_angstrom` | float | Center of the particle (x position, in *Angstroms*) in the micrograph. | +| `pos_y_img_angstrom` | float | Center of the particle (y position, in *Angstroms*) in the micrograph. | +| `phi` | float | The \( \phi \) angle which (in degrees) produced the MIP value. Orientation angles are in the ZYZ format. | +| `theta` | float | The \( \theta \) angle (in degrees) which produced the MIP value. Orientation angles are in the ZYZ format. | +| `psi` | float | The \( \psi \) angle (in degrees) which produced the MIP value. Orientation angles are in the ZYZ format. | +| `relative_defocus` | float | The relative defocus value (in Angstroms) which produced the MIP value. Relative to `defocus_u` and `defocus_v`. | +| `defocus_u` | float | Defocus value along the major axis for the **micrograph** (in Angstroms). | +| `defocus_v` | float | Defocus value along the minor axis for the **micrograph** (in Angstroms). | +| `astigmatism_angle` | float | Angle of the astigmatism (in degrees) for defocus. | +| `pixel_size` | float | Pixel size of the micrograph (in Angstroms). | +| `voltage` | float | Voltage of the microscope (in kV). | +| `spherical_aberration` | float | Spherical aberration of the microscope (in mm). | +| `amplitude_contrast_ratio` | float | Amplitude contrast ratio of the microscope. | +| `phase_shift` | float | Phase shift of the microscope (in degrees). | +| `ctf_B_factor` | float | B-factor of the CTF, in Angstroms^2. | +| `micrograph_path` | str | Path to the micrograph searched over in the match template program. Can be used to identify where particles came from when data frames from multiple match template runs are merged. | +| `template_path` | str | Path to the template used for the search. | +| `mip_path` | str | Path to the saved MIP map. | +| `scaled_mip_path` | str | Path to the saved scaled MIP map. | +| `psi_path` | str | Path to the saved psi map. | +| `theta_path` | str | Path to the saved theta map. | +| `phi_path` | str | Path to the saved phi map. | +| `defocus_path` | str | Path to the saved defocus map. | +| `correlation_average_path` | str | Path to the saved correlation mean map. | +| `correlation_variance_path` | str | Path to the saved correlation variance map. | + +!!! note "Path columns when using the HDF5 result back-end" + + When `match_template` is configured with `MatchTemplateResultHDF5` rather than `MatchTemplateResultMRC`, all of the `*_path` columns above point to the *same* single `.h5` file rather than eight distinct files. + See [Match Template Result Formats](data_formats/match_template_results.md) for details. + +## Particle stacks + +A `ParticleStack` collects the information needed to re-extract and re-score individual particles identified above — locations, orientations, defocus values, and references to the source micrograph/statistics maps — for use by `refine_template`, `constrained_search`, and the peak/frame inspection programs. +As with match template results, particle stacks support two storage back-ends: `ParticleStackCSV` (the default, built directly from the match template/refine template DataFrame above) and `ParticleStackHDF5` (bundling the particle table, and optionally the extracted particle images and per-particle statistics, into a single portable file). +See [Particle Stack Formats](data_formats/particle_stack.md) for a full comparison of the two back-ends and code examples for constructing, converting, and loading each. ## Data from the refine template program @@ -106,20 +121,77 @@ Note that refined results can be re-refined, for example with a slightly differe When using the results from a refine template run, say to construct a particle stack, columns with the `refined_` prefix will be used by default, if they are present. ### Refine template **DataFrame** + The program outputs another DataFrame with additional columns for the refined orientations, defocus values, and positions. New columns with descriptions are listed below: -| Column Name | Type | Description | -|-------------------------------|-------|-------------| -| `refined_mip` | float | New maximum cross-correlation over refinement search space. -| `refined_scaled_mip` | float | New scaled MIP value (z-score) normalized by cross-correlation mean and variance. -| `refined_pos_x` | int | The refined x position of the particle, top-left corner of the template. -| `refined_pos_y` | int | The refined y position of the particle, top-left corner of the template. -| `refined_pos_x_img` | int | The refined x position of the particle, center of the particle in the micrograph. -| `refined_pos_y_img` | int | The refined y position of the particle, center of the particle in the micrograph. -| `refined_pos_x_img_angstrom` | float | The refined x position of the particle, center of the particle in the micrograph (in Angstroms). -| `refined_pos_y_img_angstrom` | float | The refined y position of the particle, center of the particle in the micrograph (in Angstroms). -| `refined_phi` | float | The refined \( \phi \) angle (in degrees). -| `refined_theta` | float | The refined \( \theta \) angle (in degrees). -| `refined_psi` | float | The refined \( \psi \) angle (in degrees). -| `refined_relative_defocus` | float | The refined relative defocus value (in Angstroms). \ No newline at end of file +| Column Name | Type | Description | +|-------------------------------|-------|---------------------------------------------------------------------------------------------------| +| `refined_mip` | float | New maximum cross-correlation over refinement search space. | +| `refined_scaled_mip` | float | New scaled MIP value (z-score) normalized by cross-correlation mean and variance. | +| `refined_pos_x` | int | The refined x position of the particle, top-left corner of the template. | +| `refined_pos_y` | int | The refined y position of the particle, top-left corner of the template. | +| `refined_pos_x_img` | int | The refined x position of the particle, center of the particle in the micrograph. | +| `refined_pos_y_img` | int | The refined y position of the particle, center of the particle in the micrograph. | +| `refined_pos_x_img_angstrom` | float | The refined x position of the particle, center of the particle in the micrograph (in Angstroms). | +| `refined_pos_y_img_angstrom` | float | The refined y position of the particle, center of the particle in the micrograph (in Angstroms). | +| `refined_phi` | float | The refined \( \phi \) angle (in degrees). | +| `refined_theta` | float | The refined \( \theta \) angle (in degrees). | +| `refined_psi` | float | The refined \( \psi \) angle (in degrees). | +| `refined_relative_defocus` | float | The refined relative defocus value (in Angstroms). | + +## Data from peak & frame inspection + +Where `refine_template` reduces its local search to a single best-scoring hypothesis per particle, the `inspect_peaks`/`frame_inspection` programs (see [Peak & Frame Inspection](programs/inspect_peaks.md)) return the *entire* grid of local scores. +Because this is a dense tensor rather than a sparse table, results are saved as a single self-describing `.npz` file (via `leopard_em.analysis.save_inspection_result`/`load_inspection_result`) instead of a DataFrame or MRC/HDF5 statistics maps. + +### Score tensor shape + +The main array (`InspectionResult.scores`) has one of two shapes depending on `output_mode`: + +- `"cross_correlation"` mode: `(N, n_px, n_defocus, n_orient, H, W)` — local cross-correlation maps (valid-mode, same shapes as described in [the note on correlation modes](#a-note-on-correlation-modes-and-output-shapes)) for every searched hypothesis. +- `"frc"` mode: `(N, n_px, n_defocus, n_orient, n_freq)` — local Fourier ring correlation spectra instead of full 2-D maps, which is far more compact for large searches. + +Per-frame inspection (`FrameInspectionManager`) inserts an additional `frame` axis immediately after the particle axis: `(N, T, n_px, n_defocus, n_orient, H, W)` or `(N, T, n_px, n_defocus, n_orient, n_freq)` for `T` movie frames. + +| Axis | Meaning | +|--------------|--------------------------------------------------------------------------------------------------------| +| `N` | Particle (row in the source particle stack). | +| `T` | Movie frame index (per-frame inspection only). | +| `n_px` | Pixel-size offset index — indexes `pixel_size_offsets`. | +| `n_defocus` | Relative defocus offset index — indexes `defocus_offsets`. | +| `n_orient` | Local Euler-angle offset index (phi, theta, psi) — indexes `euler_angle_offsets`. | +| `H, W` | Valid-mode spatial dimensions of the local cross-correlation map (`"cross_correlation"` mode only). | +| `n_freq` | FRC frequency bin (`"frc"` mode only). | + +The `.npz` file's `axes` metadata field always lists these labels in order. + +### `.npz` file contents` + +Besides `scores`, the file stores everything needed to interpret it without separately tracking how the run was configured: + +| Array | Shape | Description | +|-----------------------|-------------------|---------------------------------------------------------------------------------------------------------------------------------------------------| +| `scores` | see above | The main score tensor. | +| `euler_angle_offsets` | `(n_orient, 3)` | ZYZ orientation offsets searched, in degrees, relative to each particle's base orientation. | +| `defocus_offsets` | `(n_defocus,)` | Relative defocus offsets searched (Angstroms). | +| `pixel_size_offsets` | `(n_px,)` | Relative pixel-size offsets searched. | +| `base_euler_angles` | `(N, 3)` | Per-particle base ZYZ angles the offsets are relative to. | +| `base_defocus` | `(N, 3)` | Per-particle base astigmatic defocus `(defocus_u, defocus_v, defocus_angle)` the offsets are relative to. | +| `particle_index` | `(N,)` | Optional; maps tensor rows back to the source particle stack's `particle_index` column, if present. | +| `frequency_bins` | `(n_freq,)` | Optional; only present in `"frc"` mode. | +| `frame_index` | `(T,)` | Optional; only present for per-frame inspection results — the movie frame index for each entry of the `frame` axis. | +| `metadata_json` | — | JSON blob with `format_version` (currently `1`), `output_mode`, `axes`, `per_frame`, and any extra run metadata (e.g. `correlation_batch_size`). | + +```python +from leopard_em.analysis import load_inspection_result + +result = load_inspection_result("results_inspect_peaks.npz") +result.scores # main tensor; see result.axes for per-dimension labels +result.axes # e.g. ("particle", "pixel_size", "defocus", "orientation", "y", "x") +result.particle_index # maps tensor rows back to the particle stack, if available +``` + +!!! note "Versioned, backwards-incompatible format" + + The `.npz` layout is tracked by `metadata["format_version"]` (currently `1`), bumped whenever the on-disk layout changes in a backwards-incompatible way. diff --git a/docs/data_formats/match_template_results.md b/docs/data_formats/match_template_results.md new file mode 100644 index 00000000..9bb21e79 --- /dev/null +++ b/docs/data_formats/match_template_results.md @@ -0,0 +1,163 @@ +--- +title: Match Template Result Formats +description: MRC vs HDF5 back-ends for match template statistics maps, and the CorrelationTable +--- + +## Match template result formats + +The `match_template` program supports two storage back-ends for its output statistics maps, plus an optional sparse `CorrelationTable` of every detection which crossed a correlation threshold during the search. +This page describes both result back-ends, when to reach for each one, and the `CorrelationTable` format, with minimal code examples for exporting and loading each. + +!!! note "Choosing a back-end" + + The back-end is selected by which class you instantiate for the `match_template_result` field: `MatchTemplateResultMRC` or `MatchTemplateResultHDF5`. + Both are subclasses of a shared (non-instantiable) base class and expose the same in-memory tensor attributes (`mip`, `scaled_mip`, `correlation_average`, `correlation_variance`, `orientation_psi`, `orientation_theta`, `orientation_phi`, `relative_defocus`) — only how they're read from and written to disk differs. + `MatchTemplateResult` remains available as a backward-compatible alias for `MatchTemplateResultMRC`. + +### MRC back-end (`MatchTemplateResultMRC`) + +This is the _only_ Leopard-EM behavior for versions ``<=v1.2``: each of the eight statistics maps is written as its own [MRC format](https://www.ccpem.ac.uk/mrc-format/mrc2014/) file. + +```yaml +match_template_result: + allow_file_overwrite: true + mip_path: ./output_mip.mrc + scaled_mip_path: ./output_scaled_mip.mrc + correlation_average_path: ./output_correlation_average.mrc + correlation_variance_path: ./output_correlation_variance.mrc + orientation_psi_path: ./output_orientation_psi.mrc + orientation_theta_path: ./output_orientation_theta.mrc + orientation_phi_path: ./output_orientation_phi.mrc + relative_defocus_path: ./output_relative_defocus.mrc +``` + +**Use the MRC back-end when...** + +- You want to open individual statistics maps directly in other cryo-EM tools (IMOD, ChimeraX, RELION, etc.) without going through Leopard-EM or Python. +- You're following an existing pipeline or set of scripts built around per-statistic MRC paths (e.g. the columns documented in [Program Output Formats](../data_formats.md#match-template-dataframe)). +- You want to inspect or overwrite a single statistic (say, just the scaled MIP) without touching the others. + +```python +from leopard_em.pydantic_models.results import MatchTemplateResultMRC + +# Loading an existing result back into memory +result = MatchTemplateResultMRC( + mip_path="./output_mip.mrc", + scaled_mip_path="./output_scaled_mip.mrc", + correlation_average_path="./output_correlation_average.mrc", + correlation_variance_path="./output_correlation_variance.mrc", + orientation_psi_path="./output_orientation_psi.mrc", + orientation_theta_path="./output_orientation_theta.mrc", + orientation_phi_path="./output_orientation_phi.mrc", + relative_defocus_path="./output_relative_defocus.mrc", +) +result.load_tensors_from_paths() + +# result.mip, result.scaled_mip, ... are now populated torch.Tensor attributes +peaks_df = result.peaks_to_dataframe() +``` + +### HDF5 back-end (`MatchTemplateResultHDF5`) + +`MatchTemplateResultHDF5` bundles all eight statistics maps, plus run metadata (`leopard_em_version`, `total_projections`, `total_orientations`, `total_defocus`), into a single `.h5` file. + +```yaml +match_template_result: + allow_file_overwrite: true + hdf5_path: ./match_template_output.h5 + compress: true # gzip level 4; set false for faster writes at the cost of file size +``` + +**Use the HDF5 back-end when...** + +- You want a single self-contained, portable result file per micrograph instead of eight — simpler to move, archive, or upload alongside a `CorrelationTable`. +- You're processing many micrographs and want to keep the number of output files per run manageable, especially if you're running in a HPC environment. +- You want run metadata (`leopard_em_version`, total search-space sizes) recorded alongside the statistics maps rather than tracked separately. + +### HDF5 file layout + +```text +/ (root) +│ attrs: leopard_em_version, total_projections, +│ total_orientations, total_defocus +└─ tensors/ + mip float32, shape (H-h+1, W-w+1), gzip-4 (if compress=True) + scaled_mip float32, shape (H-h+1, W-w+1), gzip-4 + correlation_average float32, shape (H-h+1, W-w+1), gzip-4 + correlation_variance float32, shape (H-h+1, W-w+1), gzip-4 + orientation_psi float32, shape (H-h+1, W-w+1), gzip-4 + orientation_theta float32, shape (H-h+1, W-w+1), gzip-4 + orientation_phi float32, shape (H-h+1, W-w+1), gzip-4 + relative_defocus float32, shape (H-h+1, W-w+1), gzip-4 +``` + +where `(H, W)` is the original micrograph size and `(h, w)` is the projected template size, both in units of pixels. + +```python +from leopard_em.pydantic_models.results import MatchTemplateResultHDF5 + +# Loading an existing result back into memory +result = MatchTemplateResultHDF5.from_hdf5("./match_template_output.h5") + +# result.mip, result.scaled_mip, ... are populated torch.Tensor attributes, +# same as the MRC back-end +peaks_df = result.peaks_to_dataframe() +``` + +#### Effect on the match template DataFrame + +When using `MatchTemplateResultHDF5`, the `*_path` columns in the [match template DataFrame](../data_formats.md#match-template-dataframe) (`mip_path`, `scaled_mip_path`, `psi_path`, `theta_path`, `phi_path`, `defocus_path`, `correlation_average_path`, `correlation_variance_path`) all point to the **same** `hdf5_path` instead of eight distinct file paths. + +## Correlation table (sparse detections) + +Where the statistics maps above only retain the _best_ value at each `(x, y)` position, an instance of a `CorrelationTable` object records every search index (defocus offset x out-of-plane orientation x in-plane orientation) whose cross-correlation exceeded a configured threshold, anywhere in the search space. +This is useful for downstream analysis of near-threshold or secondary peaks that don't show up in the per-pixel best-statistic maps — for example, distinguishing a single strong detection from several correlated-but-weaker hypotheses at nearby search indices. + +By default, `MatchTemplateManager.run_match_template(...)` computes a `CorrelationTable` for every run (`compute_correlation_table=True`) and stores it on `match_template_result.correlation_table`; pass `compute_correlation_table=False` to skip this and leave it empty. +Computing and storing the correlation table does incur a few percent overhead in total runtime. + +### Correlation table HDF5 layout + +`CorrelationTable` always uses its own HDF5 format for on-disk storage (independent of which `MatchTemplateResult` back-end you're using): + +```text +/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 +``` + +### Exporting and loading a correlation table + +!!! warning "Correlation table export is currently MRC-specific" + + `export_correlation_table()` and `load_correlation_table_from_path()` — which read/write via the `correlation_table_path` field — are only implemented on `MatchTemplateResultMRC`. + `MatchTemplateResultMRC.export_results()` calls `export_correlation_table()` automatically, so setting `correlation_table_path` in your MRC-backed config is enough. + For `MatchTemplateResultHDF5`, `export_results()` does **not** currently export the correlation table automatically — call `.to_hdf5(...)` on the `CorrelationTable` directly, as shown below. + +```python +# MRC back-end: correlation_table_path is exported/loaded automatically +mrc_result.correlation_table_path = "./output_correlation_table.h5" +mrc_result.export_results() # writes the 8 mrc files AND the correlation table +mrc_result.load_correlation_table_from_path() + +# HDF5 back-end: export the correlation table explicitly (not done by export_results) +hdf5_result.export_results() # writes hdf5_path only +hdf5_result.correlation_table.to_hdf5("./output_correlation_table.h5") + +# Loading a correlation table directly, regardless of back-end +from leopard_em.pydantic_models.results import CorrelationTable + +table = CorrelationTable.from_hdf5("./output_correlation_table.h5") +table_df = table.to_dataframe() # one row per detection +``` + +In both cases, `mrc_result`/`hdf5_result` are `MatchTemplateResult*` instances populated by `MatchTemplateManager.run_match_template(...)` (with `compute_correlation_table=True`, the default) — `match_template_result.correlation_table` is set automatically after the run. diff --git a/docs/data_formats/particle_stack.md b/docs/data_formats/particle_stack.md new file mode 100644 index 00000000..bb386aef --- /dev/null +++ b/docs/data_formats/particle_stack.md @@ -0,0 +1,183 @@ +--- +title: Particle Stack Formats +description: CSV vs HDF5 back-ends for particle stacks used by refine template, constrained search, and inspect peaks +--- + +## Particle stack formats + +A `ParticleStack` collects everything needed to re-extract and re-score individual particles identified by `match_template` — their locations, orientations, defocus values, and references to the source micrograph and statistics maps — for use by `refine_template`, `constrained_search`, and the peak/frame inspection programs. +As with match template results, particle stacks support two storage back-ends. + +!!! note "Choosing a back-end" + + Both back-ends are subclasses of a shared (non-instantiable) base class and expose the same in-memory API (`get_euler_angles()`, `get_relative_defocus()`, `construct_image_stack(...)`, etc.) — only how the particle table (and optionally particle images) are read from and written to disk differs. + `ParticleStack` remains available as a backward-compatible alias for `ParticleStackCSV`. + +!!! note "`refine_template` and `constrained_search` output matches the input back-end by default" + + `RefineTemplateManager.particle_stack`, `OptimizeTemplateManager.particle_stack`, and `ConstrainedSearchManager.particle_stack_reference`/`particle_stack_constrained` are all typed as `ParticleStackCSV | ParticleStackHDF5`, so a `ParticleStackHDF5` can be passed in directly wherever a particle stack input is required. + + `ConstrainedSearchManager.export_results(...)` additionally writes two small CSV sibling tables — `_parameters.csv` (search parameters and the false-positive threshold) and `_above_threshold.csv` (rows above that threshold) — derived from the base of `output_dataframe_path` regardless of the main table's `output_format`. + + `OptimizeTemplateManager`'s optional per-pixel-size diagnostic dumps (`write_individual_csv`) are intermediate debug artifacts, not the program's primary output, and remain hardcoded to CSV. + + To convert an existing CSV-backed refined table to HDF5 after the fact, use `ParticleStackCSV.to_hdf5(...)` as a separate post-processing step (see below), or see [exporting results](#exporting-refined-results) for the general-purpose helper. + +### Exporting refined results + +`RefineTemplateManager.export_results(...)` and `ConstrainedSearchManager.export_results(...)` build the refined particle table and write it to disk, then **return the newly-written particle stack instance** (a `ParticleStackCSV` or `ParticleStackHDF5`, matching whichever `output_format` was used) — reuse it directly as input to the next program in a pipeline instead of re-reading it from disk: + +```python +result_stack = refine_manager.export_results( + output_dataframe_path="/some/path/to/refined.h5", + result=refine_result, + output_format="hdf5", # omit to match the input particle_stack's back-end +) +# result_stack is a ready-to-use ParticleStackHDF5 — e.g. feed it straight into +# ConstrainedSearchManager(particle_stack_reference=result_stack, ...) +``` + + +### CSV back-end (`ParticleStackCSV`) + +This is the only Leopard-EM behavior for versions ``<=v1.2``: the particle table is a CSV file (the same DataFrame [written by `match_template` or `refine_template`](../data_formats.md#match-template-dataframe)), and particle images are extracted on demand from the micrograph/statistics-map paths referenced in each row. + +```yaml +particle_stack: + df_path: /some/path/to/particles.csv + extracted_box_size: [528, 528] + original_template_size: [512, 512] +``` + +**Use the CSV back-end when...** + +- You're feeding the direct output of `match_template` or `refine_template` into the next program in the pipeline — this is the default hand-off format documented on the [refine template](../programs/refine_template.md#particle-stack-of-particles-to-refine) and [constrained search](../programs/constrained_search.md) program pages. +- You want to inspect or edit particle metadata as a plain-text/CSV table (e.g. in a spreadsheet or with `pandas`) without unpacking an HDF5 file. +- Your source micrographs and MRC statistics maps are expected to stay available at their original paths — the CSV back-end re-reads them each time, so it stays in sync with those files rather than freezing a snapshot. + +```python +from leopard_em.pydantic_models.data_structures import ParticleStackCSV + +particle_stack = ParticleStackCSV( + df_path="/some/path/to/particles.csv", + extracted_box_size=(528, 528), + original_template_size=(512, 512), +) + +# Load the (deduplicated) referenced micrographs, then extract per-particle boxes +images, indices = particle_stack.load_images_grouped_by_column("micrograph_path") +image_stack = particle_stack.construct_image_stack( + images=images, + indices=indices, + extraction_size=particle_stack.extracted_box_size, +) +``` + +### HDF5 back-end (`ParticleStackHDF5`) + +`ParticleStackHDF5` stores the particle table in a single `.h5` file, and can optionally bundle the extracted particle images (`image_stack`) and/or per-particle local statistic maps (`local_stats`) directly into that same file, so the stack no longer depends on the original micrograph/statistics-map files being available at their recorded paths. + +```yaml +particle_stack: + hdf5_path: /some/path/to/particles.h5 + extracted_box_size: [528, 528] + original_template_size: [512, 512] +``` + +**Use the HDF5 back-end when...** + +- You want a fully portable, self-contained particle stack — one file you can archive, share, or move to another machine without also shipping every referenced micrograph and statistics map. +- You want per-particle local statistic maps stored alongside the particle table rather than recomputed. `local_stats` is a `dict[str, torch.Tensor]` keyed by `*_path` column name — any subset of `mip_path`, `scaled_mip_path`, `psi_path`, `theta_path`, `phi_path`, `defocus_path`, `correlation_average_path`, `correlation_variance_path` (or all of them) can be stored, not just correlation average/variance. + +#### Two loading modes + +`ParticleStackHDF5` supports two mutually exclusive modes, controlled by the `image_stack_stored`/`local_stats_stored` flags (set automatically by `to_hdf5(...)`, and read back by `from_hdf5(...)`): + +- **Load from referenced files** (`image_stack_stored=False`): the HDF5 file stores only the particle table; `image_stack`/`local_stats` are computed on demand from the micrograph/statistics-map paths in that table, same as the CSV back-end. +- **Load from HDF5 directly** (`image_stack_stored=True` and/or `local_stats_stored=True`): `image_stack`/`local_stats` are read from the HDF5 datasets, no access to the original micrograph files needed. + +To populate `local_stats` before writing, use `get_local_stat_maps(...)` (extracts the valid cross-correlation region around each particle for any `*_path` column) and assign the result: + +```python +particle_stack.local_stats.update(particle_stack.get_local_stat_maps()) +# or a specific subset: +particle_stack.local_stats.update( + particle_stack.get_local_stat_maps(columns=["mip_path", "correlation_average_path"]) +) +particle_stack.to_hdf5(include_local_stats=True) +``` + +#### HDF5 file layout + +```text +/ (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] + (N, valid_h, valid_w) float32 + ... -- one dataset per entry in `local_stats` at write + time, named after its column (e.g. `mip_path`, + `correlation_average_path`) +``` + +where `valid_h = extracted_box_size[0] - original_template_size[0] + 1` and `valid_w = extracted_box_size[1] - original_template_size[1] + 1` (see the [note on correlation modes](../data_formats.md#a-note-on-correlation-modes-and-output-shapes)). + +### Converting a CSV-backed stack to HDF5 + +`ParticleStackCSV.to_hdf5(...)` is the recommended migration path. +It re-uses the CSV back-end's already-configured extraction settings, generates a `particle_id` for each row, and writes a new `ParticleStackHDF5`: + +```python +from leopard_em.pydantic_models.data_structures import ParticleStackCSV + +csv_stack = ParticleStackCSV( + df_path="/some/path/to/particles.csv", + extracted_box_size=(528, 528), + original_template_size=(512, 512), +) + +# Extract particle images before baking them into the HDF5 file +images, indices = csv_stack.load_images_grouped_by_column("micrograph_path") +csv_stack.image_stack = csv_stack.construct_image_stack( + images=images, + indices=indices, + extraction_size=csv_stack.extracted_box_size, +) + +hdf5_stack = csv_stack.to_hdf5( + "/some/path/to/particles.h5", + include_image_stack=True, # requires image_stack to already be populated + include_local_stats=False, +) +``` + +### Loading an existing HDF5 particle stack + +```python +from leopard_em.pydantic_models.data_structures import ParticleStackHDF5 + +particle_stack = ParticleStackHDF5.from_hdf5("/some/path/to/particles.h5") + +if particle_stack.image_stack_stored: + image_stack = particle_stack.image_stack +else: + images, indices = particle_stack.load_images_grouped_by_column("micrograph_path") + image_stack = particle_stack.construct_image_stack( + images=images, + indices=indices, + extraction_size=particle_stack.extracted_box_size, + ) +``` diff --git a/docs/index.md b/docs/index.md index 04b21750..66772e46 100644 --- a/docs/index.md +++ b/docs/index.md @@ -11,19 +11,20 @@ Leopard-EM is a Python implementation of Two-Dimensional Template Matching (2DTM !!! note "Citing this work" - If you use Leopard-EM in your research, please cite the [Leopard-EM preprint](https://doi.org/10.1101/2025.08.26.672452): + If you use Leopard-EM in your research, please cite the [Leopard-EM Manuscript](https://doi.org/10.1107/S2059798325009982): ``` - @article {Giammar2025.08.26.672452, - author = {Giammar, Matthew David and Dickerson, Joshua Luke and Hall, Laina Nicole and Lucas, Bronwyn Ayla}, - title = {Leopard-EM: An extensible 2DTM package to accelerate in situ structural biology}, - elocation-id = {2025.08.26.672452}, - year = {2025}, - doi = {10.1101/2025.08.26.672452}, - publisher = {Cold Spring Harbor Laboratory}, - URL = {https://www.biorxiv.org/content/early/2025/08/29/2025.08.26.672452}, - eprint = {https://www.biorxiv.org/content/early/2025/08/29/2025.08.26.672452.full.pdf}, - journal = {bioRxiv} + @article{Giammar:rr5256, + author = "Giammar, Matthew D. and Dickerson, Joshua L. and Hall, Laina N. and Lucas, Bronwyn A.", + title = "{{\it Leopard-EM}: an extensible 2D template-matching package to accelerate {\it in situ} structural biology}", + journal = "Acta Crystallographica Section D", + year = "2026", + volume = "82", + number = "1", + pages = "30--42", + month = "Jan", + doi = {10.1107/S2059798325009982}, + url = {https://doi.org/10.1107/S2059798325009982}, } ``` @@ -33,7 +34,7 @@ Leopard-EM is a Python implementation of Two-Dimensional Template Matching (2DTM The general system requirements for Leopard-EM are -- Python version 3.10 or above +- Python version 3.11 or above - PyTorch 2.4.0 or above - Linux operating system diff --git a/docs/programs/constrained_search.md b/docs/programs/constrained_search.md index b69bf911..1fcf0485 100644 --- a/docs/programs/constrained_search.md +++ b/docs/programs/constrained_search.md @@ -87,5 +87,8 @@ orientation_refinement_config: These should be the same as for [Match Template](../programs/match_template.md). +## Running the constrained search program +`ConstrainedSearchManager.run_constrained_search(output_dataframe_path, ...)` writes the main refined particle table to `output_dataframe_path`, matching the back-end of `particle_stack_reference` by default (`ParticleStackCSV` or `ParticleStackHDF5`), or pass `output_format="csv"`/`"hdf5"` to override. It also writes two small, always-CSV sibling tables next to it — `_parameters.csv` (search parameters and the false-positive threshold) and `_above_threshold.csv` (rows above that threshold), named from the base of `output_dataframe_path` regardless of the main table's format. +`run_constrained_search(...)` returns the refined particle stack it wrote for the main table (a `ParticleStackCSV` or `ParticleStackHDF5`), ready to reuse directly without re-reading it from disk. See [exporting refined results](../data_formats/particle_stack.md#exporting-refined-results) for more detail. diff --git a/docs/programs/inspect_peaks.md b/docs/programs/inspect_peaks.md new file mode 100644 index 00000000..8b567049 --- /dev/null +++ b/docs/programs/inspect_peaks.md @@ -0,0 +1,132 @@ +--- +title: The Peak & Frame Inspection Programs +description: Inspecting the full local search grid around identified particles, optionally per movie frame +--- + +# Peak & frame inspection + +`refine_template` searches a local grid of orientation, defocus, and pixel-size offsets around each particle and keeps only the single best-scoring hypothesis. +The `inspect_peaks` and `frame_inspection` programs run that same local search but **skip the best-peak reduction**, returning the full grid of local scores for every particle instead. +This is useful for diagnosing *why* a particle refined the way it did — for example, visualizing how sharply peaked the correlation is around the refined orientation, or checking a particle's stability across individual movie frames before trusting its refined pose. + +!!! note "Same configuration schema as `refine_template`" + + Both `PeakInspectionManager` and `FrameInspectionManager` are subclasses of `RefineTemplateManager` and reuse its exact YAML schema — see the [refine template program details](refine_template.md) for the meaning of `particle_stack`, `defocus_refinement_config`, `orientation_refinement_config`, `pixel_size_refinement_config`, and `preprocessing_filters`. + Whatever grid those refinement configs define is the grid that gets scored and returned in full, rather than reduced to a single best value. + +!!! warning "Output tensors can be very large" + + The full local score grid is `(N, n_px, n_defocus, n_orient, H, W)` in `"cross_correlation"` mode — this can get large quickly for big particle stacks or fine-grained refinement configs. + Use `"frc"` output mode (below) for a much more compact result, and/or subset the particle stack in your config for large runs. + There is currently no automatic chunking of the particle stack to manage memory. + +## Peak inspection (`PeakInspectionManager`) + +`PeakInspectionManager` runs the refine-template backend without best-peak reduction, returning the score at every searched hypothesis for every particle. + +A default config file is available [here on the GitHub page](https://raw.githubusercontent.com/Lucaslab-Berkeley/Leopard-EM/refs/heads/main/programs/inspect_peaks/inspect_peaks_example_config.yaml) — it is identical in structure to the [refine template example config](https://raw.githubusercontent.com/Lucaslab-Berkeley/Leopard-EM/refs/heads/main/programs/refine_template/refine_template_example_config.yaml). + +### Running peak inspection + +We provide an example script, [`Leopard-EM/programs/inspect_peaks/run_inspect_peaks.py`](https://github.com/Lucaslab-Berkeley/Leopard-EM/blob/main/programs/inspect_peaks/run_inspect_peaks.py), which loads a config, runs peak inspection, and saves the result to a self-describing `.npz` file. +Edit the constants near the top of the script: + +```python +YAML_CONFIG_PATH = "/path/to/inspect_peaks_configuration.yaml" +OUTPUT_PATH = "/path/to/results_inspect_peaks.npz" +CORRELATION_BATCH_SIZE = 32 # lower if you run out of GPU memory +OUTPUT_MODE = "cross_correlation" # or "frc" +``` + +Or drive it directly from Python: + +```python +from leopard_em.pydantic_models.managers import PeakInspectionManager + +manager = PeakInspectionManager.from_yaml("/path/to/inspect_peaks_configuration.yaml") +output_path = manager.run_and_save_peak_inspection( + output_path="/path/to/results_inspect_peaks.npz", + correlation_batch_size=32, + prefer_refined_angles=True, + output_mode="cross_correlation", # or "frc" +) +``` + +`run_peak_inspection(...)` is also available if you'd rather get the raw tensor (and, in `"frc"` mode, the frequency bins) back in memory without writing a file. + +### Output modes + +- `"cross_correlation"` (default) — returns the full local cross-correlation map for every hypothesis, shape `(N, n_px, n_defocus, n_orient, H, W)`. +- `"frc"` — returns local Fourier ring correlation spectra instead of full 2-D maps, shape `(N, n_px, n_defocus, n_orient, n_freq)`, which is far more compact for large searches. + +See [Data from peak & frame inspection](../data_formats.md#data-from-peak-frame-inspection) for the full `.npz` file layout and how to load results back with `leopard_em.analysis.load_inspection_result`. + +## Per-frame peak inspection (`FrameInspectionManager`) + +`FrameInspectionManager` extends `PeakInspectionManager` to score each movie frame **independently**, rather than the motion-corrected sum used by `match_template`/`refine_template`/`inspect_peaks`. +This is useful for checking how a particle's local correlation landscape evolves across a movie's exposure — for example, to spot particles that only correlate strongly in a subset of frames. + +### Configuring the movie + +Per-frame inspection requires `movie_config.enabled: true`, plus either a per-frame deformation field or an explicit per-particle shifts CSV to align each particle box to the correct position in every frame: + +```yaml +movie_config: + enabled: true + movie_path: /some/path/to/aligned_or_unaligned_movie.mrc + # Provide exactly one of the following two: + deformation_field_path: /some/path/to/deformation_grid.csv + # particle_shifts_path: /some/path/to/particle_shifts.csv + pre_exposure: 0.0 + fluence_per_frame: 1.0 +``` + +`particle_shifts_path`, if provided, takes precedence over `deformation_field_path`. +It should be a CSV with columns `particle_index`, `frame`, `y_shift`, `x_shift`. + +A default config file is available [here on the GitHub page](https://raw.githubusercontent.com/Lucaslab-Berkeley/Leopard-EM/refs/heads/main/programs/inspect_peaks/frame_inspection_example_config.yaml). + +### Optional per-frame template dose weighting + +Setting `apply_template_dose_weighting=True` applies cumulative electron-dose filtering to the (non-dose-weighted) template separately for each frame's exposure interval, using `pre_exposure`/`fluence_per_frame` from `movie_config` — this accounts for radiation damage accumulating over the course of the movie when scoring later frames. +It is off by default. + +### Running per-frame peak inspection + +We provide an example script, [`Leopard-EM/programs/inspect_peaks/run_frame_inspection.py`](https://github.com/Lucaslab-Berkeley/Leopard-EM/blob/main/programs/inspect_peaks/run_frame_inspection.py), which mirrors `run_inspect_peaks.py`: + +```python +from leopard_em.pydantic_models.managers import FrameInspectionManager + +manager = FrameInspectionManager.from_yaml("/path/to/frame_inspection_configuration.yaml") +output_path = manager.run_and_save_peak_inspection_per_frame( + output_path="/path/to/results_frame_inspection.npz", + correlation_batch_size=32, + prefer_refined_angles=True, + apply_projection_normalization=True, + output_mode="cross_correlation", # or "frc" + apply_template_dose_weighting=False, +) +``` + +The resulting `.npz` carries an extra `frame` axis immediately after the particle axis — `(N, T, n_px, n_defocus, n_orient, H, W)` in `"cross_correlation"` mode, or `(N, T, n_px, n_defocus, n_orient, n_freq)` in `"frc"` mode, for `T` movie frames — and includes a `frame_index` array mapping that axis to actual movie frame numbers. +See [Data from peak & frame inspection](../data_formats.md#data-from-peak-frame-inspection) for the complete layout. + +## Loading and inspecting results + +Both programs write the same self-describing `.npz` format, loadable with a single helper regardless of which program produced it: + +```python +from leopard_em.analysis import load_inspection_result + +result = load_inspection_result("results_inspect_peaks.npz") + +result.scores # main tensor; see result.axes for per-dimension labels +result.axes # e.g. ("particle", "pixel_size", "defocus", "orientation", "y", "x") +result.euler_angle_offsets # (n_orient, 3) ZYZ offsets per orientation index +result.defocus_offsets # (n_defocus,) relative defocus offsets (Angstroms) +result.pixel_size_offsets # (n_px,) relative pixel-size offsets +result.particle_index # (N,) maps tensor rows back to the particle stack, if available +result.frequency_bins # (n_freq,) FRC frequencies, "frc" mode only +result.frame_index # (T,) movie frame index, per-frame inspection only +``` diff --git a/docs/programs/match_template.md b/docs/programs/match_template.md index de844ae5..1adc99d3 100644 --- a/docs/programs/match_template.md +++ b/docs/programs/match_template.md @@ -180,6 +180,56 @@ computational_config: If you encounter the error `RuntimeError: CUDA error: invalid device ordinal`, then you've probably listed more GPU devices than are on your machine! Check how many GPUs you have (for example with `nvidia-smi`) and update the `gpu_ids` field accordingly. +### Choosing a cross-correlation backend + +The `computational_config.backend` field selects which implementation computes the 2D cross-correlations for each orientation/defocus combination. +Three options are available: + +- `"streamed"` (default) — individual 2D cross-correlations are computed across multiple CUDA streams using PyTorch. +- `"batched"` — all 2D cross-correlations for a batch of orientations are computed in a single batched PyTorch call. +- `"zipfft"` — cross-correlations are computed using the optional [zipFFT](https://github.com/mgiammar/zipFFT) library, which can substantially speed up the FFT-based cross-correlation step over the two PyTorch-based backends. + +```yaml +computational_config: + gpu_ids: [0, 1] + num_cpus: 4 + backend: zipfft +``` + + + +!!! note "zipFFT and the correlation-mode shapes" + + Correlation maps produced by the `zipfft` backend are already cropped to "valid" mode internally (rather than "same" mode, like the other two backends). + Leopard-EM detects this automatically, so the output statistics maps have the same shapes regardless of backend — see the [note on correlation modes](../data_formats.md#a-note-on-correlation-modes-and-output-shapes) for details. + +#### Installing zipFFT + +zipFFT is an **optional** dependency and is distributed as a source-only package with a compiled CUDA/C-extension. It must be built locally against a matching CUDA toolkit. + +**Requirements:** + +- A CUDA-capable GPU with compute capability 8.0 or above. +- CUDA toolkit >=13.0, plus a working C/C++ compile toolchain (`nvcc` and a compatible host compiler) available on your `PATH`. + +**Installation:** + +```bash +git clone https://github.com/mgiammar/zipFFT.git +cd zipFFT +pip install . +``` + +Once installed, `zipfft` is importable from the same Python environment as Leopard-EM, and `backend: zipfft` becomes available in your `computational_config`. See the [zipFFT GitHub page](https://github.com/mgiammar/zipFFT) for additional details on configuring compiled shapes/sizes for your system. + +!!! note "Backend availability is checked at first use, not at config load" + + Leopard-EM does not validate that `zipfft` is importable when the `computational_config` is loaded. + If `backend: zipfft` is configured but the package isn't installed, the run raises a clear `ImportError` ("backend='zipfft' requires the optional 'zipfft' package...") as soon as the first cross-correlation batch is computed. + ## Running the match template program Once you've configured a YAML file, running the match template program is fairly simple. @@ -195,7 +245,6 @@ For example, # Select peaks bases on a number of false positives df = mt_manager.results_to_dataframe(locate_peaks_kwargs={"false_positives": 1.0}) - # Uses a pre-defined z-score cutoff for peak calling df = mt_manager.results_to_dataframe(locate_peaks_kwargs={"z_score_cutoff": 7.8}) ``` @@ -208,7 +257,6 @@ The provided program script will output the statistics maps over the image for t These data can be passed onto downstream analysis, for example the refine template program. More detail about these data and their formats is on the [Leopard-EM data formats page](../data_formats.md). - ## Mathematical description Described succinctly using mathematics, the match template constructs the orientational search space, \( \mathbf{R} = \{ R_1, R_2, \dots, R_n\} \), and relative defocus search space, \( \mathbf{Q} = \{ \Delta f_1, \Delta f_2, \dots, \Delta f_m\} \), to generate the CTF-convolved projections of a reference template: diff --git a/docs/programs/overview.md b/docs/programs/overview.md index 39088400..b3d9da45 100644 --- a/docs/programs/overview.md +++ b/docs/programs/overview.md @@ -3,7 +3,7 @@ title: Overview of Leopard-EM programs description: A basic overview of each of the built-in programs in Leopard-EM --- -# Programs at a glance +## Programs at a glance The Leopard-EM package currently has five main programs which are easily user-configurable through YAML files and are runnable from Python scripts. These five programs encompass a variety of 2DTM data processing workflows. Here, we provide a brief overview of each program's functionality and the necessary input data for that program. Detailed information on configuring and running each program can be found on their respective pages, linked below. @@ -23,7 +23,6 @@ The required inputs (besides config fields) for `match_template` are: * The estimated CTF defocus parameters for that micrograph, and * A simulated 3D reference template (see package [ttsim3d](https://github.com/teamtomo/ttsim3d) for simulating reference templates). - ## Refine Template After the `match_template` program identifies particles from a "coarse" search, the `refine_template` program locally refines the orientation, location, and defocus parameters on a per-particle basis. @@ -35,7 +34,6 @@ The required inputs for `refine_template` are: * A simulated 3D reference template (see package [ttsim3d](https://github.com/teamtomo/ttsim3d) for simulating reference templates), and * The csv file of particle locations and orientations output from the `match_template` program. - ## Optimize Template The 2DTM SNR is extremely sensitive to incorrect pixel size in the reference template structure. @@ -48,7 +46,6 @@ The required inputs for `optimize_template` are: * Simulation configuration for the [ttsim3d](https://github.com/teamtomo/ttsim3d) package including the reference structure, and * The csv file of particle locations and orientations output from either the `match_template` or `refine_template` program. - ## Constrained Search The `constrained_search` program uses pre-identified locations and orientations of a reference particle (lets call it particle A) to constrain the search space for another particle (called particle B); constraining the search space increases the sensitivity of 2DTM. @@ -67,6 +64,17 @@ The required inputs for `constrained_search` are: * A particle stack csv (from `match_template`) for the constrained particle, and * Estimates on the relative position & orientation as well as flexibility of the constrained particle. +## Peak & Frame Inspection + +The `inspect_peaks` and `frame_inspection` programs let you inspect the local search areas which can be useful to investigate how sensitive peaks/detections are to small changes in parameter space. +Another axis to compare against is correlation across movie frames (i.e. how much signal is present after some amount of accumulated electron dose). +The [peak & frame inspection program details](inspect_peaks.md) contains further information on configuring and running these programs and generally follows the `refine_template` program configuration. + +The required inputs for `inspect_peaks`/`frame_inspection` are: + +* A simulated 3D reference template (see package [ttsim3d](https://github.com/teamtomo/ttsim3d) for simulating reference templates), +* The csv file of particle locations and orientations output from the `match_template` or `refine_template` program, and +* (for `frame_inspection` only) The original (unaligned or aligned) movie stack, plus a deformation field or per-particle shifts describing motion across frames. ## Optimize B-Factor script diff --git a/docs/programs/refine_template.md b/docs/programs/refine_template.md index 3b69dda2..16ba5684 100644 --- a/docs/programs/refine_template.md +++ b/docs/programs/refine_template.md @@ -40,6 +40,8 @@ This csv file is written when [running the match template program](match_templat That is, the csv file for `df_path` does not need to come from match template. Running multiple refinement can be useful to compare between similar reference structures using 2DTM. + This CSV-in, CSV-out behavior is the default when the input `particle_stack` is CSV-backed (`df_path`). If the input is HDF5-backed (`hdf5_path`) instead, the output is HDF5 by default too — see [particle stack formats](../data_formats/particle_stack.md#exporting-refined-results) for the `export_results(..., output_format=...)` override. `export_results(...)` returns the refined particle stack it just wrote, so it can be passed straight into another program (e.g. `constrained_search`) without re-reading it from disk. + The next two fields are `extracted_box_size` and `original_template_size` which together are used to extract regions in the image and statistics maps around a particle. Set the `original_template_size` field to the same shape as the simulated volume, that is if the 3D mrc file for the reference template is of shape \( (512, 512, 512) \), then this filed should be `original_template_size: [512, 512]`. @@ -152,5 +154,5 @@ In addition to the YAML configuration path, there are the additional variables ` The latter variable is used to process multiple particles at once since we want to maximize hardware utilization. But this parameter also needs to balance available GPU memory. -The former variable, `DATAFRAME_OUTPUT_PATH`, will write a new particle stack csv file with new columns corresponding to the refined position, orientation, defocus, and pixel size on a per-particle basis. -More details on the particle stack csv format can be found on the [Leopard-EM data formats page](../data_formats.md). +The former variable, `DATAFRAME_OUTPUT_PATH`, will write a new particle stack file with new columns corresponding to the refined position, orientation, defocus, and pixel size on a per-particle basis — as CSV or HDF5, matching the input `particle_stack`'s back-end by default (or override via `output_format` on `run_refine_template(...)`). +More details on the particle stack csv/HDF5 formats can be found on the [Leopard-EM data formats page](../data_formats.md) and the [particle stack formats page](../data_formats/particle_stack.md). diff --git a/mkdocs.yml b/mkdocs.yml deleted file mode 100644 index bde7e320..00000000 --- a/mkdocs.yml +++ /dev/null @@ -1,56 +0,0 @@ -site_name: "Leopard-EM: Two-Dimensional Template Matching in Python" -theme: - name: material - features: - - content.code.copy -plugins: - - search - - autorefs - - mkdocs-jupyter - - inline-select-svg - - mkdocs-autoapi: - autoapi_dir: src - - mkdocstrings: - enabled: !ENV [ENABLE_MKDOCSTRINGS, true] - handlers: - python: - paths: - - src - options: - docstring_style: numpy -markdown_extensions: - - admonition - - pymdownx.details - - pymdownx.superfences - - mdx_math - - footnotes - - tables -extra_javascript: - - https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.4/MathJax.js?config=TeX-AMS-MML_HTMLorMML -nav: - - Home: index.md - - Tutorials: - - 2DTM Introduction: tutorials/match_template_intro.md - - Batch Processing: tutorials/batch_processing.md - - Distributed Computing: tutorials/distributed_match_template.md - - Examples: - - Match Template Config: examples/01_basic_configuration.ipynb - - Peaks to DataFrame: examples/02_extract_peak_info.ipynb - - Comparing Scoring Metrics: examples/03_compare_scoring_metrics.ipynb - - Plotting match template results: examples/04_plotting_2dtm_results.ipynb - - Structure re-projection plot: examples/05_structure_reprojection.ipynb - # - Constrained Search: examples/constrained_search_example.yaml - - Programs: - - Overview: programs/overview.md - - Match Template details: programs/match_template.md - - Refine Template details: programs/refine_template.md - - Optimize Template details: programs/optimize_template.md - - Constrained Search details: programs/constrained_search.md - - Program Output Formats: data_formats.md - - Benchmarks: benchmarks.md - - Common Issues: common_issues.md - # - Analyzing Match Template Results: examples/analyzing_results.md - # - Pydantic Model API: api/pydantic_models.md - # - Backend API: api/backend.md - # - Theory: theory.md - - Contributing: contributing.md \ No newline at end of file diff --git a/programs/constrained_search/run_constrained_search.py b/programs/constrained_search/run_constrained_search.py index 43020882..713379f9 100644 --- a/programs/constrained_search/run_constrained_search.py +++ b/programs/constrained_search/run_constrained_search.py @@ -14,6 +14,11 @@ YAML_CONFIG_PATH = "/path/to/constrained-search-configuration.yaml" # Path to where the dataframe with refined peak parameters will be output. +# May be a .csv or .h5 path -- output format matches the particle_stack_reference +# back-end configured in the YAML file by default (df_path -> CSV, hdf5_path -> +# HDF5), or pass output_format="csv"/"hdf5" to run_constrained_search() to +# override. The accompanying "_parameters" and "_above_threshold" sibling +# tables are always written as CSV regardless of this setting. DATAFRAME_OUTPUT_PATH = "/path/to/constrained-search-results.csv" # Number of particles to refine simultaneously. Will need to tune this parameter diff --git a/programs/inspect_peaks/frame_inspection_example_config.yaml b/programs/inspect_peaks/frame_inspection_example_config.yaml new file mode 100644 index 00000000..e2f2f71e --- /dev/null +++ b/programs/inspect_peaks/frame_inspection_example_config.yaml @@ -0,0 +1,46 @@ +# Example config for per-frame peak inspection (frame correlations). +# Uses the same schema as refine/inspect configs, but movie_config.enabled must be true. +template_volume_path: /some/path/to/template.mrc +particle_stack: + df_path: /some/path/to/particles.csv # Needs to be readable by pandas + extracted_box_size: [512, 512] + original_template_size: [512, 512] +defocus_refinement_config: + enabled: false + defocus_max: 100.0 + defocus_min: -100.0 + defocus_step: 20.0 +orientation_refinement_config: + enabled: false + psi_step_coarse: 1.5 + psi_step_fine: 0.15 + theta_step_coarse: 2.5 + theta_step_fine: 0.25 +pixel_size_refinement_config: + enabled: false + pixel_size_min: -0.005 + pixel_size_max: 0.005 + pixel_size_step: 0.001 +preprocessing_filters: + whitening_filter: + do_power_spectrum: true + enabled: true + max_freq: 0.5 + num_freq_bins: null + bandpass_filter: + enabled: false + falloff: null + high_freq_cutoff: null + low_freq_cutoff: null +computational_config: + gpu_ids: 0 + num_cpus: 1 +apply_global_filtering: true +movie_config: + enabled: true + movie_path: /some/path/to/aligned_or_unaligned_movie.mrc + # Optional: either provide particle_shifts_path OR deformation_field_path. + # deformation_field_path: /some/path/to/deformation_grid.csv + # particle_shifts_path: /some/path/to/particle_shifts.csv + pre_exposure: 0.0 + fluence_per_frame: 1.0 diff --git a/programs/inspect_peaks/inspect_peaks_example_config.yaml b/programs/inspect_peaks/inspect_peaks_example_config.yaml new file mode 100644 index 00000000..58de967d --- /dev/null +++ b/programs/inspect_peaks/inspect_peaks_example_config.yaml @@ -0,0 +1,45 @@ +# Identical to `programs/refine_template/refine_template_example_config.yaml`. +# PeakInspectionManager extends RefineTemplateManager and uses the same fields. +# Call `PeakInspectionManager.from_yaml(path)` to load this file. +template_volume_path: /some/path/to/template.mrc +particle_stack: + df_path: /some/path/to/particles.csv # Needs to be readable by pandas + extracted_box_size: [528, 528] + original_template_size: [512, 512] +defocus_refinement_config: + enabled: true + defocus_max: 100.0 # in Angstroms, relative to "best" defocus value in particle stack dataframe + defocus_min: -100.0 # in Angstroms, relative to "best" defocus value in particle stack dataframe + defocus_step: 20.0 # in Angstroms +orientation_refinement_config: + enabled: true + psi_step_coarse: 1.5 # in degrees + psi_step_fine: 0.15 # in degrees + theta_step_coarse: 2.5 # in degrees + theta_step_fine: 0.25 # in degrees +pixel_size_refinement_config: + enabled: false + pixel_size_min: -0.005 + pixel_size_max: 0.005 + pixel_size_step: 0.001 +preprocessing_filters: + whitening_filter: + do_power_spectrum: true + enabled: true + max_freq: 0.5 # In terms of Nyquist frequency + num_freq_bins: null + bandpass_filter: + enabled: false + falloff: null + high_freq_cutoff: null + low_freq_cutoff: null +computational_config: + gpu_ids: 0 + num_cpus: 1 +apply_global_filtering: true +movie_config: + enabled: false + movie_path: path/to/aligned_movie.mrc + deformation_field_path: path/to/deformation_field.csv + pre_exposure: 0.0 + fluence_per_frame: 1.0 diff --git a/programs/inspect_peaks/run_frame_inspection.py b/programs/inspect_peaks/run_frame_inspection.py new file mode 100644 index 00000000..9e2f9d52 --- /dev/null +++ b/programs/inspect_peaks/run_frame_inspection.py @@ -0,0 +1,81 @@ +"""Run per-frame peak inspection and save the local score tensor per movie frame. + +This mirrors ``run_inspect_peaks.py`` but scores each movie frame independently, so the +stored tensor carries a ``frame`` axis after ``particle``. The output ``.npz`` is +self-describing. Reload it with:: + + from leopard_em.analysis import load_inspection_result + + result = load_inspection_result("results_frame_inspection.npz") + result.scores # main tensor; see result.axes for layout + result.frame_index # (T,) movie frame index for the ``frame`` axis + result.euler_angle_offsets # (n_orient, 3) ZYZ offsets per orientation index + result.defocus_offsets # (n_defocus,) relative defocus offsets (Angstrom) + result.pixel_size_offsets # (n_px,) relative pixel-size offsets + result.base_euler_angles # (N, 3) per-particle base ZYZ angles + result.base_defocus # (N, 3) per-particle base (defocus_u, defocus_v, angle) + result.particle_index # (N,) maps tensor rows to the particle stack + result.frequency_bins # (n_freq,) FRC frequencies (FRC mode only) + +Stored tensor layout (``result.scores``): +- ``"cross_correlation"``: ``(N, T, n_px, n_defocus, n_orient, H, W)`` +- ``"frc"``: ``(N, T, n_px, n_defocus, n_orient, n_freq)`` +""" + +import time +from typing import Literal + +from leopard_em.pydantic_models.managers import FrameInspectionManager + +####################################### +### Editable parameters for program ### +####################################### + +# Edit the YAML similarly to refine/inspect configs. +YAML_CONFIG_PATH = "/path/to/frame_inspection_configuration.yaml" + +# Where to write the per-frame score tensor. A ``.npz`` suffix is appended if missing. +OUTPUT_PATH = "/path/to/results_frame_inspection.npz" + +# Batched orientations per GPU call - lower if you run out of memory. +CORRELATION_BATCH_SIZE = 32 + +# Output mode for inspect backend. +# - "cross_correlation": saves (N, T, n_px, n_defocus, n_orient, H, W) +# - "frc": saves (N, T, n_px, n_defocus, n_orient, n_freq) + freq bins +# NOTE: the per-frame cross-correlation tensor can be very large (an extra frame axis +# on top of the spatial inspection tensor). Subset the particle stack in the config for +# big runs. +OUTPUT_MODE: Literal["cross_correlation", "frc"] = "cross_correlation" + +# Set True to apply cumulative template dose filtering per frame interval. +APPLY_TEMPLATE_DOSE_WEIGHTING = False + + +def main() -> None: + """Run per-frame inspect workflow and save the score tensor.""" + manager = FrameInspectionManager.from_yaml(YAML_CONFIG_PATH) + + print("Loaded configuration (same schema as refine template).") + print( + f"Running per-frame peak inspection in {OUTPUT_MODE!r} mode " + "(may take a while)..." + ) + + start_time = time.time() + output_path = manager.run_and_save_peak_inspection_per_frame( + output_path=OUTPUT_PATH, + correlation_batch_size=CORRELATION_BATCH_SIZE, + prefer_refined_angles=True, + apply_projection_normalization=True, + output_mode=OUTPUT_MODE, + apply_template_dose_weighting=APPLY_TEMPLATE_DOSE_WEIGHTING, + ) + elapsed = time.time() - start_time + + print(f"Finished per-frame peak inspection in {elapsed:.1f} s") + print(f"Saved per-frame score tensor to: {output_path}") + + +if __name__ == "__main__": + main() diff --git a/programs/inspect_peaks/run_inspect_peaks.py b/programs/inspect_peaks/run_inspect_peaks.py new file mode 100644 index 00000000..b409a29d --- /dev/null +++ b/programs/inspect_peaks/run_inspect_peaks.py @@ -0,0 +1,75 @@ +"""Run peak inspection and save the local score tensor to a self-describing file. + +The output ``.npz`` is self-describing. Reload it with:: + + from leopard_em.analysis import load_inspection_result + + result = load_inspection_result("results_inspect_peaks.npz") + result.scores # main tensor; see result.axes for layout + result.euler_angle_offsets # (n_orient, 3) ZYZ offsets per orientation index + result.defocus_offsets # (n_defocus,) relative defocus offsets (Angstrom) + result.pixel_size_offsets # (n_px,) relative pixel-size offsets + result.particle_index # (N,) maps tensor rows to the particle stack + result.frequency_bins # (n_freq,) FRC frequencies (FRC mode only) + +Stored tensor layout (``result.scores``): +- ``"cross_correlation"``: ``(N, n_px, n_defocus, n_orient, H, W)`` +- ``"frc"``: ``(N, n_px, n_defocus, n_orient, n_freq)`` +""" + +# NOTE: The ``if __name__ == "__main__"`` guard is required for multiprocessing. +import time +from typing import Literal + +from leopard_em.pydantic_models.managers import PeakInspectionManager + +####################################### +### Editable parameters for program ### +####################################### + +# Edit the YAML the same way you would for refine template (see example config next +# to this script, identical schema to `refine_template_example_config.yaml`). +YAML_CONFIG_PATH = "/path/to/inspect_peaks_configuration.yaml" + +# Where to write the score tensor. A ``.npz`` suffix is appended if missing. +OUTPUT_PATH = "/path/to/results_inspect_peaks.npz" + +# Batched orientations per GPU call -- lower if you run out of memory. +CORRELATION_BATCH_SIZE = 32 + +# Output mode for inspect backend. +# - "cross_correlation": saves (N, n_px, n_defocus, n_orient, H, W) +# - "frc": saves (N, n_px, n_defocus, n_orient, n_freq) + frequency bins +# NOTE: the cross-correlation tensor can be very large; subset the particle stack +# in the config for big runs. +# TODO: Automatically chunk particle stack within the InspectPeaksManager through an +# option for managing memory. +OUTPUT_MODE: Literal["cross_correlation", "frc"] = "cross_correlation" + + +def main() -> None: + """Run peak inspection over the full stack and save the score tensor.""" + manager = PeakInspectionManager.from_yaml(YAML_CONFIG_PATH) + + print("Loaded configuration (same schema as refine template).") + print(f"Running peak inspection in {OUTPUT_MODE!r} mode (may take a while)...") + + start_time = time.time() + output_path = manager.run_and_save_peak_inspection( + output_path=OUTPUT_PATH, + correlation_batch_size=CORRELATION_BATCH_SIZE, + prefer_refined_angles=True, + output_mode=OUTPUT_MODE, + ) + elapsed = time.time() - start_time + + print(f"Finished peak inspection in {elapsed:.1f} s") + print(f"Saved score tensor to: {output_path}") + print( + "Load it with `leopard_em.analysis.load_inspection_result` " + "for inspection/plotting (see docs/examples)." + ) + + +if __name__ == "__main__": + main() diff --git a/programs/refine_template/run_refine_template.py b/programs/refine_template/run_refine_template.py index 4f78dbca..ecae3196 100644 --- a/programs/refine_template/run_refine_template.py +++ b/programs/refine_template/run_refine_template.py @@ -14,6 +14,10 @@ YAML_CONFIG_PATH = "/path/to/refine-template-configuration.yaml" # Path to where the dataframe with refined peak parameters will be output. +# May be a .csv or .h5 path -- output format matches the particle_stack back-end +# configured in the YAML file by default (df_path -> CSV, hdf5_path -> HDF5), +# or pass output_format="csv"/"hdf5" to RefineTemplateManager.run_refine_template() +# to override. DATAFRAME_OUTPUT_PATH = "/path/to/refine-template-results.csv" # Number of particles to refine simultaneously. Will need to tune this parameter diff --git a/pyproject.toml b/pyproject.toml index e9e459ad..8a3e520c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ name = "leopard_em" dynamic = ["version"] description = "Location & Orientation of Particles found using Two-Dimensional Template Matching" readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.11" license = { text = "BSD-3-Clause" } authors = [ { name = "Matthew Giammar", email = "matthew_giammar@berkeley.edu" }, @@ -30,7 +30,6 @@ classifiers = [ "Development Status :: 3 - Alpha", "License :: OSI Approved :: BSD License", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -48,15 +47,19 @@ dependencies = [ "types-PyYAML", "roma", "tqdm", - "torch-fourier-slice>=v0.4.0", - "torch-fourier-filter>=v0.3.4", - "torch-so3>=v0.2.0", - "ttsim3d>=v0.4.0", "lmfit", "zenodo-get", - "torch-fourier-shift", - "torch-motion-correction>=0.0.4", - "torch-grid-utils>=v0.0.9", + "h5py", + "tensordict", + # TeamTomo dependencies + "ttsim3d>=v0.4.0", + "torch-fourier-slice>=v0.5.2", + "torch-fourier-filter>=v0.5.2", + "torch-fourier-shell-correlation>=v0.5.2", + "torch-so3>=v0.5.2", + "torch-fourier-shift>=v0.5.2", + "torch-motion-correction>=v0.5.2", + "torch-grid-utils>=v0.5.2", "torch-ctf", ] @@ -76,19 +79,17 @@ dev = [ "pre-commit", "rich", # https://github.com/Textualize/rich "ruff", + "mypy", "pylint", ] # Documentation dependencies docs = [ + "zensical", "python-markdown-math", - "mkdocs", - "mkdocs-material", "mkdocs-jupyter", "mkdocs-inline-select-svg-plugin", "mkdocs-autorefs", - "mkdocs-autoapi", - "mkdocstrings[python]", - "mkdocs-gallery", + "mkdocstrings[python]>=1.0", ] [project.urls] @@ -161,6 +162,9 @@ filterwarnings = [ "error", "ignore::FutureWarning", "ignore::DeprecationWarning:torch.jit", + # mrcfile sets ndarray.dtype directly when reading/writing headers, which NumPy + # 2.5 deprecated. Pending a newer mrcfile release; safe to ignore for now. + "ignore:Setting the dtype on a NumPy array has been deprecated.*:DeprecationWarning", ] addopts = "-m 'not slow'" # Skip slow tests on default markers = ["slow: marks test as slow"] diff --git a/scripts/gen_api_docs.py b/scripts/gen_api_docs.py new file mode 100644 index 00000000..0d8e3121 --- /dev/null +++ b/scripts/gen_api_docs.py @@ -0,0 +1,48 @@ +"""Generate mkdocstrings stub pages for all public leopard_em subpackages. + +Run this before `zensical build`. Output goes to docs/api/ which is git-ignored. +Each stub file contains a single `:::module` directive that mkdocstrings expands +into full API documentation from the source docstrings. +""" + +from pathlib import Path + +SRC = Path("src/leopard_em") +OUT = Path("docs/api") + +# One page per subpackage: maps page title -> module import path. +# Add entries here whenever a new public subpackage is introduced. +PAGES = { + "index": ("leopard_em", "API Reference"), + "analysis": ("leopard_em.analysis", "Analysis"), + "backend": ("leopard_em.backend", "Backend"), + "pydantic_models": ("leopard_em.pydantic_models", "Pydantic Models"), + "pydantic_models/config": ("leopard_em.pydantic_models.config", "Config Models"), + "pydantic_models/data_structures": ( + "leopard_em.pydantic_models.data_structures", + "Data Structures", + ), + "pydantic_models/managers": ( + "leopard_em.pydantic_models.managers", + "Managers", + ), + "pydantic_models/results": ( + "leopard_em.pydantic_models.results", + "Results", + ), + "utils": ("leopard_em.utils", "Utilities"), +} + + +def main() -> None: + """Generate the stub pages.""" + OUT.mkdir(parents=True, exist_ok=True) + for page_path, (module, title) in PAGES.items(): + stub = OUT / f"{page_path}.md" + stub.parent.mkdir(parents=True, exist_ok=True) + stub.write_text(f"# {title}\n\n:::{module}\n") + print(f" wrote {stub}") + + +if __name__ == "__main__": + main() diff --git a/src/leopard_em/analysis/__init__.py b/src/leopard_em/analysis/__init__.py index 7eed542b..3c365636 100644 --- a/src/leopard_em/analysis/__init__.py +++ b/src/leopard_em/analysis/__init__.py @@ -1,5 +1,10 @@ """Submodule for analyzing results during the template matching pipeline.""" +from .inspect_peaks_result import ( + InspectionResult, + load_inspection_result, + save_inspection_result, +) from .match_template_peaks import ( MatchTemplatePeaks, match_template_peaks_to_dataframe, @@ -12,6 +17,9 @@ ) __all__ = [ + "InspectionResult", + "load_inspection_result", + "save_inspection_result", "MatchTemplatePeaks", "match_template_peaks_to_dict", "match_template_peaks_to_dataframe", diff --git a/src/leopard_em/analysis/inspect_peaks.py b/src/leopard_em/analysis/inspect_peaks.py new file mode 100644 index 00000000..66cd88ef --- /dev/null +++ b/src/leopard_em/analysis/inspect_peaks.py @@ -0,0 +1,622 @@ +"""Functions for inspecting local cross-correlations around identified peaks.""" + +# Kwargs/arity mirror the refine-template distributed API (many explicit tensors). +# pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals +# pylint: disable=duplicate-code +import math +from collections.abc import Iterator +from typing import Literal + +import roma +import torch + +from leopard_em.backend.core_refine_template import ( + _device_stream_context, + _iter_refine_particle_correlation_batches, + _make_device_streams, + _move_refine_template_stack_to_device, + _synchronize_device_streams, + _tqdm_for_refine_particle_loop, + construct_multi_gpu_refine_template_kwargs, +) +from leopard_em.backend.cross_correlation import do_batched_orientation_frc +from leopard_em.backend.distributed import run_multiprocess_jobs +from leopard_em.backend.utils import EULER_ANGLE_FMT +from leopard_em.utils.ctf_utils import ( + calculate_ctf_filter_stack_full_args, + move_ctf_kwargs_tensors_to_device, +) + + +def core_inspect_template( + particle_stack_dft: torch.Tensor, # (N, H, W) + template_dft: torch.Tensor, # (d, h, w) + euler_angles: torch.Tensor, # (N, 3) + euler_angle_offsets: torch.Tensor, # (k, 3) + defocus_offsets: torch.Tensor, # (l,) + defocus_u: torch.Tensor, # (N,) + defocus_v: torch.Tensor, # (N,) + defocus_angle: torch.Tensor, # (N,) + pixel_size_offsets: torch.Tensor, # (m,) + corr_mean: torch.Tensor, # (N, H - h + 1, W - w + 1) + corr_std: torch.Tensor, # (N, H - h + 1, W - w + 1) + ctf_kwargs: dict, + projective_filters: torch.Tensor, # (N, h, w) + device: torch.device | list[torch.device], + batch_size: int = 32, + num_cuda_streams: int = 1, + mag_matrix: torch.Tensor | None = None, + apply_projection_normalization: bool = True, + output_mode: Literal["cross_correlation", "frc"] = "cross_correlation", +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Inspect all local hypotheses for each particle. + + Parameters + ---------- + particle_stack_dft : torch.Tensor + Particle image stack in RFFT form with shape ``(N, H, W_rfft)``. + template_dft : torch.Tensor + Template volume in RFFT form. + euler_angles : torch.Tensor + Per-particle base Euler angles with shape ``(N, 3)``. + euler_angle_offsets : torch.Tensor + Orientation offsets searched around each base orientation. + defocus_offsets : torch.Tensor + Relative defocus offsets searched per particle. + defocus_u : torch.Tensor + Per-particle defocus U values. + defocus_v : torch.Tensor + Per-particle defocus V values. + defocus_angle : torch.Tensor + Per-particle astigmatism angles (degrees). + pixel_size_offsets : torch.Tensor + Relative pixel-size offsets searched per particle. + corr_mean : torch.Tensor + Per-particle correlation means used for z-score normalization. + corr_std : torch.Tensor + Per-particle correlation standard deviations used for z-score normalization. + ctf_kwargs : dict + CTF keyword arguments passed to filter construction. + projective_filters : torch.Tensor + Per-particle projective filter stack. + device : torch.device | list[torch.device] + One or more devices used for distributed execution. + batch_size : int, optional + Number of orientation offsets processed per batch. + num_cuda_streams : int, optional + CUDA streams per device worker. + mag_matrix : torch.Tensor | None, optional + Optional anisotropic magnification matrix. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + output_mode : Literal["cross_correlation", "frc"], optional + ``"cross_correlation"`` returns local CC maps. + ``"frc"`` returns local FRC spectra. + + Returns + ------- + torch.Tensor | tuple[torch.Tensor, torch.Tensor] + - ``"cross_correlation"``: ``(N, n_px, n_defocus, n_orient, H, W)``. + - ``"frc"``: ``(frc_tensor, frequency_bins)`` where + ``frc_tensor`` is ``(N, n_px, n_defocus, n_orient, n_freq)`` and + ``frequency_bins`` is ``(n_freq,)``. + """ + if isinstance(device, torch.device): + device = [device] + + kwargs_per_device = construct_multi_gpu_refine_template_kwargs( + particle_stack_dft=particle_stack_dft, + template_dft=template_dft, + euler_angles=euler_angles, + euler_angle_offsets=euler_angle_offsets, + defocus_u=defocus_u, + defocus_v=defocus_v, + defocus_angle=defocus_angle, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + corr_mean=corr_mean, + corr_std=corr_std, + ctf_kwargs=ctf_kwargs, + projective_filters=projective_filters, + batch_size=batch_size, + devices=device, + num_cuda_streams=num_cuda_streams, + mag_matrix=mag_matrix, + ) + for kwargs in kwargs_per_device: + kwargs["apply_projection_normalization"] = apply_projection_normalization + kwargs["output_mode"] = output_mode + + results = run_multiprocess_jobs( + target=_core_inspect_template_single_gpu, + kwargs_list=kwargs_per_device, + ) + + for dev in device: + if dev.type == "cuda": + torch.cuda.synchronize(dev) + + ordered_results = [results[k] for k in sorted(results.keys(), key=int)] + + if output_mode == "cross_correlation": + inspection_stack = torch.cat( + [torch.from_numpy(r["inspection_stack"]) for r in ordered_results] + ) + particle_indices = torch.cat( + [torch.from_numpy(r["particle_indices"]) for r in ordered_results] + ) + sort_indices = torch.argsort(particle_indices) + return inspection_stack[sort_indices] + + frc_stack = torch.cat([torch.from_numpy(r["frc_stack"]) for r in ordered_results]) + particle_indices = torch.cat( + [torch.from_numpy(r["particle_indices"]) for r in ordered_results] + ) + sort_indices = torch.argsort(particle_indices) + sorted_frc = frc_stack[sort_indices] + frequency_bins = torch.from_numpy(ordered_results[0]["frequency_bins"]) + return sorted_frc, frequency_bins + + +def _core_inspect_template_single_gpu( + result_dict: dict, + device_id: int, + particle_stack_dft: torch.Tensor, + particle_indices: torch.Tensor, + template_dft: torch.Tensor, + euler_angles: torch.Tensor, + euler_angle_offsets: torch.Tensor, + defocus_u: torch.Tensor, + defocus_v: torch.Tensor, + defocus_angle: torch.Tensor, + defocus_offsets: torch.Tensor, + pixel_size_offsets: torch.Tensor, + corr_mean: torch.Tensor, + corr_std: torch.Tensor, + projective_filters: torch.Tensor, + ctf_kwargs: dict, + batch_size: int, + device: torch.device, + num_cuda_streams: int = 1, + mag_matrix: torch.Tensor | None = None, + apply_projection_normalization: bool = True, + output_mode: Literal["cross_correlation", "frc"] = "cross_correlation", +) -> None: + """Inspect all local hypotheses for one device subset. + + Parameters + ---------- + result_dict : dict + Shared multiprocessing dictionary to collect outputs. + device_id : int + Worker/device index used as the result key and tqdm position. + particle_stack_dft : torch.Tensor + Device-local particle stack chunk in RFFT form. + particle_indices : torch.Tensor + Original global particle indices for sorting merged outputs. + template_dft : torch.Tensor + Template volume in RFFT form. + euler_angles : torch.Tensor + Device-local base Euler angles. + euler_angle_offsets : torch.Tensor + Orientation offsets searched around each base orientation. + defocus_u : torch.Tensor + Device-local defocus U values. + defocus_v : torch.Tensor + Device-local defocus V values. + defocus_angle : torch.Tensor + Device-local astigmatism angles. + defocus_offsets : torch.Tensor + Relative defocus offsets searched per particle. + pixel_size_offsets : torch.Tensor + Relative pixel-size offsets searched per particle. + corr_mean : torch.Tensor + Device-local correlation means. + corr_std : torch.Tensor + Device-local correlation standard deviations. + projective_filters : torch.Tensor + Device-local projective filters. + ctf_kwargs : dict + CTF keyword arguments. + batch_size : int + Number of orientation offsets processed per batch. + device : torch.device + Device for this worker. + num_cuda_streams : int, optional + CUDA streams per worker. + mag_matrix : torch.Tensor | None, optional + Optional anisotropic magnification matrix. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + output_mode : Literal["cross_correlation", "frc"], optional + Score mode for this worker. + """ + if device.type == "cuda": + torch.cuda.set_device(device) + + streams = _make_device_streams(device, num_cuda_streams) + + refine_stack = _move_refine_template_stack_to_device( + device, + particle_stack_dft, + particle_indices, + template_dft, + euler_angles, + euler_angle_offsets, + defocus_u, + defocus_v, + defocus_angle, + defocus_offsets, + pixel_size_offsets, + corr_mean, + corr_std, + projective_filters, + mag_matrix, + ) + + num_particles = refine_stack.particle_stack_dft.shape[0] + pbar_iter = _tqdm_for_refine_particle_loop( + num_particles, device, device_id, "Inspecting" + ) + + inspection_results = [] + frc_frequency_bins: torch.Tensor | None = None + for i in pbar_iter: + stream = streams[i % len(streams)] + with _device_stream_context(stream): + inspection_stack = _core_inspect_template_single_thread( + particle_image_dft=refine_stack.particle_stack_dft[i], + template_dft=refine_stack.template_dft, + euler_angles=refine_stack.euler_angles[i, :], + euler_angle_offsets=refine_stack.euler_angle_offsets, + defocus_u=refine_stack.defocus_u[i], + defocus_v=refine_stack.defocus_v[i], + defocus_angle=refine_stack.defocus_angle[i], + defocus_offsets=refine_stack.defocus_offsets, + pixel_size_offsets=refine_stack.pixel_size_offsets, + corr_mean=refine_stack.corr_mean[i], + corr_std=refine_stack.corr_std[i], + ctf_kwargs=ctf_kwargs, + projective_filter=refine_stack.projective_filters[i], + batch_size=batch_size, + mag_matrix=refine_stack.mag_matrix, + apply_projection_normalization=apply_projection_normalization, + output_mode=output_mode, + ) + if output_mode == "cross_correlation": + inspection_results.append(inspection_stack) + else: + frc_tensor, frequency_bins = inspection_stack + inspection_results.append(frc_tensor) + if frc_frequency_bins is None: + frc_frequency_bins = frequency_bins + + _synchronize_device_streams(streams) + + if output_mode == "cross_correlation": + result_dict[device_id] = { + "inspection_stack": torch.stack(inspection_results).cpu().numpy(), + "particle_indices": refine_stack.particle_indices.cpu().numpy(), + } + else: + if frc_frequency_bins is None: + raise ValueError("No FRC frequencies were generated.") + result_dict[device_id] = { + "frc_stack": torch.stack(inspection_results).cpu().numpy(), + "frequency_bins": frc_frequency_bins.cpu().numpy(), + "particle_indices": refine_stack.particle_indices.cpu().numpy(), + } + + +def _core_inspect_template_single_thread( + particle_image_dft: torch.Tensor, + template_dft: torch.Tensor, + euler_angles: torch.Tensor, + euler_angle_offsets: torch.Tensor, + defocus_u: float, + defocus_v: float, + defocus_angle: float, + defocus_offsets: torch.Tensor, + pixel_size_offsets: torch.Tensor, + corr_mean: torch.Tensor, + corr_std: torch.Tensor, + ctf_kwargs: dict, + projective_filter: torch.Tensor, + batch_size: int = 32, + mag_matrix: torch.Tensor | None = None, + apply_projection_normalization: bool = True, + output_mode: Literal["cross_correlation", "frc"] = "cross_correlation", +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run inspect scoring for one particle without best-value reduction. + + Parameters + ---------- + particle_image_dft : torch.Tensor + Single particle image in RFFT form. + template_dft : torch.Tensor + Template volume in RFFT form. + euler_angles : torch.Tensor + Base Euler angle for this particle. + euler_angle_offsets : torch.Tensor + Orientation offsets searched around the base orientation. + defocus_u : float + Particle defocus U. + defocus_v : float + Particle defocus V. + defocus_angle : float + Particle astigmatism angle. + defocus_offsets : torch.Tensor + Relative defocus offsets searched per particle. + pixel_size_offsets : torch.Tensor + Relative pixel-size offsets searched per particle. + corr_mean : torch.Tensor + Per-particle correlation mean map used for z-score normalization. + corr_std : torch.Tensor + Per-particle correlation std map used for z-score normalization. + ctf_kwargs : dict + CTF keyword arguments. + projective_filter : torch.Tensor + Particle-specific projective filter. + batch_size : int, optional + Number of orientation offsets processed per batch. + mag_matrix : torch.Tensor | None, optional + Optional anisotropic magnification matrix. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + output_mode : Literal["cross_correlation", "frc"], optional + Score mode (CC map or FRC spectrum). + + Returns + ------- + torch.Tensor | tuple[torch.Tensor, torch.Tensor] + - ``"cross_correlation"``: tensor ``(n_px, n_defocus, n_orient, H, W)``. + - ``"frc"``: tuple ``(frc_tensor, frequency_bins)`` where + ``frc_tensor`` is ``(n_px, n_defocus, n_orient, n_freq)``. + """ + # TODO: corr_mean/corr_std are documented for z-score normalization but are + # currently unused here; raw (non-normalized) scores are returned. + _ = corr_mean + _ = corr_std + if output_mode == "frc": + frc_batches = _iter_refine_particle_frc_batches( + particle_image_dft=particle_image_dft, + template_dft=template_dft, + euler_angles=euler_angles, + euler_angle_offsets=euler_angle_offsets, + defocus_u=defocus_u, + defocus_v=defocus_v, + defocus_angle=defocus_angle, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + ctf_kwargs=ctf_kwargs, + projective_filter=projective_filter, + batch_size=batch_size, + mag_matrix=mag_matrix, + apply_projection_normalization=apply_projection_normalization, + ) + return _reduce_refine_all_frc( + frc_batches=frc_batches, + num_orientations=euler_angle_offsets.shape[0], + ) + + correlation_batches = _iter_refine_particle_correlation_batches( + particle_image_dft=particle_image_dft, + template_dft=template_dft, + euler_angles=euler_angles, + euler_angle_offsets=euler_angle_offsets, + defocus_u=defocus_u, + defocus_v=defocus_v, + defocus_angle=defocus_angle, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + ctf_kwargs=ctf_kwargs, + projective_filter=projective_filter, + batch_size=batch_size, + mag_matrix=mag_matrix, + apply_projection_normalization=apply_projection_normalization, + ) + + return _reduce_refine_all( + correlation_batches=correlation_batches, + num_orientations=euler_angle_offsets.shape[0], + ) + + +def _iter_refine_particle_frc_batches( + particle_image_dft: torch.Tensor, + template_dft: torch.Tensor, + euler_angles: torch.Tensor, + euler_angle_offsets: torch.Tensor, + defocus_u: float, + defocus_v: float, + defocus_angle: float, + defocus_offsets: torch.Tensor, + pixel_size_offsets: torch.Tensor, + ctf_kwargs: dict, + projective_filter: torch.Tensor, + batch_size: int = 32, + mag_matrix: torch.Tensor | None = None, + apply_projection_normalization: bool = True, +) -> Iterator[tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]]: + """Yield FRC batches for a single particle over orientation offsets. + + Parameters + ---------- + particle_image_dft : torch.Tensor + Single particle image in RFFT form. + particle_index : int + Original global particle index (for progress labeling only). + template_dft : torch.Tensor + Template volume in RFFT form. + euler_angles : torch.Tensor + Base Euler angle for this particle. + euler_angle_offsets : torch.Tensor + Orientation offsets searched around the base orientation. + defocus_u : float + Particle defocus U. + defocus_v : float + Particle defocus V. + defocus_angle : float + Particle astigmatism angle. + defocus_offsets : torch.Tensor + Relative defocus offsets searched per particle. + pixel_size_offsets : torch.Tensor + Relative pixel-size offsets searched per particle. + ctf_kwargs : dict + CTF keyword arguments. + projective_filter : torch.Tensor + Particle-specific projective filter. + batch_size : int, optional + Number of orientation offsets processed per batch. + device_id : int, optional + Worker/device index for tqdm positioning. + mag_matrix : torch.Tensor | None, optional + Optional anisotropic magnification matrix. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + + Yields + ------ + tuple[int, torch.Tensor, torch.Tensor, torch.Tensor] + ``(start_idx, angle_offsets_batch, frc_values, frequency_bins)`` where + ``frc_values`` has shape ``(n_px, n_defocus, n_orient_batch, n_freq)``. + """ + default_rot_matrix = roma.euler_to_rotmat( + EULER_ANGLE_FMT, euler_angles, degrees=True, device=particle_image_dft.device + ) + default_rot_matrix = default_rot_matrix.to(torch.float32) + + ctf_dev_kwargs = move_ctf_kwargs_tensors_to_device( + ctf_kwargs, particle_image_dft.device + ) + ctf_filters = calculate_ctf_filter_stack_full_args( + defocus_u=defocus_u, + defocus_v=defocus_v, + astigmatism_angle=defocus_angle, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + **ctf_dev_kwargs, + ) + combined_projective_filter = projective_filter[None, None, ...] * ctf_filters + + num_batches = math.ceil(euler_angle_offsets.shape[0] / batch_size) + for i in range(num_batches): + start_idx = i * batch_size + end_idx = min((i + 1) * batch_size, euler_angle_offsets.shape[0]) + euler_angle_offsets_batch = euler_angle_offsets[start_idx:end_idx] + rot_matrix_batch = roma.euler_to_rotmat( + EULER_ANGLE_FMT, + euler_angle_offsets_batch, + degrees=True, + device=particle_image_dft.device, + ) + rot_matrix_batch = rot_matrix_batch.to(torch.float32) + rot_matrix_batch = roma.rotmat_composition( + (rot_matrix_batch, default_rot_matrix) + ) + + frc_values, frequency_bins = do_batched_orientation_frc( + image_dft=particle_image_dft, + template_dft=template_dft, + rotation_matrices=rot_matrix_batch, + projective_filters=combined_projective_filter, + apply_normalization=apply_projection_normalization, + mag_matrix=mag_matrix, + ) + yield ( + start_idx, + euler_angle_offsets_batch, + frc_values, + frequency_bins, + ) + + +def _reduce_refine_all_frc( + frc_batches: Iterator[tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]], + num_orientations: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Stitch batched FRC results into a full orientation tensor. + + Parameters + ---------- + frc_batches : Iterator[tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]] + Iterator from :func:`_iter_refine_particle_frc_batches`. + num_orientations : int + Total number of orientation offsets across all batches. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + ``(frc_tensor, frequency_bins)`` where + ``frc_tensor`` is ``(n_px, n_defocus, n_orient, n_freq)`` and + ``frequency_bins`` is ``(n_freq,)``. + """ + output = None + frequency_bins = None + for start_idx, angle_offsets_batch, frc_values, batch_frequency_bins in frc_batches: + if output is None: + output = torch.empty( + ( + frc_values.shape[0], + frc_values.shape[1], + num_orientations, + frc_values.shape[-1], + ), + dtype=frc_values.dtype, + device=frc_values.device, + ) + frequency_bins = batch_frequency_bins + end_idx = start_idx + len(angle_offsets_batch) + output[:, :, start_idx:end_idx] = frc_values + + if output is None or frequency_bins is None: + raise ValueError("No FRC batches were generated.") + + return output, frequency_bins + + +def _reduce_refine_all( + correlation_batches: Iterator[tuple[int, torch.Tensor, torch.Tensor, int, int]], + num_orientations: int, +) -> torch.Tensor: + """Stitch local CC batches into a full orientation tensor. + + Parameters + ---------- + correlation_batches : Iterator[tuple[int, torch.Tensor, torch.Tensor, int, int]] + Iterator from :func:`_iter_refine_particle_correlation_batches`. + num_orientations : int + Total number of orientation offsets across all batches. + + Returns + ------- + torch.Tensor + Tensor with shape ``(n_px, n_defocus, n_orient, H, W)``. + """ + output = None + for ( + start_idx, + angle_offsets_batch, + cross_correlation, + crop_h, + crop_w, + ) in correlation_batches: + if output is None: + output = torch.empty( + ( + cross_correlation.shape[0], + cross_correlation.shape[1], + num_orientations, + crop_h, + crop_w, + ), + dtype=cross_correlation.dtype, + device=cross_correlation.device, + ) + end_idx = start_idx + len(angle_offsets_batch) + output[:, :, start_idx:end_idx] = cross_correlation + + if output is None: + raise ValueError("No orientation batches were generated.") + + return output diff --git a/src/leopard_em/analysis/inspect_peaks_result.py b/src/leopard_em/analysis/inspect_peaks_result.py new file mode 100644 index 00000000..51d98d4c --- /dev/null +++ b/src/leopard_em/analysis/inspect_peaks_result.py @@ -0,0 +1,303 @@ +"""Data class and helper functions for saving/loading peak-inspection results. + +Peak inspection produces a large local score tensor per particle (cross-correlation +maps or FRC spectra over a grid of orientation/defocus/pixel-size hypotheses). This +module bundles that tensor with the axis metadata needed to interpret it into a +self-describing ``.npz`` file, so downstream analysis (e.g. notebooks) can reload the +array without knowing how the run was configured. + + +Shape of the cross-correlation tensor ``result.scores`` +------------------------------------------------------- + +For the cross-correlation mode, have a 6-D tensor with last two dimensions corresponding +to the spatial dimensions of the valid correlation map. +``` +scores.shape = (N, n_px, n_def, n_orient, H, W) + | | | | | | + | | | | | +-- valid same-size CC map (x) + | | | | +----- valid same-size CC map (y) + | | | +--------------- local Euler *offsets* (phi,theta,psi) + | | +---------------------- relative defocus search index + | +---------------------------- relative pixel-size search index + +-------------------------------- particle (row in stack / CSV) +``` + +For the FRC mode, have a 5-D tensor where the last dimension corresponds to the +FRC frequency bins. +``` +scores.shape = (N, n_px, n_def, n_orient, num_freq) + | | | | | + | | | | +----- FRC frequency bins + | | | +--------------- local Euler *offsets* (phi,theta,psi) + | | +---------------------- relative defocus search index + | +---------------------------- relative pixel-size search index + +-------------------------------- particle (row in stack / CSV) +``` + +Per-frame inspection +-------------------- + +The per-frame manager scores each movie frame independently, inserting a ``frame`` +axis immediately after ``particle`` (``per_frame=True``). The cross-correlation tensor +is then 7-D ``(N, T, n_px, n_def, n_orient, H, W)`` and the FRC tensor is 6-D +``(N, T, n_px, n_def, n_orient, num_freq)``, where ``T`` is the number of frames. +""" + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch + +# Bump when the on-disk layout changes in a backwards-incompatible way. +INSPECTION_FORMAT_VERSION = 1 + +# Axis labels for the stored tensors. These describe what each dimension of the +# main array means so downstream analysis (notebooks) does not have to hard-code +# the 6-D / 5-D layout. +CROSS_CORRELATION_AXES = ( + "particle", + "pixel_size", + "defocus", + "orientation", + "y", + "x", +) +FRC_AXES = ( + "particle", + "pixel_size", + "defocus", + "orientation", + "frequency", +) + +# Per-frame variants insert a ``frame`` axis right after ``particle``. +CROSS_CORRELATION_FRAME_AXES = ( + "particle", + "frame", + "pixel_size", + "defocus", + "orientation", + "y", + "x", +) +FRC_FRAME_AXES = ( + "particle", + "frame", + "pixel_size", + "defocus", + "orientation", + "frequency", +) + + +@dataclass +class InspectionResult: + """Self-describing container for a saved peak-inspection run. + + Attributes + ---------- + output_mode : str + Either ``"cross_correlation"`` or ``"frc"``. + scores : np.ndarray + The main score tensor. Shape is ``(N, n_px, n_defocus, n_orient, H, W)`` for + cross-correlation mode and ``(N, n_px, n_defocus, n_orient, n_freq)`` for FRC + mode. See :attr:`axes` for per-dimension labels. + axes : tuple[str, ...] + Label for each dimension of :attr:`scores`. + euler_angle_offsets : np.ndarray + ZYZ orientation offsets searched per particle, shape ``(n_orient, 3)``. Indexes + the ``orientation`` axis of :attr:`scores`. + defocus_offsets : np.ndarray + Relative defocus offsets (Angstroms), shape ``(n_defocus,)``. Indexes the + ``defocus`` axis. + pixel_size_offsets : np.ndarray + Relative pixel-size offsets, shape ``(n_px,)``. Indexes the ``pixel_size`` axis. + base_euler_angles : np.ndarray + Per-particle base ZYZ angles the offsets are relative to, shape ``(N, 3)``. + base_defocus : np.ndarray + Per-particle base astigmatic defocus the ``defocus_offsets`` are relative to, + shape ``(N, 3)`` as ``(defocus_u, defocus_v, defocus_angle)``. + particle_index : np.ndarray | None + Global particle index for each row of the ``particle`` axis, shape ``(N,)``, or + ``None`` if the source dataframe had no ``particle_index`` + column. Maps tensor rows back to the particle stack dataframe. + frequency_bins : np.ndarray | None + FRC frequency bins, shape ``(n_freq,)``, in FRC mode; ``None`` otherwise. + frame_index : np.ndarray | None + Movie frame index for each entry of the ``frame`` axis, shape ``(T,)``, when the + result was produced by per-frame inspection; ``None`` otherwise. + metadata : dict[str, Any] + Free-form metadata stored alongside the arrays (includes the format version and + any caller-supplied ``extra_metadata``). + """ + + output_mode: Literal["cross_correlation", "frc"] + scores: np.ndarray + axes: tuple[str, ...] + euler_angle_offsets: np.ndarray + defocus_offsets: np.ndarray + pixel_size_offsets: np.ndarray + base_euler_angles: np.ndarray + base_defocus: np.ndarray + particle_index: np.ndarray | None + frequency_bins: np.ndarray | None + frame_index: np.ndarray | None + metadata: dict[str, Any] + + +def _to_numpy(tensor: torch.Tensor | np.ndarray) -> np.ndarray: + """Return a detached CPU numpy view of a tensor (passthrough for ndarrays).""" + if isinstance(tensor, torch.Tensor): + return tensor.detach().cpu().numpy() + return np.asarray(tensor) + + +# pylint: disable=too-many-arguments,too-many-positional-arguments +def save_inspection_result( + output_path: str | Path, + *, + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + output_mode: Literal["cross_correlation", "frc"], + euler_angle_offsets: torch.Tensor, + defocus_offsets: torch.Tensor, + pixel_size_offsets: torch.Tensor, + base_euler_angles: torch.Tensor, + base_defocus: torch.Tensor, + particle_index: torch.Tensor | np.ndarray | None = None, + frame_index: torch.Tensor | np.ndarray | None = None, + per_frame: bool = False, + extra_metadata: dict[str, Any] | None = None, +) -> Path: + """Write a peak-inspection result to a self-describing ``.npz`` file. + + Parameters + ---------- + output_path : str | Path + Destination path. A ``.npz`` suffix is appended if not present. + result : torch.Tensor | tuple[torch.Tensor, torch.Tensor] + Output of the inspect backend. A tensor in ``"cross_correlation"`` mode, or + ``(frc_tensor, frequency_bins)`` in ``"frc"`` mode. When ``per_frame`` is True + the score tensor carries an extra ``frame`` axis after ``particle``. + output_mode : Literal["cross_correlation", "frc"] + Score mode used to produce ``result``. + euler_angle_offsets : torch.Tensor + Orientation offsets searched, shape ``(n_orient, 3)``. + defocus_offsets : torch.Tensor + Relative defocus offsets searched, shape ``(n_defocus,)``. + pixel_size_offsets : torch.Tensor + Relative pixel-size offsets searched, shape ``(n_px,)``. + base_euler_angles : torch.Tensor + Per-particle base ZYZ angles the offsets are relative to, shape ``(N, 3)``. + base_defocus : torch.Tensor + Per-particle base astigmatic defocus the offsets are relative to, shape + ``(N, 3)`` as ``(defocus_u, defocus_v, defocus_angle)``. + particle_index : torch.Tensor | np.ndarray | None, optional + Global particle index for each tensor row, shape ``(N,)``. + frame_index : torch.Tensor | np.ndarray | None, optional + Movie frame index for each entry of the ``frame`` axis, shape ``(T,)``. Only + meaningful when ``per_frame`` is True. + per_frame : bool, optional + If True, the score tensor carries a ``frame`` axis after ``particle`` and the + stored axis labels use the per-frame variants. + extra_metadata : dict[str, Any] | None, optional + Additional JSON-serializable metadata to store alongside the arrays. + + Returns + ------- + Path + The path the result was written to (with ``.npz`` suffix). + """ + output_path = Path(output_path) + if output_path.suffix != ".npz": + output_path = output_path.with_suffix(".npz") + output_path.parent.mkdir(parents=True, exist_ok=True) + + axes: tuple[str, ...] + + if output_mode == "frc": + if not (isinstance(result, tuple) and len(result) == 2): + raise ValueError( + "FRC mode expects a (frc_tensor, frequency_bins) tuple result." + ) + scores_tensor, frequency_bins = result + axes = FRC_FRAME_AXES if per_frame else FRC_AXES + elif output_mode == "cross_correlation": + if not isinstance(result, torch.Tensor): + raise ValueError("Cross-correlation mode expects a single tensor result.") + scores_tensor = result + frequency_bins = None + axes = CROSS_CORRELATION_FRAME_AXES if per_frame else CROSS_CORRELATION_AXES + else: + raise ValueError(f"Unknown output_mode: {output_mode!r}") + + arrays: dict[str, np.ndarray] = { + "scores": _to_numpy(scores_tensor), + "euler_angle_offsets": _to_numpy(euler_angle_offsets), + "defocus_offsets": _to_numpy(defocus_offsets), + "pixel_size_offsets": _to_numpy(pixel_size_offsets), + "base_euler_angles": _to_numpy(base_euler_angles), + "base_defocus": _to_numpy(base_defocus), + } + + if frequency_bins is not None: + arrays["frequency_bins"] = _to_numpy(frequency_bins) + if particle_index is not None: + arrays["particle_index"] = _to_numpy(particle_index) + if per_frame and frame_index is not None: + arrays["frame_index"] = _to_numpy(frame_index) + + metadata: dict[str, Any] = { + "format_version": INSPECTION_FORMAT_VERSION, + "output_mode": output_mode, + "axes": list(axes), + "per_frame": per_frame, + } + + # Store metadata as a JSON string in a 0-d array so it survives the round trip. + if extra_metadata: + metadata.update(extra_metadata) + arrays["metadata_json"] = np.array(json.dumps(metadata)) + + np.savez(output_path, **arrays) + + return output_path + + +def load_inspection_result(path: str | Path) -> InspectionResult: + """Load a ``.npz`` written by :func:`save_inspection_result`. + + Parameters + ---------- + path : str | Path + Path to the ``.npz`` file. + + Returns + ------- + InspectionResult + Self-describing container with the score tensor and its axis metadata. + """ + with np.load(path, allow_pickle=False) as data: + metadata = json.loads(str(data["metadata_json"])) + output_mode = metadata["output_mode"] + return InspectionResult( + output_mode=output_mode, + scores=data["scores"], + axes=tuple(metadata["axes"]), + euler_angle_offsets=data["euler_angle_offsets"], + defocus_offsets=data["defocus_offsets"], + pixel_size_offsets=data["pixel_size_offsets"], + base_euler_angles=data["base_euler_angles"], + base_defocus=data["base_defocus"], + particle_index=( + data["particle_index"] if "particle_index" in data else None + ), + frequency_bins=( + data["frequency_bins"] if "frequency_bins" in data else None + ), + frame_index=(data["frame_index"] if "frame_index" in data else None), + metadata=metadata, + ) diff --git a/src/leopard_em/backend/core_differentiable_refine.py b/src/leopard_em/backend/core_differentiable_refine.py index def99a89..0e983ad3 100644 --- a/src/leopard_em/backend/core_differentiable_refine.py +++ b/src/leopard_em/backend/core_differentiable_refine.py @@ -17,7 +17,10 @@ do_batched_orientation_cross_correlate, ) from leopard_em.backend.utils import EULER_ANGLE_FMT, combine_euler_angles -from leopard_em.utils.ctf_utils import calculate_ctf_filter_stack_full_args +from leopard_em.utils.ctf_utils import ( + calculate_ctf_filter_stack_full_args, + move_ctf_kwargs_tensors_to_device, +) # NOTE: Disabling pylint for too many arguments because we are taking a data-oriented @@ -544,13 +547,16 @@ def _core_refine_template_single_thread( default_rot_matrix = default_rot_matrix.to(torch.float32) # Calculate the CTF filters with the relative offsets + ctf_dev_kwargs = move_ctf_kwargs_tensors_to_device( + ctf_kwargs, particle_image_dft.device + ) ctf_filters = calculate_ctf_filter_stack_full_args( defocus_u=defocus_u, # in Angstrom defocus_v=defocus_v, # in Angstrom astigmatism_angle=defocus_angle, # in degrees defocus_offsets=defocus_offsets, # in Angstrom pixel_size_offsets=pixel_size_offsets, # in Angstrom - **ctf_kwargs, + **ctf_dev_kwargs, ) # Combine the single projective filter with the CTF filter diff --git a/src/leopard_em/backend/core_match_template.py b/src/leopard_em/backend/core_match_template.py index 1aa358a4..16d4fd7c 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,8 @@ def core_match_template( num_cuda_streams: int = 1, backend: str = "streamed", mag_matrix: torch.Tensor | None = None, -) -> dict[str, torch.Tensor]: + compute_correlation_table: bool = True, +) -> 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 @@ -213,10 +221,14 @@ def core_match_template( mag_matrix : torch.Tensor | None, optional Anisotropic magnification matrix of shape (2, 2). If None, no magnification transform is applied. Default is None. + compute_correlation_table : bool, optional + Whether to track cross-correlation values which surpass the correlation table + threshold. If False, this (comparatively expensive) computation is skipped and + the returned "correlation_table" will be empty. Default is True. 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 +238,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. """ @@ -316,6 +330,7 @@ def core_match_template( "backend": backend, "device": d, "mag_matrix": mag_matrix, + "compute_correlation_table": compute_correlation_table, } kwargs_per_device.append(kwargs) @@ -335,7 +350,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 +363,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 +380,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 +404,10 @@ 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]: + compute_correlation_table: bool = True, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, tensordict.TensorDict +]: """Single-GPU call for template matching. Parameters @@ -422,6 +449,10 @@ def _core_match_template_single_gpu( mag_matrix : torch.Tensor | None, optional Anisotropic magnification matrix of shape (2, 2). If None, no magnification transform is applied. Default is None. + compute_correlation_table : bool, optional + Whether to track cross-correlation values which surpass the correlation table + threshold. If False, this (comparatively expensive) computation is skipped and + the returned correlation table will be empty. Default is True. Returns ------- @@ -433,8 +464,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 +510,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 +599,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 +608,35 @@ 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, + ) + else: + raise ValueError( + f"Unknown backend '{backend}'. Must be one of 'batched', " + "'streamed', or 'zipfft'." + ) - # 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"), + compute_correlation_table=compute_correlation_table, ) + except Exception as e: index_queue.set_error_flag() print(f"Error occurred in process {rank}: {e}") @@ -559,7 +648,14 @@ def _core_match_template_single_gpu( torch.cuda.synchronize(device) - return mip, best_global_index, correlation_sum, correlation_squared_sum + # pylint: disable=duplicate-code + return ( + mip, + best_global_index, + correlation_sum, + correlation_squared_sum, + correlation_table, + ) def _core_match_template_multiprocess_wrapper( @@ -573,9 +669,14 @@ 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] - ) + # pylint: disable=duplicate-code + ( + 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 +685,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..2558354d 100644 --- a/src/leopard_em/backend/core_match_template_distributed.py +++ b/src/leopard_em/backend/core_match_template_distributed.py @@ -4,8 +4,9 @@ import random import socket from datetime import timedelta -from typing import Optional +from typing import Any, Optional +import tensordict import torch import torch.distributed as dist @@ -19,6 +20,7 @@ from leopard_em.backend.process_results import ( aggregate_distributed_results, decode_global_search_index, + process_correlation_table, scale_mip, ) @@ -350,6 +352,41 @@ def _gather_tensors_to_rank_zero( ) +def _gather_correlation_table_to_rank_zero( + world_size: int, + rank: int, + correlation_table: tensordict.TensorDict, +) -> Optional[list[dict[str, Any]]]: + """Gather the (variable-length) per-rank correlation tables onto rank zero. + + Parameters + ---------- + world_size : int + Total number of processes in the distributed job. + rank : int + Global rank of this process. + correlation_table : tensordict.TensorDict + This rank's correlation table. + + Returns + ------- + Optional[list[dict[str, Any]]] + List of per-rank correlation table dictionaries on rank zero, or None on + other ranks. + """ + correlation_table_cpu = correlation_table.cpu().to_dict() + + gather_list: Optional[list[dict[str, Any]]] = ( + [None] * world_size if rank == 0 else None # type: ignore[list-item] + ) + + dist.barrier() + dist.gather_object(correlation_table_cpu, gather_list, dst=0) + dist.barrier() + + return gather_list + + # pylint: disable=too-many-locals def core_match_template_distributed( world_size: int, @@ -359,6 +396,7 @@ def core_match_template_distributed( orientation_batch_size: int = 1, num_cuda_streams: int = 1, backend: str = "streamed", + compute_correlation_table: bool = True, **kwargs: dict, ) -> dict[str, torch.Tensor]: """Distributed multi-node core function for the match template program. @@ -381,6 +419,10 @@ def core_match_template_distributed( backend : str, optional The backend to use for computation. Defaults to 'streamed'. Must be 'streamed' or 'batched'. + compute_correlation_table : bool, optional + Whether to track cross-correlation values which surpass the correlation table + threshold. If False, this (comparatively expensive) computation is skipped. + Default is True. **kwargs : dict[str, torch.Tensor] Additional keyword arguments passed to the single-GPU core function. For the zeroth rank this should be a dictionary of Tensor objects with the following @@ -454,21 +496,27 @@ 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, - ) + # pylint: disable=duplicate-code + ( + mip, + best_global_index, + correlation_sum, + correlation_squared_sum, + correlation_table, + ) = _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, + compute_correlation_table=compute_correlation_table, ) dist.barrier() @@ -487,6 +535,13 @@ def core_match_template_distributed( correlation_squared_sum=correlation_squared_sum, ) + # Gather the variable-length correlation tables to rank zero + gather_correlation_table = _gather_correlation_table_to_rank_zero( + world_size=world_size, + rank=rank, + correlation_table=correlation_table, + ) + ################################################## ### Final aggregation step on the main process ### ################################################## @@ -499,6 +554,7 @@ def core_match_template_distributed( assert gather_best_global_index is not None assert gather_correlation_sum is not None assert gather_correlation_squared_sum is not None + assert gather_correlation_table is not None aggregated_results = aggregate_distributed_results( results=[ @@ -507,12 +563,14 @@ def core_match_template_distributed( "best_global_index": gidx, "correlation_sum": corr_sum, "correlation_squared_sum": corr_sq_sum, + "correlation_table": corr_table, } - for mip, gidx, corr_sum, corr_sq_sum in zip( + for mip, gidx, corr_sum, corr_sq_sum, corr_table in zip( gather_mip, gather_best_global_index, gather_correlation_sum, gather_correlation_squared_sum, + gather_correlation_table, ) ] ) @@ -534,10 +592,18 @@ 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 ) + # Process the correlation table into a more interpretable format + correlation_table = process_correlation_table( + aggregated_results["correlation_table"], + pixel_values, + defocus_values, + euler_angles, + ) + mip_scaled = torch.empty_like(mip) mip, mip_scaled, correlation_mean, correlation_variance = scale_mip( mip=mip, @@ -556,6 +622,7 @@ def core_match_template_distributed( "best_defocus": best_defocus.cpu(), "correlation_mean": correlation_mean.cpu(), "correlation_variance": correlation_variance.cpu(), + "correlation_table": correlation_table, "total_projections": total_projections, "total_orientations": euler_angles.shape[0], "total_defocus": defocus_values.shape[0], diff --git a/src/leopard_em/backend/core_refine_template.py b/src/leopard_em/backend/core_refine_template.py index c5cc8efd..723b7620 100644 --- a/src/leopard_em/backend/core_refine_template.py +++ b/src/leopard_em/backend/core_refine_template.py @@ -1,10 +1,12 @@ """Backend functions related to correlating and refining particle stacks.""" # Following pylint error ignored because torc.fft.* is not recognized as callable -# pylint: disable=E1102 +# pylint: disable=E1102,too-many-lines import math -from typing import Literal +from collections.abc import Iterator +from contextlib import AbstractContextManager, nullcontext +from typing import Literal, NamedTuple, cast import roma import torch @@ -22,7 +24,110 @@ normalize_template_projection, ) from leopard_em.utils.cross_correlation import handle_correlation_mode -from leopard_em.utils.ctf_utils import calculate_ctf_filter_stack_full_args +from leopard_em.utils.ctf_utils import ( + calculate_ctf_filter_stack_full_args, + move_ctf_kwargs_tensors_to_device, +) + + +def _make_device_streams( + device: torch.device, num_cuda_streams: int +) -> list[torch.cuda.Stream | None]: + """Create CUDA streams when running on CUDA, otherwise a CPU placeholder.""" + if device.type == "cuda": + return [torch.cuda.Stream(device=device) for _ in range(num_cuda_streams)] + return [None] + + +def _device_stream_context( + stream: torch.cuda.Stream | None, +) -> AbstractContextManager[None]: + """Return the appropriate stream context for CUDA or CPU execution.""" + if stream is None: + return nullcontext() + return cast(AbstractContextManager[None], torch.cuda.stream(stream)) + + +def _synchronize_device_streams(streams: list[torch.cuda.Stream | None]) -> None: + """Synchronize CUDA streams and no-op for CPU placeholders.""" + for stream in streams: + if stream is not None: + stream.synchronize() + + +class _RefineTemplateStackOnDevice(NamedTuple): + """All per-device tensors for one refine- or inspect-template GPU chunk.""" + + particle_stack_dft: torch.Tensor + particle_indices: torch.Tensor + template_dft: torch.Tensor + euler_angles: torch.Tensor + euler_angle_offsets: torch.Tensor + defocus_u: torch.Tensor + defocus_v: torch.Tensor + defocus_angle: torch.Tensor + defocus_offsets: torch.Tensor + pixel_size_offsets: torch.Tensor + corr_mean: torch.Tensor + corr_std: torch.Tensor + projective_filters: torch.Tensor + mag_matrix: torch.Tensor | None + + +# pylint: disable=too-many-arguments,too-many-positional-arguments +def _move_refine_template_stack_to_device( + device: torch.device, + particle_stack_dft: torch.Tensor, + particle_indices: torch.Tensor, + template_dft: torch.Tensor, + euler_angles: torch.Tensor, + euler_angle_offsets: torch.Tensor, + defocus_u: torch.Tensor, + defocus_v: torch.Tensor, + defocus_angle: torch.Tensor, + defocus_offsets: torch.Tensor, + pixel_size_offsets: torch.Tensor, + corr_mean: torch.Tensor, + corr_std: torch.Tensor, + projective_filters: torch.Tensor, + mag_matrix: torch.Tensor | None, +) -> _RefineTemplateStackOnDevice: + """Move stack, template, and search grids to ``device`` (refine and inspect).""" + return _RefineTemplateStackOnDevice( + particle_stack_dft=particle_stack_dft.to(device), + particle_indices=particle_indices.to(device), + template_dft=template_dft.to(device), + euler_angles=euler_angles.to(device), + euler_angle_offsets=euler_angle_offsets.to(device), + defocus_u=defocus_u.to(device), + defocus_v=defocus_v.to(device), + defocus_angle=defocus_angle.to(device), + defocus_offsets=defocus_offsets.to(device), + pixel_size_offsets=pixel_size_offsets.to(device), + corr_mean=corr_mean.to(device), + corr_std=corr_std.to(device), + projective_filters=projective_filters.to(device), + mag_matrix=mag_matrix.to(device) if mag_matrix is not None else None, + ) + + +def _tqdm_for_refine_particle_loop( + num_particles: int, + device: torch.device, + device_id: int, + desc_verb: str, +) -> tqdm.tqdm: + """Progress bar over particle index for a single device (refine or inspect).""" + return tqdm.tqdm( + range(num_particles), + total=num_particles, + desc=f"{desc_verb} particles on device {device.index}...", + leave=True, + position=device_id, + dynamic_ncols=True, + unit="particle", + smoothing=0.1, + ) # NOTE: Disabling pylint for too many arguments because we are taking a data-oriented @@ -146,35 +251,37 @@ def core_refine_template( img_w = 2 * (img_w - 1) template_w = 2 * (template_w - 1) + ordered_results = [results[k] for k in sorted(results.keys(), key=int)] + # Concatenate results from all devices refined_cross_correlation = torch.cat( - [torch.from_numpy(r["refined_cross_correlation"]) for r in results.values()] + [torch.from_numpy(r["refined_cross_correlation"]) for r in ordered_results] ) refined_z_score = torch.cat( - [torch.from_numpy(r["refined_z_score"]) for r in results.values()] + [torch.from_numpy(r["refined_z_score"]) for r in ordered_results] ) refined_euler_angles = torch.cat( - [torch.from_numpy(r["refined_euler_angles"]) for r in results.values()] + [torch.from_numpy(r["refined_euler_angles"]) for r in ordered_results] ) refined_defocus_offset = torch.cat( - [torch.from_numpy(r["refined_defocus_offset"]) for r in results.values()] + [torch.from_numpy(r["refined_defocus_offset"]) for r in ordered_results] ) refined_pixel_size_offset = torch.cat( - [torch.from_numpy(r["refined_pixel_size_offset"]) for r in results.values()] + [torch.from_numpy(r["refined_pixel_size_offset"]) for r in ordered_results] ) refined_pos_y = torch.cat( - [torch.from_numpy(r["refined_pos_y"]) for r in results.values()] + [torch.from_numpy(r["refined_pos_y"]) for r in ordered_results] ) refined_pos_x = torch.cat( - [torch.from_numpy(r["refined_pos_x"]) for r in results.values()] + [torch.from_numpy(r["refined_pos_x"]) for r in ordered_results] ) # Ensure the results are sorted back to the original particle order # (If particles were split across devices, we need to reorder the results) particle_indices = torch.cat( - [torch.from_numpy(r["particle_indices"]) for r in results.values()] + [torch.from_numpy(r["particle_indices"]) for r in ordered_results] ) - angle_idx = torch.cat([torch.from_numpy(r["angle_idx"]) for r in results.values()]) + angle_idx = torch.cat([torch.from_numpy(r["angle_idx"]) for r in ordered_results]) sort_indices = torch.argsort(particle_indices) refined_cross_correlation = refined_cross_correlation[sort_indices] @@ -291,26 +398,32 @@ def construct_multi_gpu_refine_template_kwargs( # Split tensors for this device. All these tensors are per-particle, that is # the i-th element in each tensor corresponds to the i-th particle in the stack. - device_particle_stack_dft = particle_stack_dft[start_idx:end_idx] - device_euler_angles = euler_angles[start_idx:end_idx] - device_defocus_u = defocus_u[start_idx:end_idx] - 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] + # + # Move to CPU before passing to child processes: Python multiprocessing uses + # fork on Linux, and CUDA tensors shared across forked processes via CUDA IPC + # are unreliable (they silently read as zeros for non-primary GPUs). Sending CPU + # tensors avoids CUDA IPC entirely; each worker does a clean CPU-->GPU transfer + # via _move_refine_template_stack_to_device. + device_particle_stack_dft = particle_stack_dft[start_idx:end_idx].cpu() + device_euler_angles = euler_angles[start_idx:end_idx].cpu() + device_defocus_u = defocus_u[start_idx:end_idx].cpu() + device_defocus_v = defocus_v[start_idx:end_idx].cpu() + device_defocus_angle = defocus_angle[start_idx:end_idx].cpu() + device_projective_filters = projective_filters[start_idx:end_idx].cpu() + device_corr_mean = corr_mean[start_idx:end_idx].cpu() + device_corr_std = corr_std[start_idx:end_idx].cpu() kwargs = { "particle_stack_dft": device_particle_stack_dft, "particle_indices": particle_indices, - "template_dft": template_dft, + "template_dft": template_dft.cpu(), "euler_angles": device_euler_angles, - "euler_angle_offsets": euler_angle_offsets, + "euler_angle_offsets": euler_angle_offsets.cpu(), "defocus_u": device_defocus_u, "defocus_v": device_defocus_v, "defocus_angle": device_defocus_angle, - "defocus_offsets": defocus_offsets, - "pixel_size_offsets": pixel_size_offsets, + "defocus_offsets": defocus_offsets.cpu(), + "pixel_size_offsets": pixel_size_offsets.cpu(), "corr_mean": device_corr_mean, "corr_std": device_corr_std, "projective_filters": device_projective_filters, @@ -318,7 +431,7 @@ def construct_multi_gpu_refine_template_kwargs( "batch_size": batch_size, "num_cuda_streams": num_cuda_streams, "device": device, - "mag_matrix": mag_matrix, + "mag_matrix": mag_matrix.cpu() if mag_matrix is not None else None, } kwargs_per_device.append(kwargs) @@ -396,47 +509,32 @@ def _core_refine_template_single_gpu( Anisotropic magnification matrix of shape (2, 2). If None, no magnification transform is applied. Default is None. """ - streams = [torch.cuda.Stream(device=device) for _ in range(num_cuda_streams)] - - ###################################### - ### Send all tensors to the device ### - ###################################### - - particle_stack_dft = particle_stack_dft.to(device) - particle_indices = particle_indices.to(device) - template_dft = template_dft.to(device) - euler_angles = euler_angles.to(device) - euler_angle_offsets = euler_angle_offsets.to(device) - defocus_u = defocus_u.to(device) - defocus_v = defocus_v.to(device) - defocus_angle = defocus_angle.to(device) - defocus_offsets = defocus_offsets.to(device) - pixel_size_offsets = pixel_size_offsets.to(device) - corr_mean = corr_mean.to(device) - corr_std = corr_std.to(device) - projective_filters = projective_filters.to(device) - mag_matrix = mag_matrix.to(device) if mag_matrix is not None else None - - ######################################## - ### Setup constants and progress bar ### - ######################################## - - num_particles, _, img_w = particle_stack_dft.shape - _, _, template_w = template_dft.shape - # account for RFFT - img_w = 2 * (img_w - 1) - template_w = 2 * (template_w - 1) + if device.type == "cuda": + torch.cuda.set_device(device) + + streams = _make_device_streams(device, num_cuda_streams) + + refine_stack = _move_refine_template_stack_to_device( + device, + particle_stack_dft, + particle_indices, + template_dft, + euler_angles, + euler_angle_offsets, + defocus_u, + defocus_v, + defocus_angle, + defocus_offsets, + pixel_size_offsets, + corr_mean, + corr_std, + projective_filters, + mag_matrix, + ) - # tqdm progress bar - pbar_iter = tqdm.tqdm( - range(num_particles), - total=num_particles, - desc=f"Refining particles on device {device.index}...", - leave=True, - position=device_id, - dynamic_ncols=True, - unit="particle", - smoothing=0.1, + num_particles = refine_stack.particle_stack_dft.shape[0] + pbar_iter = _tqdm_for_refine_particle_loop( + num_particles, device, device_id, "Refining" ) ############################################################################# @@ -445,36 +543,32 @@ def _core_refine_template_single_gpu( refined_statistics = [] for i in pbar_iter: - particle_image_dft = particle_stack_dft[i] - particle_index = int(particle_indices[i]) # Original particle index + particle_image_dft = refine_stack.particle_stack_dft[i] # Distribute different particles across streams - stream = streams[i % num_cuda_streams] - with torch.cuda.stream(stream): + stream = streams[i % len(streams)] + with _device_stream_context(stream): refined_stats = _core_refine_template_single_thread( particle_image_dft=particle_image_dft, - particle_index=particle_index, - template_dft=template_dft, - euler_angles=euler_angles[i, :], - euler_angle_offsets=euler_angle_offsets, - defocus_u=defocus_u[i], - defocus_v=defocus_v[i], - defocus_angle=defocus_angle[i], - defocus_offsets=defocus_offsets, - pixel_size_offsets=pixel_size_offsets, + template_dft=refine_stack.template_dft, + euler_angles=refine_stack.euler_angles[i, :], + euler_angle_offsets=refine_stack.euler_angle_offsets, + defocus_u=refine_stack.defocus_u[i], + defocus_v=refine_stack.defocus_v[i], + defocus_angle=refine_stack.defocus_angle[i], + defocus_offsets=refine_stack.defocus_offsets, + pixel_size_offsets=refine_stack.pixel_size_offsets, ctf_kwargs=ctf_kwargs, - corr_mean=corr_mean[i], - corr_std=corr_std[i], - projective_filter=projective_filters[i], + corr_mean=refine_stack.corr_mean[i], + corr_std=refine_stack.corr_std[i], + projective_filter=refine_stack.projective_filters[i], batch_size=batch_size, - mag_matrix=mag_matrix, - device_id=device_id, + mag_matrix=refine_stack.mag_matrix, ) refined_statistics.append(refined_stats) # Wait for all streams to finish - for stream in streams: - stream.synchronize() + _synchronize_device_streams(streams) # For each particle, calculate the new best orientation, defocus, and position refined_cross_correlation = torch.tensor( @@ -511,10 +605,10 @@ def _core_refine_template_single_gpu( stats["refined_theta_offset"], stats["refined_psi_offset"], ], - dtype=euler_angles.dtype, + dtype=refine_stack.euler_angles.dtype, device=device, ), - euler_angles[i, :], # original angle + refine_stack.euler_angles[i, :], # original angle ) refined_euler_angles[i, :] = composed_refined_angle @@ -544,17 +638,15 @@ def _core_refine_template_single_gpu( "refined_pixel_size_offset": refined_pixel_size_offset.cpu().numpy(), "refined_pos_y": refined_pos_y.cpu().numpy(), "refined_pos_x": refined_pos_x.cpu().numpy(), - "particle_indices": particle_indices.cpu().numpy(), # Original idxs for sorting + "particle_indices": refine_stack.particle_indices.cpu().numpy(), # sort keys "angle_idx": angle_idx.cpu().numpy(), } result_dict[device_id] = result -# pylint: disable=too-many-locals, too-many-statements -def _core_refine_template_single_thread( +def _iter_refine_particle_correlation_batches( particle_image_dft: torch.Tensor, - particle_index: int, template_dft: torch.Tensor, euler_angles: torch.Tensor, euler_angle_offsets: torch.Tensor, @@ -563,119 +655,46 @@ def _core_refine_template_single_thread( defocus_angle: float, defocus_offsets: torch.Tensor, pixel_size_offsets: torch.Tensor, - corr_mean: torch.Tensor, - corr_std: torch.Tensor, ctf_kwargs: dict, projective_filter: torch.Tensor, batch_size: int = 32, - device_id: int = 0, mag_matrix: torch.Tensor | None = None, -) -> dict[str, float | int]: - """Run the single-threaded core refine template function. - - Parameters - ---------- - particle_image_dft : torch.Tensor - The real-Fourier transformed particle image. Shape of (H, W). - particle_index : int - The index of the particle in the stack. - template_dft : torch.Tensor - The template volume to extract central slices from. Real-Fourier transformed - and fftshifted. - euler_angles : torch.Tensor - The previous best euler angle for the particle. Shape of (3,). - euler_angle_offsets : torch.Tensor - The Euler angle offsets to apply to each particle. Shape of (k, 3). - defocus_u : float - The defocus along the major axis for the particle. - defocus_v : float - The defocus along the minor for the particle. - defocus_angle : float - The defocus astigmatism angle for the particle. - defocus_offsets : torch.Tensor - The defocus offsets to search over for each particle. Shape of (l,). - pixel_size_offsets : torch.Tensor - The pixel size offsets to search over for each particle. Shape of (m,). - corr_mean : torch.Tensor - The mean of the cross-correlation values from the full orientation search - for the pixels around the center of the particle. - corr_std : torch.Tensor - The standard deviation of the cross-correlation values from the full - orientation search for the pixels around the center of the particle. - ctf_kwargs : dict - Keyword arguments to pass to the CTF calculation function. - projective_filter : torch.Tensor - Projective filters to apply to the Fourier slice particle. Shape of (h, w). - batch_size : int, optional - The number of orientations to cross-correlate at once. Default is 32. - device_id : int, optional - The ID of the device/process. Default is 0. - mag_matrix : torch.Tensor | None, optional - Anisotropic magnification matrix of shape (2, 2). If None, - no magnification transform is applied. Default is None. - - Returns - ------- - dict[str, float | int] - The refined statistics for the particle. + apply_projection_normalization: bool = True, +) -> Iterator[tuple[int, torch.Tensor, torch.Tensor, int, int]]: + """Yield batched local correlations for one particle using refine semantics. + + Note + ---- + Only correlations are computed here, z-score normalization must happen externally to + this function. """ img_h, img_w = particle_image_dft.shape _, template_h, template_w = template_dft.shape - # account for RFFT img_w = 2 * (img_w - 1) template_w = 2 * (template_w - 1) - # valid crop shape crop_h = img_h - template_h + 1 crop_w = img_w - template_w + 1 - # Output best statistics - max_cc = -1e9 - max_z_score = -1e9 - refined_phi_offset = 0.0 - refined_theta_offset = 0.0 - refined_psi_offset = 0.0 - full_angle_idx = 0 - refined_defocus_offset = 0.0 - refined_pixel_size_offset = 0.0 - refined_pos_y = 0 - refined_pos_x = 0 - - # The "best" Euler angle from the match template program default_rot_matrix = roma.euler_to_rotmat( EULER_ANGLE_FMT, euler_angles, degrees=True, device=particle_image_dft.device ) - default_rot_matrix = default_rot_matrix.to(torch.float32) - # Calculate the CTF filters with the relative offsets + + ctf_dev_kwargs = move_ctf_kwargs_tensors_to_device( + ctf_kwargs, particle_image_dft.device + ) ctf_filters = calculate_ctf_filter_stack_full_args( - defocus_u=defocus_u, # in Angstrom - defocus_v=defocus_v, # in Angstrom - astigmatism_angle=defocus_angle, # in degrees - defocus_offsets=defocus_offsets, # in Angstrom - pixel_size_offsets=pixel_size_offsets, # in Angstrom - **ctf_kwargs, + defocus_u=defocus_u, + defocus_v=defocus_v, + astigmatism_angle=defocus_angle, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + **ctf_dev_kwargs, ) - - # Combine the single projective filter with the CTF filter combined_projective_filter = projective_filter[None, None, ...] * ctf_filters - # Iterate over the Euler angle offsets in batches - # The tqdm iterator is over batches, but we want to report cross-correlations/sec. - # We therefore scale by the number of cross-correlations per batch. num_batches = math.ceil(euler_angle_offsets.shape[0] / batch_size) - cross_corr_per_batch = len(defocus_offsets) * len(pixel_size_offsets) * batch_size - - tqdm_iter = tqdm.tqdm( - range(num_batches), - total=num_batches, - desc=f"Refining particle {particle_index} on device {device_id}", - leave=False, - position=device_id + torch.cuda.device_count(), - unit="corr", - unit_scale=cross_corr_per_batch, - ) - - for i in tqdm_iter: + for i in range(num_batches): start_idx = i * batch_size end_idx = min((i + 1) * batch_size, euler_angle_offsets.shape[0]) euler_angle_offsets_batch = euler_angle_offsets[start_idx:end_idx] @@ -686,21 +705,17 @@ def _core_refine_template_single_thread( device=particle_image_dft.device, ) rot_matrix_batch = rot_matrix_batch.to(torch.float32) - - # Rotate the default (best) orientation by the offsets rot_matrix_batch = roma.rotmat_composition( (rot_matrix_batch, default_rot_matrix) ) - # Calculate the cross-correlation if particle_image_dft.device.type == "cuda": - # NOTE: Here we are setting to only a single stream, but this can easily - # be extended to multiple streams if needed. cross_correlation = do_batched_orientation_cross_correlate( image_dft=particle_image_dft, template_dft=template_dft, rotation_matrices=rot_matrix_batch, projective_filters=combined_projective_filter, + apply_normalization=apply_projection_normalization, mag_matrix=mag_matrix, ) else: @@ -709,37 +724,61 @@ def _core_refine_template_single_thread( template_dft=template_dft, rotation_matrices=rot_matrix_batch, projective_filters=combined_projective_filter, + apply_normalization=apply_projection_normalization, mag_matrix=mag_matrix, ) - cross_correlation = cross_correlation[..., :crop_h, :crop_w] # valid crop + cross_correlation = cross_correlation[..., :crop_h, :crop_w] - # Scale cross_correlation to be "z-score"-like + yield ( + start_idx, + euler_angle_offsets_batch, + cross_correlation, + crop_h, + crop_w, + ) + + +def _reduce_refine_best_zscore( + correlation_batches: Iterator[tuple[int, torch.Tensor, torch.Tensor, int, int]], + corr_mean: torch.Tensor, + corr_std: torch.Tensor, + defocus_offsets: torch.Tensor, + pixel_size_offsets: torch.Tensor, +) -> dict[str, float | int]: + """Reduce local correlation batches to the current refine-template best result.""" + max_cc = -1e9 + max_z_score = -1e9 + refined_phi_offset = 0.0 + refined_theta_offset = 0.0 + refined_psi_offset = 0.0 + full_angle_idx = 0 + refined_defocus_offset = 0.0 + refined_pixel_size_offset = 0.0 + refined_pos_y = 0 + refined_pos_x = 0 + + for ( + start_idx, + euler_angle_offsets_batch, + cross_correlation, + crop_h, + crop_w, + ) in correlation_batches: z_score = (cross_correlation - corr_mean) / corr_std - # shape xc is (num_Cs, num_defocus, num_orientations, y, x) - # where num_Cs is the number of different pixel size offsets, - # num_defocus is the number of defocus offsets, - # and num_orientations is the number of Euler angle offsets. - # Update the best refined statistics (only if max is greater than previous) if z_score.max() > max_z_score: max_cc = cross_correlation.max() max_z_score = z_score.max() - - # Find the maximum value and its indices max_values, max_indices = torch.max(z_score.view(-1, crop_h, crop_w), dim=0) - - # Get the overall maximum value and its position _, max_pos = torch.max(max_values.view(-1), dim=0) y_idx, x_idx = max_pos // crop_w, max_pos % crop_w - # Calculate the indices for each dimension flat_idx = max_indices[y_idx, x_idx] - px_idx = flat_idx // (len(defocus_offsets) * len(euler_angle_offsets_batch)) - defocus_idx = (flat_idx // len(euler_angle_offsets_batch)) % len( - defocus_offsets - ) - angle_idx = flat_idx % len(euler_angle_offsets_batch) + num_angles_batch = len(euler_angle_offsets_batch) + px_idx = flat_idx // (len(defocus_offsets) * num_angles_batch) + defocus_idx = (flat_idx // num_angles_batch) % len(defocus_offsets) + angle_idx = flat_idx % num_angles_batch refined_phi_offset = euler_angle_offsets_batch[angle_idx, 0] refined_theta_offset = euler_angle_offsets_batch[angle_idx, 1] @@ -750,8 +789,7 @@ def _core_refine_template_single_thread( refined_pos_x = x_idx full_angle_idx = angle_idx + start_idx - # Return the refined statistics - refined_stats = { + return { "max_cc": max_cc, "max_z_score": max_z_score, "refined_phi_offset": refined_phi_offset, @@ -764,7 +802,92 @@ def _core_refine_template_single_thread( "angle_idx": full_angle_idx, } - return refined_stats + +# pylint: disable=too-many-locals, too-many-statements +def _core_refine_template_single_thread( + particle_image_dft: torch.Tensor, + template_dft: torch.Tensor, + euler_angles: torch.Tensor, + euler_angle_offsets: torch.Tensor, + defocus_u: float, + defocus_v: float, + defocus_angle: float, + defocus_offsets: torch.Tensor, + pixel_size_offsets: torch.Tensor, + corr_mean: torch.Tensor, + corr_std: torch.Tensor, + ctf_kwargs: dict, + projective_filter: torch.Tensor, + batch_size: int = 32, + mag_matrix: torch.Tensor | None = None, +) -> dict[str, float | int]: + """Run the single-threaded core refine template function. + + Parameters + ---------- + particle_image_dft : torch.Tensor + The real-Fourier transformed particle image. Shape of (H, W). + template_dft : torch.Tensor + The template volume to extract central slices from. Real-Fourier transformed + and fftshifted. + euler_angles : torch.Tensor + The previous best euler angle for the particle. Shape of (3,). + euler_angle_offsets : torch.Tensor + The Euler angle offsets to apply to each particle. Shape of (k, 3). + defocus_u : float + The defocus along the major axis for the particle. + defocus_v : float + The defocus along the minor for the particle. + defocus_angle : float + The defocus astigmatism angle for the particle. + defocus_offsets : torch.Tensor + The defocus offsets to search over for each particle. Shape of (l,). + pixel_size_offsets : torch.Tensor + The pixel size offsets to search over for each particle. Shape of (m,). + corr_mean : torch.Tensor + The mean of the cross-correlation values from the full orientation search + for the pixels around the center of the particle. + corr_std : torch.Tensor + The standard deviation of the cross-correlation values from the full + orientation search for the pixels around the center of the particle. + ctf_kwargs : dict + Keyword arguments to pass to the CTF calculation function. + projective_filter : torch.Tensor + Projective filters to apply to the Fourier slice particle. Shape of (h, w). + batch_size : int, optional + The number of orientations to cross-correlate at once. Default is 32. + mag_matrix : torch.Tensor | None, optional + Anisotropic magnification matrix of shape (2, 2). If None, + no magnification transform is applied. Default is None. + + Returns + ------- + dict[str, float | int] + The refined statistics for the particle. + """ + correlation_batches = _iter_refine_particle_correlation_batches( + particle_image_dft=particle_image_dft, + template_dft=template_dft, + euler_angles=euler_angles, + euler_angle_offsets=euler_angle_offsets, + defocus_u=defocus_u, + defocus_v=defocus_v, + defocus_angle=defocus_angle, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + ctf_kwargs=ctf_kwargs, + projective_filter=projective_filter, + batch_size=batch_size, + mag_matrix=mag_matrix, + ) + + return _reduce_refine_best_zscore( + correlation_batches=correlation_batches, + corr_mean=corr_mean, + corr_std=corr_std, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + ) # pylint: disable=too-many-locals diff --git a/src/leopard_em/backend/cross_correlation.py b/src/leopard_em/backend/cross_correlation.py index c078f881..e05a6dfe 100644 --- a/src/leopard_em/backend/cross_correlation.py +++ b/src/leopard_em/backend/cross_correlation.py @@ -1,6 +1,7 @@ """File containing Fourier-slice based cross-correlation functions for 2DTM.""" import torch +from torch_fourier_shell_correlation import fsc from torch_fourier_slice import extract_central_slices_rfft_3d, transform_slice_2d from leopard_em.backend.utils import ( @@ -8,6 +9,26 @@ 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 = [ # NOTE: restrictive to 4k images and 512 templates + 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 + ZIPFFT_AVAILABLE = True +except ImportError: + zipfft = None + ZIPFFT_SUPPORTED_BATCH_SIZES = [] + ZIPFFT_SUPPORTED_CONFIGS = [] + ZIPFFT_AVAILABLE = False + # pylint: disable=too-many-locals,E1102 def do_streamed_orientation_cross_correlate( @@ -16,6 +37,7 @@ def do_streamed_orientation_cross_correlate( rotation_matrices: torch.Tensor, projective_filters: torch.Tensor, streams: list[torch.cuda.Stream], + apply_normalization: bool = True, mag_matrix: torch.Tensor | None = None, ) -> torch.Tensor: """Calculates a grid of 2D cross-correlations over multiple CUDA streams. @@ -47,6 +69,8 @@ def do_streamed_orientation_cross_correlate( streams : list[torch.cuda.Stream] List of CUDA streams to use for parallel computation. Each stream will handle a separate cross-correlation. + apply_normalization : bool, optional + Whether to apply normalization to the template projections, by default True mag_matrix : torch.Tensor | None, optional Anisotropic magnification matrix of shape (2, 2). If None, no magnification transform is applied. Default is None. @@ -115,10 +139,19 @@ def do_streamed_orientation_cross_correlate( fourier_slice_filtered = fourier_slice * projective_filters[k, j] projection = torch.fft.irfft2(fourier_slice_filtered) projection = torch.fft.ifftshift(projection, dim=(-2, -1)) - projection = normalize_template_projection_compiled( - projection, - projection_shape_real, - image_shape_real, + + if apply_normalization: + projection = normalize_template_projection_compiled( + projection, + projection_shape_real, + image_shape_real, + ) + + # NOTE: Decomposing 2D FFT into component 1D FFTs. Saves on first + # pass where many lines are zeros. Approx 6-8% speedup. + temp_fft = torch.fft.rfft(projection, n=image_shape_real[1], dim=-1) + projection_dft = torch.fft.fft( + temp_fft, n=image_shape_real[0], dim=-2 ) # NOTE: Decomposing 2D FFT into component 1D FFTs. Saves on first @@ -160,6 +193,7 @@ def do_batched_orientation_cross_correlate( template_dft: torch.Tensor, rotation_matrices: torch.Tensor, projective_filters: torch.Tensor, + apply_normalization: bool = True, requires_grad: bool = False, mag_matrix: torch.Tensor | None = None, ) -> torch.Tensor: @@ -188,6 +222,8 @@ def do_batched_orientation_cross_correlate( 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. + apply_normalization : bool, optional + Whether to apply normalization to the template projections, by default True requires_grad : bool, optional Whether the input is requires_grad. Default is False. If True, the input will be cloned before any in-place operations. @@ -250,11 +286,13 @@ def do_batched_orientation_cross_correlate( # 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, - ) + + if apply_normalization: + projections = normalize_template_projection_compiled( + projections, + projection_shape_real, + image_shape_real, + ) for j in range(num_defocus): for k in range(num_Cs): @@ -289,6 +327,7 @@ def do_batched_orientation_cross_correlate_cpu( template_dft: torch.Tensor, rotation_matrices: torch.Tensor, projective_filters: torch.Tensor, + apply_normalization: bool = True, mag_matrix: torch.Tensor | None = None, ) -> torch.Tensor: """Same as `do_streamed_orientation_cross_correlate` but on the CPU. @@ -313,6 +352,8 @@ def do_batched_orientation_cross_correlate_cpu( projective_filters : torch.Tensor Multiplied 'ctf_filters' with 'whitening_filter_template'. Has shape (defocus_batch, h, w // 2 + 1). Is RFFT and not fftshifted. + apply_normalization : bool, optional + Whether to apply normalization to the template projections, by default True mag_matrix : torch.Tensor | None, optional Anisotropic magnification matrix of shape (2, 2). If None, no magnification transform is applied. Default is None. @@ -352,11 +393,13 @@ def do_batched_orientation_cross_correlate_cpu( # 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( - projections, - projection_shape_real, - image_shape_real, - ) + + if apply_normalization: + projections = normalize_template_projection( + projections, + projection_shape_real, + image_shape_real, + ) # Padded forward Fourier transform for cross-correlation projections_dft = torch.fft.rfftn(projections, dim=(-2, -1), s=image_shape_real) @@ -367,3 +410,243 @@ def do_batched_orientation_cross_correlate_cpu( cross_correlation = torch.fft.irfftn(projections_dft, dim=(-2, -1)) return cross_correlation + + +# pylint: disable=too-many-locals,E1102 +def do_batched_orientation_frc( + image_dft: torch.Tensor, + template_dft: torch.Tensor, + rotation_matrices: torch.Tensor, + projective_filters: torch.Tensor, + apply_normalization: bool = True, + mag_matrix: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Batched projection and Fourier ring correlation with fixed filters. + + Parameters + ---------- + image_dft : torch.Tensor + Real-Fourier transform (RFFT) of the image with large-image filters applied. + template_dft : torch.Tensor + Real-Fourier transform (RFFT) of the template volume. + rotation_matrices : torch.Tensor + Rotation matrices for orientation offsets. Shape (num_orientations, 3, 3). + projective_filters : torch.Tensor + Filter stack with shape (num_Cs, num_defocus, h, w // 2 + 1). + apply_normalization : bool, optional + Whether to normalize real-space projections before FRC. + mag_matrix : torch.Tensor | None, optional + Optional anisotropic magnification matrix. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + - frc_values: shape (num_Cs, num_defocus, num_orientations, num_freq_bins) + - frequency_bins: shape (num_freq_bins,) + """ + projection_shape_real = (template_dft.shape[1], template_dft.shape[2] * 2 - 2) + image_shape_real = (image_dft.shape[0], image_dft.shape[1] * 2 - 2) + + num_orientations = rotation_matrices.shape[0] + num_cs = projective_filters.shape[0] + num_defocus = projective_filters.shape[1] + + fourier_slice = extract_central_slices_rfft_3d( + volume_rfft=template_dft, + rotation_matrices=rotation_matrices, + ) + if mag_matrix is not None: + rfft_shape = (template_dft.shape[1], template_dft.shape[2]) + stack_shape = (num_orientations,) + fourier_slice = transform_slice_2d( + projection_image_dfts=fourier_slice, + rfft_shape=rfft_shape, + stack_shape=stack_shape, + transform_matrix=mag_matrix, + ) + fourier_slice = torch.fft.ifftshift(fourier_slice, dim=(-2,)) + fourier_slice[..., 0, 0] = 0 + 0j + fourier_slice *= -1 + + fourier_slice = fourier_slice[None, None, ...] * projective_filters[:, :, None, ...] + + projections = torch.fft.irfftn(fourier_slice, dim=(-2, -1)) + projections = torch.fft.ifftshift(projections, dim=(-2, -1)) + + if apply_normalization: + if image_dft.device.type == "cuda": + projections = normalize_template_projection_compiled( + projections, + projection_shape_real, + image_shape_real, + ) + else: + projections = normalize_template_projection( + projections, + projection_shape_real, + image_shape_real, + ) + + image_real = torch.fft.irfftn(image_dft, dim=(-2, -1), s=image_shape_real).real + frequency_bins = torch.fft.rfftfreq(image_shape_real[-1], device=image_dft.device) + num_freq_bins = int(frequency_bins.shape[0]) + frc_values = torch.empty( + (num_cs, num_defocus, num_orientations, num_freq_bins), + dtype=image_real.dtype, + device=image_dft.device, + ) + + for j in range(num_defocus): + for k in range(num_cs): + for i in range(num_orientations): + frc_values[k, j, i] = fsc( + projections[k, j, i].real, + image_real, + ) + + return frc_values, frequency_bins + + +# 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). + + Raises + ------ + ImportError + If the optional ``zipfft`` package is not installed. + """ + if zipfft is None: + raise ImportError( + "backend='zipfft' requires the optional 'zipfft' package, which is not " + "installed in this environment. Install it from " + "https://github.com/mgiammar/zipFFT, or select a different backend " + "('batched' or 'streamed')." + ) + + # 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, + 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..e3a0d77f 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,41 @@ def aggregate_distributed_results( correlation_sum = torch.from_numpy(correlation_sum) correlation_squared_sum = torch.from_numpy(correlation_squared_sum) + # Concatenate the per-device/per-rank correlation table entries + per_key_values: dict[str, list[torch.Tensor]] = {} + threshold = None + for result in results: + correlation_table = result["correlation_table"] + correlation_table = ( + correlation_table.cpu().to_dict() + if isinstance(correlation_table, tensordict.TensorDict) + else dict(correlation_table) + ) + threshold = correlation_table.pop("threshold") + for key, value in correlation_table.items(): + per_key_values.setdefault(key, []).append(torch.as_tensor(value)) + + full_correlation_table = { + key: torch.cat(values) for key, values in per_key_values.items() + } + 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 +105,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 +114,63 @@ 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 + + # 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(), + } - 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..e4efc799 100644 --- a/src/leopard_em/backend/utils.py +++ b/src/leopard_em/backend/utils.py @@ -6,6 +6,7 @@ from typing import Any, Callable, TypeVar import roma +import tensordict import torch # Suppress the specific deprecation warnings from PyTorch internals @@ -188,25 +189,135 @@ 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, + compute_correlation_table: 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. + compute_correlation_table : bool, optional + Whether to find and return threshold exceedances for the correlation table. + """ + # 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) + if compute_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] + else: + empty_int = torch.empty(0, dtype=torch.int64, device=cc_reshaped.device) + y_idxs = x_idxs = empty_int + values = torch.empty(0, dtype=cc_reshaped.dtype, device=cc_reshaped.device) + global_idxs = current_indexes[empty_int] + + 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, + compute_correlation_table: 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 +333,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 +344,62 @@ 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. + compute_correlation_table : bool, optional + Whether to threshold the cross-correlation values and add exceedances to the + correlation table. """ - 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, + compute_correlation_table=compute_correlation_table, ) - 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/config/correlation_filters.py b/src/leopard_em/pydantic_models/config/correlation_filters.py index c138ead6..8226685d 100644 --- a/src/leopard_em/pydantic_models/config/correlation_filters.py +++ b/src/leopard_em/pydantic_models/config/correlation_filters.py @@ -336,6 +336,50 @@ def calculate_arbitrary_curve_filter( return filter_ndim +class RandomFourierDropoutConfig(BaseModel2DTM): + """Configuration for random Fourier dropout (masking). + + NOTE: A manual random seed *must* be provided to ensure the post-processing stage + knows which Fourier coefficients were dropped during template matching. You can + choose a seed here randomly. + + Attributes + ---------- + enabled : bool + If True, apply random Fourier dropout to the input image and template + projections. Default is False. + dropout_fraction : float + Fraction of Fourier coefficients to randomly set to zero. Must be between + 0 and 1. + """ + + enabled: bool = False + dropout_fraction: Annotated[float, Field(ge=0.0, le=1.0)] = 0.1 + seed: int + + def calculate_dropout_mask(self, output_shape: tuple[int, ...]) -> torch.Tensor: + """Calculates a random dropout mask for the desired output shape. + + Parameters + ---------- + output_shape : tuple[int, ...] + Desired output shape of the dropout mask (in terms of fourier-space shape). + + Returns + ------- + torch.Tensor + The dropout mask for the desired output shape. + """ + if not self.enabled: + return torch.ones(output_shape, dtype=torch.float32) + + gen = torch.Generator().manual_seed(self.seed) + probs = torch.rand(size=output_shape, generator=gen) + mask = probs > self.dropout_fraction + + return mask.float() + + class PreprocessingFilters(BaseModel2DTM): """Configuration class for all preprocessing filters. @@ -362,9 +406,15 @@ class PreprocessingFilters(BaseModel2DTM): PhaseRandomizationFilterConfig() ) arbitrary_curve_filter: ArbitraryCurveFilterConfig = ArbitraryCurveFilterConfig() + random_fourier_dropout: RandomFourierDropoutConfig = RandomFourierDropoutConfig( + seed=0 + ) def get_combined_filter( - self, ref_img_rfft: torch.Tensor, output_shape: tuple[int, ...] + self, + ref_img_rfft: torch.Tensor, + output_shape: tuple[int, ...], + apply_random_dropout: bool = True, ) -> torch.Tensor: """Combine all filters into a single filter. @@ -376,6 +426,10 @@ def get_combined_filter( Desired output shape of the combined filter in Fourier space. This is the filter shape in Fourier space *not* real space (like in the torch_fourier_filter package). + apply_random_dropout : bool, optional + Whether to include the random Fourier dropout mask in the combined filter. + Pass ``False`` to skip dropout (e.g. for the image side of match-template) + without mutating the configured dropout state. Defaults to ``True``. Returns ------- @@ -407,4 +461,10 @@ def get_combined_filter( * arbitrary_curve_filter_tensor ) + if apply_random_dropout: + random_dropout_mask = self.random_fourier_dropout.calculate_dropout_mask( + output_shape=output_shape + ).to(device) + combined_filter = combined_filter * random_dropout_mask + return combined_filter diff --git a/src/leopard_em/pydantic_models/config/movie_config.py b/src/leopard_em/pydantic_models/config/movie_config.py index 6d91c46f..413f8990 100644 --- a/src/leopard_em/pydantic_models/config/movie_config.py +++ b/src/leopard_em/pydantic_models/config/movie_config.py @@ -1,7 +1,7 @@ """Serialization and validation of movie parameters for 2DTM.""" import torch -from torch_motion_correction.data_io import read_deformation_field_from_csv +from torch_motion_correction.deformation_field import DeformationField from leopard_em.pydantic_models.custom_types import BaseModel2DTM from leopard_em.utils.data_io import load_mrc_volume @@ -53,4 +53,5 @@ def deformation_field(self) -> torch.Tensor | None: if self.particle_shifts_path: # Particle shifts take precedence, so don't load deformation field return None - return read_deformation_field_from_csv(self.deformation_field_path) + deformation_field = DeformationField.from_csv(self.deformation_field_path) + return deformation_field.data diff --git a/src/leopard_em/pydantic_models/custom_types.py b/src/leopard_em/pydantic_models/custom_types.py index 5372d640..b65484da 100644 --- a/src/leopard_em/pydantic_models/custom_types.py +++ b/src/leopard_em/pydantic_models/custom_types.py @@ -14,6 +14,12 @@ Annotated[Optional[torch.Tensor], Field(default=None, exclude=True)] ] +# Same as ExcludedTensor, but for a name -> tensor mapping (e.g. per-column and +# per-particle statistics maps of possibly differing shapes). +ExcludedTensorDict = SkipJsonSchema[ + Annotated[dict[str, torch.Tensor], Field(default_factory=dict, exclude=True)] +] + class BaseModel2DTM(BaseModel): """Implementation of a Pydantic BaseModel with additional, useful methods. diff --git a/src/leopard_em/pydantic_models/data_structures/__init__.py b/src/leopard_em/pydantic_models/data_structures/__init__.py index 6538ed01..03081128 100644 --- a/src/leopard_em/pydantic_models/data_structures/__init__.py +++ b/src/leopard_em/pydantic_models/data_structures/__init__.py @@ -1,9 +1,17 @@ """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, + export_particle_stack, +) __all__ = [ "OpticsGroup", "ParticleStack", + "ParticleStackCSV", + "ParticleStackHDF5", + "export_particle_stack", ] 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..3ab4a46f 100644 --- a/src/leopard_em/pydantic_models/data_structures/particle_stack.py +++ b/src/leopard_em/pydantic_models/data_structures/particle_stack.py @@ -1,29 +1,50 @@ -"""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 os 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 from torch_grid_utils import coordinate_grid from torch_motion_correction.correct_motion import get_pixel_shifts -from torch_motion_correction.deformation_field_utils import ( - evaluate_deformation_field_at_t, -) +from torch_motion_correction.deformation_field import DeformationField +from typing_extensions import Self from leopard_em.pydantic_models.config import PreprocessingFilters from leopard_em.pydantic_models.custom_types import ( BaseModel2DTM, ExcludedTensor, + ExcludedTensorDict, +) +from leopard_em.pydantic_models.formats import ( + MATCH_TEMPLATE_DF_COLUMN_ORDER, + STATISTIC_MAP_PATH_COLUMNS, ) -from leopard_em.pydantic_models.formats import MATCH_TEMPLATE_DF_COLUMN_ORDER from leopard_em.utils.data_io import load_mrc_image from leopard_em.utils.image_processing import dose_weight_movie_to_micrograph @@ -33,21 +54,174 @@ "replicate": "edge", } +_HDF5_PARTICLES_GROUP = "particles" +_HDF5_LOCAL_STATS_GROUP = "local_stats" +_HDF5_IMAGE_STACK_DATASET = "image_stack" +_HDF5_STRING_DTYPE = h5py.string_dtype() + +# Full-micrograph 2DTM result maps extracted per-particle by +# ``get_local_stat_maps`` when no explicit columns are requested. +_DEFAULT_LOCAL_STAT_COLUMNS = tuple(STATISTIC_MAP_PATH_COLUMNS) + + +# 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 _check_output_path(path: str, allow_file_overwrite: bool) -> None: + """Ensure ``path``'s parent directory is writable and overwrite policy is met. + + Creates the parent directory if it does not already exist. + + Raises + ------ + ValueError + If the parent directory is not writable, or the file already exists + and ``allow_file_overwrite`` is False. """ - return bool(s.isna().any() or s.isin([float("inf"), float("-inf")]).any()) + directory = str(Path(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 '{path}'." + ) + if not allow_file_overwrite and os.path.exists(path): + raise ValueError( + f"File '{path}' already exists but 'allow_file_overwrite' is False." + ) + + +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"]``. + """ + 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 +304,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 +370,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 +422,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 +438,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 +445,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 : ExcludedTensorDict + Per-particle local statistic maps, keyed by the ``*_path`` DataFrame column they + were derived from (e.g. ``"mip_path"``, ``"correlation_average_path"``). Each + value has 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``. + Populated on demand via :meth:`get_local_stat_maps` -- assign its return value + (or a subset of it) here to make those maps available for + :meth:`_stored_local_stat_map` lookups and, for ``ParticleStackHDF5``, for + ``to_hdf5(include_local_stats=True)``. + 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: ExcludedTensorDict - 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 +762,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,11 +770,93 @@ 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 - def construct_image_stack( + def _stored_local_stat_map(self, column: str) -> torch.Tensor | None: + """Return an already-computed local stat map for ``column``, if any. + + Parameters + ---------- + column : str + Path column name (e.g. ``"correlation_average_path"``). + + Returns + ------- + torch.Tensor | None + The stored ``(num_particles, valid_h, valid_w)`` map, or None if it + must be derived from the referenced result file. + """ + return self.local_stats.get(column) + + def get_local_stat_maps( + self, + columns: list[str] | None = None, + device: torch.device | str = "cpu", + valid_size: tuple[int, int] | None = None, + padding_value: float = 0.0, + ) -> dict[str, torch.Tensor]: + """Extract per-particle local sub-images for arbitrary result-map columns. + + Generalizes the correlation-statistics extraction to any full-micrograph result + map exposed as a ``*_path`` column. Returned regions are cropped to the valid + cross-correlation region based on extracted box size and original template size, + unless ``valid_size`` is explicitly provided. + + Parameters + ---------- + columns : list[str] | None + Path columns to extract. Defaults to the six standard 2DTM statistic + maps (``mip``, ``scaled_mip``, ``psi``, ``theta``, ``phi``, ``defocus``). + device : torch.device | str + Target device for the returned tensors. Defaults to ``"cpu"``. + valid_size : tuple[int, int] | None + Extraction size ``(height, width)``. Defaults to the valid + cross-correlation region + ``extracted_box_size - original_template_size + 1``. + padding_value : float + Constant pad value for out-of-bounds regions. Defaults to ``0.0``. + + Returns + ------- + dict[str, torch.Tensor] + Maps each requested column name to its + ``(num_particles, valid_h, valid_w)`` sub-image stack on ``device``. + """ + if columns is None: + columns = list(_DEFAULT_LOCAL_STAT_COLUMNS) + + use_default_valid_size = valid_size is None + if valid_size is None: + box_h, box_w = self.extracted_box_size + h, w = self.original_template_size + valid_size = (box_h - h + 1, box_w - w + 1) + + device = torch.device(device) + + stat_maps: dict[str, torch.Tensor] = {} + for column in columns: + stored = ( + self._stored_local_stat_map(column) if use_default_valid_size else None + ) + if stored is not None: + stat_maps[column] = stored.to(device) + continue + + images, indices = self.load_images_grouped_by_column(column) + stat_maps[column] = self._crop_particle_regions( + images=images, + indices=indices, + extraction_size=valid_size, + pos_reference="top-left", + handle_bounds="pad", + padding_mode="constant", + padding_value=padding_value, + ).to(device) + + return stat_maps + + def _crop_particle_regions( self, images: torch.Tensor, indices: list[pd.Index], @@ -447,7 +866,7 @@ def construct_image_stack( padding_mode: Literal["constant", "reflect", "replicate"] = "constant", padding_value: float = 0.0, ) -> torch.Tensor: - """Construct stack of images from the DataFrame (updates image_stack in-place). + """Crop a per-particle sub-image from each source image, read-only. This method preferentially selects refined position columns by default (refined_pos_x, refined_pos_y) if they are present in the DataFrame, falling @@ -493,86 +912,55 @@ 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) + region_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, @@ -583,18 +971,68 @@ def construct_image_stack( padding_mode=padding_mode, padding_value=padding_value, ) - image_stack[indexes] = cropped_images - - self.image_stack = image_stack + # ``indexes`` holds DataFrame index *labels* (string ``particle_id`` for + # HDF5-backed stacks); ``region_stack`` is positional, so map labels to + # 0-based row positions. For a RangeIndex (CSV-backed stacks) this is an + # identity map. + positions = self._df.index.get_indexer(indexes) + region_stack[positions] = cropped_images - return image_stack + return region_stack - def construct_image_filters( + def construct_image_stack( self, - preprocess_filters: PreprocessingFilters, - output_shape: tuple[int, int], - images_dft: torch.Tensor, - ) -> torch.Tensor: + images: torch.Tensor, + indices: list[pd.Index], + extraction_size: tuple[int, int], + pos_reference: Literal["center", "top-left"] = "top-left", + handle_bounds: Literal["pad", "error"] = "pad", + padding_mode: Literal["constant", "reflect", "replicate"] = "constant", + padding_value: float = 0.0, + ) -> torch.Tensor: + """Construct stack of particle images from the DataFrame. + + Parameters + ---------- + images : torch.Tensor + A tensor of loaded images with shape (N, H, W). + indices : list[pd.Index] + Row indexes for particles from each corresponding image. + extraction_size : tuple[int, int] + Size of the extracted boxes in pixels (height, width). + pos_reference : Literal["center", "top-left"], optional + 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 + Padding mode when ``handle_bounds="pad"``, by default "constant". + padding_value : float, optional + Constant padding value, by default 0.0. + + Returns + ------- + torch.Tensor + Stack of extracted images ``(N, extraction_h, extraction_w)``. + """ + image_stack = self._crop_particle_regions( + images=images, + indices=indices, + extraction_size=extraction_size, + pos_reference=pos_reference, + handle_bounds=handle_bounds, + padding_mode=padding_mode, + padding_value=padding_value, + ) + self.image_stack = image_stack + + return image_stack + + def construct_image_filters( + self, + preprocess_filters: PreprocessingFilters, + output_shape: tuple[int, int], + images_dft: torch.Tensor, + ) -> torch.Tensor: """Get stack of Fourier filters from filter config and reference images. Note that here the filters are assumed to be applied globally (i.e. no local @@ -621,7 +1059,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 +1090,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 +1118,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 +1131,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 +1166,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), @@ -951,7 +1183,7 @@ def _process_single_frame_with_shifts_checkpoint( def compute_frame_particle_shifts_from_deformation( self, movie_frame: torch.Tensor, - deformation_field: CubicCatmullRomGrid3d, + deformation_field: DeformationField, normalized_t_value: torch.Tensor, pixel_grid: torch.Tensor, pixel_spacing: float, @@ -960,38 +1192,36 @@ 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, - t=normalized_t_value.item(), + frame_deformation_field = deformation_field.evaluate_at_t( + t=float(normalized_t_value.item()), grid_shape=(10 * gh, 10 * gw), ) @@ -1011,78 +1241,76 @@ def compute_frame_particle_shifts_from_deformation( # pylint: disable=too-many-positional-arguments # pylint: disable=too-many-statements # pylint: disable=too-many-branches - def construct_image_stack_from_movie( + def _construct_particle_movie_rfft_stack( self, movie: torch.Tensor, - deformation_field: CubicCatmullRomGrid3d | None = None, + deformation_field: DeformationField | None = None, particle_shifts: torch.Tensor | None = None, pos_reference: Literal["center", "top-left"] = "top-left", handle_bounds: Literal["pad", "error"] = "pad", padding_mode: Literal["constant", "reflect", "replicate"] = "constant", padding_value: float = 0.0, - pre_exposure: float = 0.0, - fluence_per_frame: float = 0.0, use_gradient_checkpointing: bool = True, particle_indices: list[int] | None = None, - ) -> torch.Tensor: - """Construct a stack of images from a movie file. + require_motion_source: bool = True, + normalized_t_values: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[Any]]: + """Construct per-particle movie frame DFTs after optional motion shifts. Parameters ---------- movie : torch.Tensor The movie tensor. - deformation_field : CubicCatmullRomGrid3d | None, optional + deformation_field : DeformationField | 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. + require_motion_source : bool, optional + If True, raises an error if neither ``deformation_field`` nor + ``particle_shifts`` is provided. If False, assumes the movie is already + aligned and extracts each frame without shifts. + normalized_t_values : torch.Tensor | None, optional + Normalized time values for each frame, shape ``(t,)``. If None, a linear + ramp from 0 to 1 is constructed and 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): + if deformation_field is not None and particle_shifts is not None: raise ValueError( - "One of `deformation_field` or `particle_shifts` must be provided." + "Only one of `deformation_field` or `particle_shifts` can be provided." + ) + if deformation_field is None and particle_shifts is None: + if require_motion_source: + raise ValueError( + "One of `deformation_field` or `particle_shifts` must be provided." + ) + warnings.warn( + "No deformation field or particle shifts were provided. Assuming the " + "movie is already aligned and extracting each frame without shifts.", + stacklevel=2, ) 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 @@ -1090,23 +1318,29 @@ def construct_image_stack_from_movie( _, _, gh, gw = deformation_field.data.shape else: gh = gw = 0 - normalized_t = torch.linspace(0, 1, steps=t, device=movie.device) + if normalized_t_values is None: + normalized_t = torch.linspace(0, 1, steps=t, device=movie.device) + else: + if normalized_t_values.numel() != t: + raise ValueError( + "normalized_t_values must have one entry per movie frame." + ) + normalized_t = normalized_t_values.to(device=movie.device).reshape(t) pixel_grid = coordinate_grid( 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] + particle_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() + particle_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() + pos_y = self._df.loc[particle_indexes, y_col].to_numpy().copy() + pos_x = self._df.loc[particle_indexes, x_col].to_numpy().copy() # 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": @@ -1127,15 +1361,19 @@ 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) + if particle_indices is not None: + frame_shifts = frame_shifts[particle_indices] + elif deformation_field is None: + frame_shifts = torch.zeros( + (num_particles_to_process, 2), + dtype=movie.dtype, + device=movie.device, + ) else: frame_shifts = self.compute_frame_particle_shifts_from_deformation( movie_frame=movie_frame, @@ -1149,9 +1387,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,12 +1412,99 @@ 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() + return aligned_particle_movies_rfft, particle_indexes + + # pylint: disable=too-many-arguments + # pylint: disable=too-many-positional-arguments + def construct_particle_movie_stack( + self, + movie: torch.Tensor, + deformation_field: DeformationField | None = None, + particle_shifts: torch.Tensor | None = None, + pos_reference: Literal["center", "top-left"] = "top-left", + handle_bounds: Literal["pad", "error"] = "pad", + padding_mode: Literal["constant", "reflect", "replicate"] = "constant", + padding_value: float = 0.0, + use_gradient_checkpointing: bool = True, + particle_indices: list[int] | None = None, + normalized_t_values: torch.Tensor | None = None, + ) -> torch.Tensor: + """Construct per-frame particle images from a movie without dose summing. + + If neither ``deformation_field`` nor ``particle_shifts`` is provided, the + movie is assumed to already be aligned and frames are extracted directly. + + Returns + ------- + torch.Tensor + Real-space particle movie stack with shape ``(T, N, H, W)``. + """ + particle_movie_rfft, _ = self._construct_particle_movie_rfft_stack( + movie=movie, + deformation_field=deformation_field, + particle_shifts=particle_shifts, + pos_reference=pos_reference, + handle_bounds=handle_bounds, + padding_mode=padding_mode, + padding_value=padding_value, + use_gradient_checkpointing=use_gradient_checkpointing, + particle_indices=particle_indices, + require_motion_source=False, + normalized_t_values=normalized_t_values, + ) + particle_movie = torch.fft.irfftn( # pylint: disable=not-callable + particle_movie_rfft, + s=self.extracted_box_size, + dim=(-2, -1), + ) + return particle_movie.permute(1, 0, 2, 3).contiguous() + + # pylint: disable=too-many-arguments + # pylint: disable=too-many-positional-arguments + def construct_image_stack_from_movie( + self, + movie: torch.Tensor, + deformation_field: DeformationField | None = None, + particle_shifts: torch.Tensor | None = None, + pos_reference: Literal["center", "top-left"] = "top-left", + handle_bounds: Literal["pad", "error"] = "pad", + padding_mode: Literal["constant", "reflect", "replicate"] = "constant", + padding_value: float = 0.0, + pre_exposure: float = 0.0, + fluence_per_frame: float = 0.0, + use_gradient_checkpointing: bool = True, + particle_indices: list[int] | None = None, + ) -> torch.Tensor: + """Construct a dose-weighted particle image stack from a movie file. + + 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. + """ + pixel_sizes = self.get_pixel_size() + box_h, box_w = self.extracted_box_size + aligned_particle_movies_rfft, particle_indexes = ( + self._construct_particle_movie_rfft_stack( + movie=movie, + deformation_field=deformation_field, + particle_shifts=particle_shifts, + pos_reference=pos_reference, + handle_bounds=handle_bounds, + padding_mode=padding_mode, + padding_value=padding_value, + use_gradient_checkpointing=use_gradient_checkpointing, + particle_indices=particle_indices, + require_motion_source=True, + ) + ) + num_particles_to_process = aligned_particle_movies_rfft.shape[0] + # Dose weight the aligned particle images aligned_particle_images = torch.zeros( (num_particles_to_process, box_h, box_w), @@ -1192,19 +1514,520 @@ def construct_image_stack_from_movie( particle_dft = aligned_particle_movies_rfft[particle_index] # Get the actual dataframe index for this particle - df_idx = paticle_indexes[particle_index] + df_idx = particle_indexes[particle_index] df_loc = self._df.index.get_loc(df_idx) dw_sum = dose_weight_movie_to_micrograph( movie_fft=particle_dft, - pixel_size=pixel_sizes[df_loc], + pixel_size=float(pixel_sizes[df_loc].item()), 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 + + # --------------------------------------------------------------------------- + # I/O methods + # --------------------------------------------------------------------------- + + def export_results(self, allow_file_overwrite: bool = False) -> None: + """Write the particle table to ``df_path`` as CSV. + + Parameters + ---------- + allow_file_overwrite : bool + Whether to overwrite an existing file at ``df_path``. Default is + False. + + Raises + ------ + ValueError + If the parent directory is not writable, or ``df_path`` already + exists and ``allow_file_overwrite`` is False. + """ + _check_output_path(self.df_path, allow_file_overwrite) + self._df.to_csv(self.df_path) + + 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 every entry currently in :attr:`local_stats` to the HDF5 file, by + default False. Raises ``ValueError`` if :attr:`local_stats` is empty. + + 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=self.local_stats if include_local_stats else {}, + skip_df_load=True, + ) + hdf5_stack._df = df # pylint: disable=protected-access + 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 statistic + maps 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``). + + Any subset of the ``*_path`` statistic-map columns can be stored as ``local_stats`` + -- not just correlation average/variance. Populate ``self.local_stats`` (e.g. via + ``self.local_stats.update(self.get_local_stat_maps())``) before calling + ``to_hdf5(include_local_stats=True)``, and every entry present at that point is + written, each to its own dataset under ``/local_stats`` named after its column + (e.g. ``mip_path``, ``correlation_average_path``). + + 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] + (N, valid_h, valid_w) float32 + ... -- one dataset per entry in + `local_stats` at write time, + e.g. `mip_path`, + `correlation_average_path` + + 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. + """ + _check_output_path(self.hdf5_path, self.allow_file_overwrite) + 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. + """ + 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 export_results( + self, + include_image_stack: bool = False, + include_local_stats: bool = False, + ) -> None: + """Write the particle table (and optional tensors) to ``hdf5_path``. + + Alias for ``to_hdf5``, kept for API symmetry with ``ParticleStackCSV``. + """ + self.to_hdf5( + include_image_stack=include_image_stack, + include_local_stats=include_local_stats, + ) + + 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 every entry currently in :attr:`local_stats` to its own + dataset under ``/local_stats``, by default False. + """ + 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 stat maps -- every column currently held in + # `local_stats` is written, whatever it is (correlation average/variance, + # MIP, orientations, defocus, ...). + if include_local_stats: + if not self.local_stats: + raise ValueError( + "local_stats is empty; cannot write to HDF5. Populate it " + "first, e.g. " + "self.local_stats.update(self.get_local_stat_maps())." + ) + local_grp = f.create_group(_HDF5_LOCAL_STATS_GROUP) + for column, stat_map in self.local_stats.items(): + local_grp.create_dataset( + column, + data=stat_map.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")) + # pylint: disable=not-an-iterable + 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"] + ) + # pylint: enable=not-an-iterable + 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_stats: dict[str, torch.Tensor] = {} + 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_stats = { + column: torch.from_numpy(local_grp[column][:]) + for column in local_grp + } + + 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=local_stats, + skip_df_load=True, + ) + instance._df = df + return instance + + +# --------------------------------------------------------------------------- +# Shared result-export helper +# --------------------------------------------------------------------------- + + +def export_particle_stack( + df: pd.DataFrame, + output_path: str, + source_particle_stack: "_ParticleStackBase", + output_format: Literal["csv", "hdf5"] | None = None, + allow_file_overwrite: bool = False, +) -> "ParticleStackCSV | ParticleStackHDF5": + """Wrap a particle-result DataFrame in a ParticleStack and write it to disk. + + Notes + ----- + Used by the refine/optimize/constrained-search managers so that their output back- + end matches the back-end of the input particle stack by default, while still + allowing an explicit override. The DataFrame is only an intermediate — the returned + object is the actual particle stack, reusable directly (e.g. fed into the next + program) without re-reading from disk. + + Parameters + ---------- + df : pd.DataFrame + The particle table to write (e.g. a refined result table). Must be a superset of + the columns a `ParticleStackCSV`/`ParticleStackHDF5` expects; extra columns + (e.g. `refined_*`) are preserved as-is. + output_path : str + Destination file path. + source_particle_stack : _ParticleStackBase + The particle stack `df` was derived from. Supplies the default output format + (matches its own back-end) and the shared box-size/pre-processing metadata to + carry over to the new instance. + output_format : Literal["csv", "hdf5"] | None + Explicit output back-end. If None (default), inferred from + ``type(source_particle_stack)``: ``ParticleStackHDF5`` -> "hdf5", otherwise + "csv". + allow_file_overwrite : bool + Whether to overwrite an existing file at ``output_path``. Default is False. + + Returns + ------- + ParticleStackCSV | ParticleStackHDF5 + The newly constructed particle stack, already written to + ``output_path``. + + Raises + ------ + ValueError + If ``output_format`` is not one of "csv" or "hdf5". + """ + if output_format is None: + output_format = ( + "hdf5" if isinstance(source_particle_stack, ParticleStackHDF5) else "csv" + ) + + spc = source_particle_stack + shared_kwargs: dict[str, Any] = { + "extracted_box_size": spc.extracted_box_size, + "original_template_size": spc.original_template_size, + "leopard_em_version": spc.leopard_em_version, + "global_whitening_applied": spc.global_whitening_applied, + "local_whitening_applied": spc.local_whitening_applied, + "global_normalization_applied": (spc.global_normalization_applied), + "local_normalization_applied": spc.local_normalization_applied, + "skip_df_load": True, + } + + if output_format == "csv": + csv_stack = ParticleStackCSV(df_path=output_path, **shared_kwargs) + csv_stack._df = df # pylint: disable=protected-access + csv_stack.export_results(allow_file_overwrite=allow_file_overwrite) + return csv_stack + + if output_format != "hdf5": + raise ValueError( + f"Unknown output_format '{output_format}'; expected 'csv' or 'hdf5'." + ) + + df_out = df + if df_out.index.name != "particle_id" and "particle_id" not in df_out.columns: + df_out = df_out.copy() + particle_ids = _generate_particle_ids(df_out) + df_out.insert(0, "particle_id", particle_ids) + df_out = df_out.set_index("particle_id") + df_out.index.name = "particle_id" + + hdf5_stack = ParticleStackHDF5( + hdf5_path=output_path, + allow_file_overwrite=allow_file_overwrite, + **shared_kwargs, + ) + hdf5_stack._df = df_out # pylint: disable=protected-access + hdf5_stack.export_results() + return hdf5_stack + + +# --------------------------------------------------------------------------- +# Backward-compatibility alias +# --------------------------------------------------------------------------- + +# Existing code that imports `ParticleStack` continues to receive +# `ParticleStackCSV` unchanged. +ParticleStack = ParticleStackCSV diff --git a/src/leopard_em/pydantic_models/formats.py b/src/leopard_em/pydantic_models/formats.py index 1e84a2e4..3e2fbb89 100644 --- a/src/leopard_em/pydantic_models/formats.py +++ b/src/leopard_em/pydantic_models/formats.py @@ -1,5 +1,16 @@ """Submodule for shared formats used in pydantic models.""" +# Full-micrograph 2DTM result map paths, shared across the column-order +# lists below and by ``_DEFAULT_LOCAL_STAT_COLUMNS`` in particle_stack.py. +STATISTIC_MAP_PATH_COLUMNS = [ + "mip_path", + "scaled_mip_path", + "psi_path", + "theta_path", + "phi_path", + "defocus_path", +] + MATCH_TEMPLATE_DF_COLUMN_ORDER = [ "particle_index", "mip", diff --git a/src/leopard_em/pydantic_models/managers/__init__.py b/src/leopard_em/pydantic_models/managers/__init__.py index 75d8f939..3cc0392a 100644 --- a/src/leopard_em/pydantic_models/managers/__init__.py +++ b/src/leopard_em/pydantic_models/managers/__init__.py @@ -1,13 +1,17 @@ """Pydantic models for Leopard-EM program managers.""" from .constrained_search_manager import ConstrainedSearchManager +from .frame_inspection_manager import FrameInspectionManager from .match_template_manager import MatchTemplateManager from .optimize_template_manager import OptimizeTemplateManager +from .peak_inspection_manager import PeakInspectionManager from .refine_template_manager import RefineTemplateManager __all__ = [ + "ConstrainedSearchManager", + "FrameInspectionManager", "MatchTemplateManager", - "RefineTemplateManager", "OptimizeTemplateManager", - "ConstrainedSearchManager", + "PeakInspectionManager", + "RefineTemplateManager", ] diff --git a/src/leopard_em/pydantic_models/managers/constrained_search_manager.py b/src/leopard_em/pydantic_models/managers/constrained_search_manager.py index cc343a32..cf687058 100644 --- a/src/leopard_em/pydantic_models/managers/constrained_search_manager.py +++ b/src/leopard_em/pydantic_models/managers/constrained_search_manager.py @@ -1,7 +1,8 @@ # pylint: disable=duplicate-code """Pydantic model for running the constrained search program.""" -from typing import Any, ClassVar +import os +from typing import Any, ClassVar, Literal import numpy as np import pandas as pd @@ -18,7 +19,11 @@ PreprocessingFilters, ) from leopard_em.pydantic_models.custom_types import BaseModel2DTM, ExcludedTensor -from leopard_em.pydantic_models.data_structures import ParticleStack +from leopard_em.pydantic_models.data_structures import ( + ParticleStackCSV, + ParticleStackHDF5, + export_particle_stack, +) from leopard_em.pydantic_models.formats import CONSTRAINED_DF_COLUMN_ORDER from leopard_em.utils.backend_setup import ( _setup_correlation_stacks_from_micrographs, @@ -42,9 +47,12 @@ class ConstrainedSearchManager(BaseModel2DTM): Path to the template volume MRC file. center_vector : list[float] The centre vector of the template volume. - particle_stack_reference : ParticleStack - Particle stack object containing particle data reference particles. - particle_stack_constrained : ParticleStack + particle_stack_reference : ParticleStackCSV | ParticleStackHDF5 + Particle stack object containing particle data reference particles. Use + ``ParticleStackCSV`` for a CSV-backed particle table or + ``ParticleStackHDF5`` for an HDF5-backed one. Both expose the same in-memory + API. + particle_stack_constrained : ParticleStackCSV | ParticleStackHDF5 Particle stack object containing particle data constrained particles. defocus_refinement_config : DefocusSearchConfig Configuration for defocus refinement. @@ -68,6 +76,13 @@ class ConstrainedSearchManager(BaseModel2DTM): Create the kwargs for the backend refine_template core function. run_constrained_search(self, orientation_batch_size: int = 64) -> None Run the constrained search program. + refine_result_to_dataframe(self, result) -> pd.DataFrame + Build the refined particle DataFrame from a backend result (no I/O). + export_results(self, output_dataframe_path: str, result, ...) -> None + Build the refined DataFrame and write it (plus CSV parameter/ + above-threshold siblings) to disk, matching the input + particle_stack_reference's back-end by default (override with + `output_format`). """ model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True) @@ -75,8 +90,8 @@ class ConstrainedSearchManager(BaseModel2DTM): template_volume_path: str # In df per-particle, but ensure only one reference center_vector: list[float] = Field(default=[0.0, 0.0, 0.0]) - particle_stack_reference: ParticleStack - particle_stack_constrained: ParticleStack + particle_stack_reference: ParticleStackCSV | ParticleStackHDF5 + particle_stack_constrained: ParticleStackCSV | ParticleStackHDF5 defocus_refinement_config: DefocusSearchConfig orientation_refinement_config: ConstrainedOrientationConfig preprocessing_filters: PreprocessingFilters @@ -209,8 +224,10 @@ def run_constrained_search( output_dataframe_path: str, false_positives: float = 0.005, orientation_batch_size: int = 64, + output_format: Literal["csv", "hdf5"] | None = None, + allow_file_overwrite: bool = False, ) -> None: - """Run the constrained search program and saves the resultant DataFrame to csv. + """Run the constrained search program and export the resultant DataFrame. Parameters ---------- @@ -220,15 +237,26 @@ def run_constrained_search( The number of false positives to allow per particle. orientation_batch_size : int Number of orientations to process at once. Defaults to 64. + output_format : Literal["csv", "hdf5"] | None + Output back-end for the main refined table. Defaults to None, which matches + the back-end of ``self.particle_stack_reference`` (CSV in, CSV out; HDF5 in, + HDF5 out). Pass "csv" or "hdf5" to override. The accompanying "_parameters" + and "_above_threshold" sibling tables are always written as CSV regardless + of this setting. + allow_file_overwrite : bool + Whether to overwrite an existing file at ``output_dataframe_path``. Defaults + to False. """ backend_kwargs = self.make_backend_core_function_kwargs() result = self.get_refine_result(backend_kwargs, orientation_batch_size) - self.refine_result_to_dataframe( + self.export_results( output_dataframe_path=output_dataframe_path, result=result, false_positives=false_positives, + output_format=output_format, + allow_file_overwrite=allow_file_overwrite, ) def get_refine_result( @@ -269,20 +297,20 @@ def get_refine_result( # pylint: disable=too-many-locals def refine_result_to_dataframe( self, - output_dataframe_path: str, result: dict[str, np.ndarray], - false_positives: float = 0.005, - ) -> None: - """Convert refine template result to dataframe. + ) -> pd.DataFrame: + """Convert constrained search result to a DataFrame. Parameters ---------- - output_dataframe_path : str - Path to save the refined particle data. result : dict[str, np.ndarray] - The result of the refine template program. - false_positives : float - The number of false positives to allow per particle. + The result of the constrained search program. + + Returns + ------- + pd.DataFrame + The refined particle data. Not written to disk; use ``export_results`` to do + both in one call. """ df_refined = self.particle_stack_reference.get_dataframe_copy() @@ -338,8 +366,51 @@ def refine_result_to_dataframe( # Reorder the columns df_refined = df_refined.reindex(columns=CONSTRAINED_DF_COLUMN_ORDER) - # Save the refined DataFrame to disk - df_refined.to_csv(output_dataframe_path) + return df_refined + + # pylint: disable=too-many-locals + def export_results( + self, + output_dataframe_path: str, + result: dict[str, np.ndarray], + false_positives: float = 0.005, + output_format: Literal["csv", "hdf5"] | None = None, + allow_file_overwrite: bool = False, + ) -> ParticleStackCSV | ParticleStackHDF5: + """Build the refined DataFrame and write it, plus two CSV siblings, to disk. + + Parameters + ---------- + output_dataframe_path : str + Path to save the refined particle data. + result : dict[str, np.ndarray] + The result of the constrained search program. + false_positives : float + The number of false positives to allow per particle. + output_format : Literal["csv", "hdf5"] | None + Output back-end for the main refined table. Defaults to None, + which matches the back-end of ``self.particle_stack_reference``. + Pass "csv" or "hdf5" to override. + allow_file_overwrite : bool + Whether to overwrite an existing file at ``output_dataframe_path``. + Defaults to False. + + Returns + ------- + ParticleStackCSV | ParticleStackHDF5 + The refined particle stack (main table only), already written to + ``output_dataframe_path``. + """ + df_refined = self.refine_result_to_dataframe(result=result) + + # Save the main refined DataFrame, matching the input back-end by default + result_stack = export_particle_stack( + df=df_refined, + output_path=output_dataframe_path, + source_particle_stack=self.particle_stack_reference, + output_format=output_format, + allow_file_overwrite=allow_file_overwrite, + ) # Save a second dataframe # I also want the original user input offsets back somewhere @@ -362,6 +433,13 @@ def refine_result_to_dataframe( num_correlations, float(false_positives) ) + # The "_parameters" and "_above_threshold" siblings are always CSV, + # named from the base of output_dataframe_path regardless of its + # extension (which may be .h5 when output_format is "hdf5"). + output_base, _ = os.path.splitext(output_dataframe_path) + parameters_path = f"{output_base}_parameters.csv" + above_threshold_path = f"{output_base}_above_threshold.csv" + # Save all parameters to CSV including false-positives params_df = pd.DataFrame( { @@ -372,7 +450,7 @@ def refine_result_to_dataframe( "threshold": [threshold], } ) - params_df.to_csv(output_dataframe_path.replace(".csv", "_parameters.csv")) + params_df.to_csv(parameters_path) print( f"Threshold: {threshold} which gives {false_positives} " @@ -389,10 +467,7 @@ def refine_result_to_dataframe( df_refined_above_threshold["refined_scaled_mip"] != np.nan ] # Save the above threshold dataframe - print( - f"Saving above threshold dataframe to " - f"{output_dataframe_path.replace('.csv', '_above_threshold.csv')}" - ) - df_refined_above_threshold.to_csv( - output_dataframe_path.replace(".csv", "_above_threshold.csv") - ) + print(f"Saving above threshold dataframe to {above_threshold_path}") + df_refined_above_threshold.to_csv(above_threshold_path) + + return result_stack diff --git a/src/leopard_em/pydantic_models/managers/frame_inspection_manager.py b/src/leopard_em/pydantic_models/managers/frame_inspection_manager.py new file mode 100644 index 00000000..0743eebc --- /dev/null +++ b/src/leopard_em/pydantic_models/managers/frame_inspection_manager.py @@ -0,0 +1,875 @@ +"""Pydantic model for per-frame local peak inspection.""" + +# pylint: disable=duplicate-code + +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import torch +import tqdm +from torch_fourier_filter.dose_weight import cumulative_dose_filter_3d +from torch_motion_correction.correct_motion import correct_motion +from torch_motion_correction.deformation_field import ( # pyright: ignore[reportMissingImports] + DeformationField, +) + +from leopard_em.analysis.inspect_peaks_result import save_inspection_result +from leopard_em.pydantic_models.managers.peak_inspection_manager import ( + PeakInspectionManager, +) +from leopard_em.utils.backend_setup import ( + setup_frame_filters_particle_stack, + setup_static_particle_kwargs, +) +from leopard_em.utils.data_io import load_template_tensor, read_particle_shifts_from_csv +from leopard_em.utils.image_processing import get_image_normalization_factor + + +@dataclass(frozen=True) +class FixedFrameFilters: + """Fixed whitening/normalization artifacts shared by every movie frame. + + Attributes + ---------- + image_filters : torch.Tensor + Fixed cumulative Fourier filters applied to each frame's particle images. + projective_filters : torch.Tensor + Fixed projective (template-shaped) Fourier filters. + normalization_factor : torch.Tensor + Fixed per-particle variance normalization factor. + full_image_shape : tuple[int, int] + Shape used as the "full image" reference during normalization. + extracted_box_shape : tuple[int, int] + Shape of the extracted particle boxes. + """ + + image_filters: torch.Tensor + projective_filters: torch.Tensor + normalization_factor: torch.Tensor + full_image_shape: tuple[int, int] + extracted_box_shape: tuple[int, int] + + +@dataclass(frozen=True) +class FrameInspectionContext: + """Immutable inputs threaded through one per-frame inspection run. + + Bundles the movie/motion inputs, the shared template and fixed filters, the frame- + independent backend kwargs, and the scoring settings so the per-frame helpers + compose by passing one object instead of re-deriving or stashing transient state on + the manager. + + Attributes + ---------- + movie : torch.Tensor + Movie tensor with leading frame dimension, on the primary GPU device. + deformation_field : Any | None + Optional deformation-field object used to derive per-frame shifts. + particle_shifts : torch.Tensor | None + Optional explicit per-frame/per-particle shifts (takes precedence over + deformation-derived shifts). + template : torch.Tensor + Template volume shared by every frame (optionally dose-weighted per frame). + fixed_filters : FixedFrameFilters + Fixed whitening/normalization artifacts shared across frames. + frame_independent_kwargs : dict[str, Any] + Backend kwargs that do not vary frame to frame. + correlation_batch_size : int + Number of orientation offsets processed per backend batch. + apply_projection_normalization : bool + Whether to normalize each projection before scoring. + apply_template_dose_weighting : bool + Whether to dose-weight the template separately for each frame interval. + output_mode : Literal["cross_correlation", "frc"] + Score mode used by the inspect backend. + """ + + movie: torch.Tensor + deformation_field: Any | None + particle_shifts: torch.Tensor | None + template: torch.Tensor + fixed_filters: FixedFrameFilters + frame_independent_kwargs: dict[str, Any] + correlation_batch_size: int + apply_projection_normalization: bool + apply_template_dose_weighting: bool + output_mode: Literal["cross_correlation", "frc"] + + +class FrameInspectionManager(PeakInspectionManager): + """Run peak inspection independently for each frame in a movie.""" + + def _prepare_frame_template( + self, + template_tensor: torch.Tensor | None = None, + ) -> torch.Tensor: + """Load a template volume once for all frame-level inspections. + + Parameters + ---------- + template_tensor : torch.Tensor | None, optional + Optional in-memory template volume override. When ``None``, load from + manager-configured template inputs. + + Returns + ------- + torch.Tensor + Template volume on the primary configured GPU device. + """ + if template_tensor is None: + template = load_template_tensor( + template_volume=self.template_volume, + template_volume_path=self.template_volume_path, + ) + else: + template = load_template_tensor(template_volume=template_tensor) + return template.to(self.computational_config.gpu_devices[0]) + + def _setup_frame_independent_kwargs( + self, + template: torch.Tensor, + prefer_refined_angles: bool = True, + ) -> dict[str, Any]: + """Build inspect-backend kwargs that remain constant across frames. + + Parameters + ---------- + template : torch.Tensor + Template volume used to derive static CTF/setup dimensions. + prefer_refined_angles : bool, optional + If ``True``, use refined Euler angles when available in particle metadata. + + Returns + ------- + dict[str, Any] + Backend kwargs shared by all frame calls. + """ + return setup_static_particle_kwargs( + particle_stack=self.particle_stack, + template=template, + euler_angles=self.particle_stack.get_euler_angles(prefer_refined_angles), + euler_angle_offsets=self.orientation_refinement_config.euler_angles_offsets, + defocus_offsets=self.defocus_refinement_config.defocus_values, + pixel_size_offsets=self.pixel_size_refinement_config.pixel_size_values, + device_list=self.computational_config.gpu_devices, + ) + + def _load_and_setup_frame_inspection( + self, + ) -> tuple[torch.Tensor, Any | None, torch.Tensor | None]: + """Load movie input and resolve motion information for frame processing. + + Returns + ------- + tuple[torch.Tensor, Any | None, torch.Tensor | None] + ``(movie, deformation_field, particle_shifts)`` where movie is on the + primary GPU, and either deformation field or particle shifts may be set. + """ + if not self.movie_config.enabled: + raise ValueError("Per-frame peak inspection requires movie_config.enabled.") + if not self.movie_config.movie_path: + raise ValueError( + "Per-frame peak inspection requires movie_config.movie_path." + ) + + device = self.computational_config.gpu_devices[0] + movie = self.movie_config.movie + if movie is None: + raise ValueError("Per-frame peak inspection requires a loaded movie.") + movie = movie.to(device) + + particle_shifts = None + deformation_field = None + if self.movie_config.particle_shifts_path: + particle_shifts = read_particle_shifts_from_csv( + csv_path=self.movie_config.particle_shifts_path, + num_frames=movie.shape[0], + num_particles=self.particle_stack.num_particles, + ).to(device) + elif self.movie_config.deformation_field_path: + deformation_field_tensor = self.movie_config.deformation_field + if deformation_field_tensor is not None: + deformation_field_data = deformation_field_tensor.to(device) + deformation_field = DeformationField( + data=deformation_field_data, + grid_type="catmull_rom", + ) + + return movie, deformation_field, particle_shifts + + def _iter_frame_particle_stacks( + self, + movie: torch.Tensor, + deformation_field: Any | None, + particle_shifts: torch.Tensor | None, + ) -> Iterator[tuple[int, torch.Tensor]]: + """Yield ``(frame_idx, particle_image_stack)`` for each movie frame. + + Parameters + ---------- + movie : torch.Tensor + Movie tensor with leading frame dimension. + deformation_field : Any | None + Optional deformation-field object used to derive per-frame shifts. + particle_shifts : torch.Tensor | None + Optional explicit per-frame/per-particle shifts; takes precedence over + deformation-derived shifts when provided. + + Yields + ------ + tuple[int, torch.Tensor] + The frame index and its particle image stack ``(N, H_box, W_box)``. + """ + num_frames = movie.shape[0] + for frame_idx in range(num_frames): + normalized_t_values = torch.tensor( + [0.0 if num_frames == 1 else frame_idx / (num_frames - 1)], + device=movie.device, + ) + frame_movie = movie[frame_idx : frame_idx + 1] + frame_particle_shifts = None + if particle_shifts is not None: + frame_particle_shifts = particle_shifts[frame_idx : frame_idx + 1] + frame_particle_stack = self.particle_stack.construct_particle_movie_stack( + movie=frame_movie, + deformation_field=deformation_field, + particle_shifts=frame_particle_shifts, + pos_reference="top-left", + handle_bounds="pad", + padding_mode="reflect", + padding_value=0.0, + use_gradient_checkpointing=False, + normalized_t_values=normalized_t_values, + )[0] + yield frame_idx, frame_particle_stack + + def _build_summed_particle_stack_from_movie( + self, + movie: torch.Tensor, + deformation_field: Any | None, + particle_shifts: torch.Tensor | None, + ) -> torch.Tensor: + """Accumulate per-frame particle crops into one summed particle stack. + + Parameters + ---------- + movie : torch.Tensor + Movie tensor with leading frame dimension. + deformation_field : Any | None + Optional deformation-field object used to derive per-frame shifts. + particle_shifts : torch.Tensor | None + Optional explicit per-frame/per-particle shifts; takes precedence over + deformation-derived shifts when provided. + + Returns + ------- + torch.Tensor + Summed particle image stack with shape ``(N, H_box, W_box)``. + """ + summed_particle_images = None + for _, frame_particle_stack in self._iter_frame_particle_stacks( + movie, deformation_field, particle_shifts + ): + if summed_particle_images is None: + summed_particle_images = frame_particle_stack + else: + summed_particle_images = summed_particle_images + frame_particle_stack + + if summed_particle_images is None: + raise ValueError("Movie has no frames to sum for particle stack.") + return summed_particle_images + + def _apply_template_dose_filter( + self, + template: torch.Tensor, + start_exposure: float, + end_exposure: float, + ) -> torch.Tensor: + """Apply cumulative dose weighting over an exposure interval. + + Parameters + ---------- + template : torch.Tensor + Non-dose-weighted template volume. + start_exposure : float + Start of the exposure interval (electrons / Angstrom^2). + end_exposure : float + End of the exposure interval (electrons / Angstrom^2). + + Returns + ------- + torch.Tensor + Dose-weighted template in real space. + """ + pixel_size = float(self.particle_stack.get_pixel_size()[0].item()) + template_rfft = torch.fft.rfftn( # pylint: disable=not-callable + template, dim=(-3, -2, -1) + ) + dose_filter = cumulative_dose_filter_3d( + volume_shape=template.shape, + pixel_size=pixel_size, + start_exposure=start_exposure, + end_exposure=end_exposure, + crit_exposure_bfactor=-1, + rfft=True, + fftshift=False, + device=template.device, + ) + dose_weighted_template_rfft = template_rfft * dose_filter + return torch.fft.irfftn( # pylint: disable=not-callable + dose_weighted_template_rfft, + s=template.shape, + dim=(-3, -2, -1), + ) + + def _setup_frame_kwargs( + self, + frame_particle_stack: torch.Tensor, + template: torch.Tensor, + fixed_filters: FixedFrameFilters, + ) -> dict[str, torch.Tensor]: + """Build backend kwargs that vary per frame (or summed stack). + + Parameters + ---------- + frame_particle_stack : torch.Tensor + Particle image stack for one frame (or summed across frames). + template : torch.Tensor + Frame-specific template (optionally dose-filtered for this frame). + fixed_filters : FixedFrameFilters + Fixed whitening/normalization artifacts reused across every frame. + + Returns + ------- + dict[str, torch.Tensor] + Frame-local backend kwargs containing particle/image/template DFT tensors + and projective filters. + """ + particle_images_dft, template_dft, projective_filters = ( + setup_frame_filters_particle_stack( + particle_stack=self.particle_stack, + preprocessing_filters=self.preprocessing_filters, + template=template, + particle_images=frame_particle_stack.to(template.device), + apply_global_filtering=self.apply_global_filtering, + fixed_image_filters=fixed_filters.image_filters, + fixed_projective_filters=fixed_filters.projective_filters, + fixed_normalization_factor=fixed_filters.normalization_factor, + full_image_shape=fixed_filters.full_image_shape, + extracted_box_shape=fixed_filters.extracted_box_shape, + ) + ) + return { + "particle_stack_dft": particle_images_dft, + "template_dft": template_dft, + "projective_filters": projective_filters, + } + + def _setup_fixed_frame_normalization_filters( # pylint: disable=too-many-locals + self, + movie: torch.Tensor, + deformation_field: Any | None, + particle_shifts: torch.Tensor | None, + template: torch.Tensor, + ) -> FixedFrameFilters: + """Compute fixed image/projective filters and normalization for all frames. + + Parameters + ---------- + movie : torch.Tensor + Input movie tensor with shape ``(T, H, W)``. + deformation_field : Any | None + Optional deformation-field object for motion correction. + particle_shifts : torch.Tensor | None + Optional explicit per-frame particle shifts. + template : torch.Tensor + Template volume used to derive template Fourier output shape. + + Returns + ------- + FixedFrameFilters + Fixed image/projective filters, normalization factor, and the full-image / + extracted-box shapes used by per-frame preprocessing. + """ + box_h, box_w = self.particle_stack.extracted_box_size + template_output_shape = (template.shape[-2], template.shape[-1] // 2 + 1) + has_dataframe = hasattr(self.particle_stack, "_df") + if self.apply_global_filtering and has_dataframe: + # Reference micrograph for global whitening: sum frames along time. + # - No deformation field (aligned movie or shifts-only path): raw sum. + # - Deformation grid: motion-correct each frame like torch_motion_correction + # (correct_motion), then sum corrected frames. + if deformation_field is not None and particle_shifts is None: + pixel_spacing = float( + self.particle_stack.get_pixel_size().mean().item() + ) + corrected_movie = correct_motion( + image=movie, + deformation_field=deformation_field, + pixel_spacing=pixel_spacing, + device=movie.device, + ) + summed_movie_image = corrected_movie.sum(dim=0, keepdim=True) + else: + summed_movie_image = movie.sum(dim=0, keepdim=True) + summed_movie_dft = torch.fft.rfftn( # pylint: disable=not-callable + summed_movie_image, dim=(-2, -1) + ) + summed_movie_dft[..., 0, 0] = 0.0 + 0.0j + df_index = self.particle_stack.get_dataframe_copy().index + fixed_projective_filters = self.particle_stack.construct_projective_filters( + self.preprocessing_filters, + output_shape=template_output_shape, + images_dft=summed_movie_dft.detach(), + indices=[df_index], + ).to(template.device) + summed_particle_images = self.particle_stack.construct_image_stack( + images=summed_movie_image, + indices=[df_index], + extraction_size=self.particle_stack.extracted_box_size, + pos_reference="top-left", + handle_bounds="pad", + padding_mode="reflect", + padding_value=0.0, + ) + summed_particle_images_dft = torch.fft.rfftn( # pylint: disable=not-callable + summed_particle_images, dim=(-2, -1) + ) + summed_particle_images_dft[..., 0, 0] = 0.0 + 0.0j + particle_rfft_shape = summed_particle_images_dft.shape[-2:] + fixed_image_filters = self.particle_stack.construct_projective_filters( + self.preprocessing_filters, + output_shape=particle_rfft_shape, + images_dft=summed_movie_dft.detach(), + indices=[df_index], + ).to(template.device) + bandpass_filter = ( + self.preprocessing_filters.bandpass_filter.calculate_bandpass_filter( + summed_particle_images_dft.shape[-2:] + ).to(template.device) + ) + normalization_factor = get_image_normalization_factor( + image_rfft=summed_particle_images_dft, + cumulative_fourier_filters=fixed_image_filters, + bandpass_filter=bandpass_filter, + full_image_shape=(box_h, box_w), + extracted_box_shape=(box_h, box_w), + ) + return FixedFrameFilters( + image_filters=fixed_image_filters, + projective_filters=fixed_projective_filters, + normalization_factor=normalization_factor, + full_image_shape=(box_h, box_w), + extracted_box_shape=(box_h, box_w), + ) + + summed_particle_images = self._build_summed_particle_stack_from_movie( + movie=movie, + deformation_field=deformation_field, + particle_shifts=particle_shifts, + ) + summed_particle_images_dft = torch.fft.rfftn( # pylint: disable=not-callable + summed_particle_images, dim=(-2, -1) + ) + summed_particle_images_dft[..., 0, 0] = 0.0 + 0.0j + particle_rfft_shape = summed_particle_images_dft.shape[-2:] + fixed_image_filters = self.particle_stack.construct_image_filters( + self.preprocessing_filters, + output_shape=particle_rfft_shape, + images_dft=summed_particle_images_dft.detach(), + ).to(template.device) + fixed_projective_filters = self.particle_stack.construct_image_filters( + self.preprocessing_filters, + output_shape=template_output_shape, + images_dft=summed_particle_images_dft.detach(), + ).to(template.device) + bandpass_filter = ( + self.preprocessing_filters.bandpass_filter.calculate_bandpass_filter( + summed_particle_images_dft.shape[-2:] + ).to(template.device) + ) + normalization_factor = get_image_normalization_factor( + image_rfft=summed_particle_images_dft, + cumulative_fourier_filters=fixed_image_filters, + bandpass_filter=bandpass_filter, + full_image_shape=(box_h, box_w), + extracted_box_shape=(box_h, box_w), + ) + return FixedFrameFilters( + image_filters=fixed_image_filters, + projective_filters=fixed_projective_filters, + normalization_factor=normalization_factor, + full_image_shape=(box_h, box_w), + extracted_box_shape=(box_h, box_w), + ) + + @staticmethod + def _stack_frame_results( + frame_results: list[torch.Tensor | tuple[torch.Tensor, torch.Tensor]], + output_mode: Literal["cross_correlation", "frc"], + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Stack per-frame backend outputs with frame as leading axis. + + Parameters + ---------- + frame_results : list[torch.Tensor | tuple[torch.Tensor, torch.Tensor]] + Per-frame outputs returned by the inspect backend. + output_mode : Literal["cross_correlation", "frc"] + Determines expected type and final stacked return format. + + Returns + ------- + torch.Tensor | tuple[torch.Tensor, torch.Tensor] + Cross-correlation mode returns ``(T, N, ...)`` tensor; FRC mode returns + ``(stacked_frc, frequency_bins)``. + """ + if not frame_results: + raise ValueError("No frame results were generated.") + if output_mode == "cross_correlation": + cc_results = [] + for result in frame_results: + if not isinstance(result, torch.Tensor): + raise TypeError( + "Expected tensor results for cross-correlation mode." + ) + cc_results.append(result) + return torch.stack(cc_results) + + frc_results = [] + frequency_bins = None + for result in frame_results: + if isinstance(result, torch.Tensor): + raise TypeError("Expected tuple results for FRC mode.") + frc_tensor, current_frequency_bins = result + frc_results.append(frc_tensor) + if frequency_bins is None: + frequency_bins = current_frequency_bins + if frequency_bins is None: + raise ValueError("No FRC frequency bins were generated.") + return torch.stack(frc_results), frequency_bins + + def _frame_dose_template( + self, + template: torch.Tensor, + frame_idx: int, + apply_template_dose_weighting: bool, + ) -> torch.Tensor: + """Return the template to use for a single frame (cumulative dose filter). + + Parameters + ---------- + template : torch.Tensor + Shared non-dose-weighted template volume. + frame_idx : int + Index of the frame being processed. + apply_template_dose_weighting : bool + Whether to dose-weight the template for this frame's exposure interval. + + Returns + ------- + torch.Tensor + Frame-specific template. + """ + if not apply_template_dose_weighting: + return template + start_exposure = self.movie_config.pre_exposure + ( + frame_idx * self.movie_config.fluence_per_frame + ) + end_exposure = start_exposure + self.movie_config.fluence_per_frame + return self._apply_template_dose_filter(template, start_exposure, end_exposure) + + def _inspect_particle_stack( + self, + ctx: FrameInspectionContext, + frame_particle_stack: torch.Tensor, + template: torch.Tensor, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run the inspect backend on one particle stack (single frame or summed). + + Parameters + ---------- + ctx : FrameInspectionContext + Shared inputs and settings for this inspection run. + frame_particle_stack : torch.Tensor + Particle image stack to score. + template : torch.Tensor + Template to score against (optionally dose-weighted). + + Returns + ------- + torch.Tensor | tuple[torch.Tensor, torch.Tensor] + Inspect backend output (CC tensor or FRC tuple). + """ + frame_kwargs = self._setup_frame_kwargs( + frame_particle_stack=frame_particle_stack, + template=template, + fixed_filters=ctx.fixed_filters, + ) + backend_kwargs = {**ctx.frame_independent_kwargs, **frame_kwargs} + return self.get_peak_inspection_result( + backend_kwargs=backend_kwargs, + correlation_batch_size=ctx.correlation_batch_size, + apply_projection_normalization=ctx.apply_projection_normalization, + output_mode=ctx.output_mode, + ) + + def _iter_frame_inspection_results( + self, + ctx: FrameInspectionContext, + ) -> Iterator[tuple[int, torch.Tensor | tuple[torch.Tensor, torch.Tensor]]]: + """Yield ``(frame_idx, inspect_result)`` for each movie frame. + + Parameters + ---------- + ctx : FrameInspectionContext + Shared inputs and settings for this inspection run. + + Yields + ------ + tuple[int, torch.Tensor | tuple[torch.Tensor, torch.Tensor]] + The frame index and its inspect backend output. + """ + num_frames = ctx.movie.shape[0] + frame_iter = tqdm.tqdm( + self._iter_frame_particle_stacks( + ctx.movie, ctx.deformation_field, ctx.particle_shifts + ), + total=num_frames, + desc="Inspecting frames", + unit="frame", + dynamic_ncols=True, + ) + for frame_idx, frame_particle_stack in frame_iter: + frame_template = self._frame_dose_template( + ctx.template, frame_idx, ctx.apply_template_dose_weighting + ) + yield ( + frame_idx, + self._inspect_particle_stack( + ctx, + frame_particle_stack=frame_particle_stack, + template=frame_template, + ), + ) + + def _build_inspection_context( + self, + correlation_batch_size: int, + prefer_refined_angles: bool, + apply_projection_normalization: bool, + template_tensor: torch.Tensor | None, + output_mode: Literal["cross_correlation", "frc"], + apply_template_dose_weighting: bool, + ) -> FrameInspectionContext: + """Load inputs and assemble the immutable per-frame inspection context. + + Loads the movie/motion inputs, prepares the shared template and fixed whitening + filters, and builds the frame-independent backend kwargs, bundling them into a + :class:`FrameInspectionContext`. + + Parameters + ---------- + correlation_batch_size : int + Number of orientation offsets processed per backend batch. + prefer_refined_angles : bool + If True, use refined Euler angles from the particle stack when available. + apply_projection_normalization : bool + Whether to normalize each projection before scoring. + template_tensor : torch.Tensor | None + Optional template volume override. + output_mode : Literal["cross_correlation", "frc"] + Score mode (CC maps or FRC spectra). + apply_template_dose_weighting : bool + If True, apply cumulative dose filtering to the template per frame interval. + + Returns + ------- + FrameInspectionContext + Immutable inputs threaded through the per-frame inspection run. + """ + movie, deformation_field, particle_shifts = ( + self._load_and_setup_frame_inspection() + ) + template = self._prepare_frame_template(template_tensor=template_tensor) + fixed_filters = self._setup_fixed_frame_normalization_filters( + movie=movie, + deformation_field=deformation_field, + particle_shifts=particle_shifts, + template=template, + ) + frame_independent_kwargs = self._setup_frame_independent_kwargs( + template=template, + prefer_refined_angles=prefer_refined_angles, + ) + return FrameInspectionContext( + movie=movie, + deformation_field=deformation_field, + particle_shifts=particle_shifts, + template=template, + fixed_filters=fixed_filters, + frame_independent_kwargs=frame_independent_kwargs, + correlation_batch_size=correlation_batch_size, + apply_projection_normalization=apply_projection_normalization, + apply_template_dose_weighting=apply_template_dose_weighting, + output_mode=output_mode, + ) + + def _collect_stacked_results( + self, + ctx: FrameInspectionContext, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Score every frame and stack the results with frame as the leading axis. + + Parameters + ---------- + ctx : FrameInspectionContext + Shared inputs and settings for this inspection run. + + Returns + ------- + torch.Tensor | tuple[torch.Tensor, torch.Tensor] + Stacked per-frame results: a ``(T, N, n_px, n_def, n_orient, H, W)`` CC + tensor, or ``(stacked_frc, frequency_bins)`` in FRC mode where + ``stacked_frc`` has shape ``(T, N, n_px, n_def, n_orient, n_freq)``. + """ + frame_results = [ + result for _, result in self._iter_frame_inspection_results(ctx) + ] + return self._stack_frame_results(frame_results, ctx.output_mode) + + def run_peak_inspection_per_frame( + self, + correlation_batch_size: int = 32, + prefer_refined_angles: bool = True, + apply_projection_normalization: bool = True, + template_tensor: torch.Tensor | None = None, + output_mode: Literal["cross_correlation", "frc"] = "cross_correlation", + apply_template_dose_weighting: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run peak inspection independently for every movie frame. + + Composes the per-frame pipeline: load the movie/motion inputs, prepare the + shared template and fixed whitening filters, build the frame-independent + backend kwargs, then score and stack every frame. + + Parameters + ---------- + correlation_batch_size : int, optional + Number of orientation offsets processed per backend batch. + prefer_refined_angles : bool, optional + If True, use refined Euler angles from the particle stack when available. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + template_tensor : torch.Tensor | None, optional + Optional template volume override. + output_mode : Literal["cross_correlation", "frc"], optional + Score mode (CC maps or FRC spectra). + apply_template_dose_weighting : bool, optional + If True, apply cumulative dose filtering to the provided non-dose- + weighted template separately for each frame interval. + + Returns + ------- + torch.Tensor | tuple[torch.Tensor, torch.Tensor] + Stacked per-frame results: a ``(T, N, n_px, n_def, n_orient, H, W)`` CC + tensor, or ``(stacked_frc, frequency_bins)`` in FRC mode. + """ + ctx = self._build_inspection_context( + correlation_batch_size=correlation_batch_size, + prefer_refined_angles=prefer_refined_angles, + apply_projection_normalization=apply_projection_normalization, + template_tensor=template_tensor, + output_mode=output_mode, + apply_template_dose_weighting=apply_template_dose_weighting, + ) + return self._collect_stacked_results(ctx) + + def run_and_save_peak_inspection_per_frame( + self, + output_path: str | Path, + correlation_batch_size: int = 32, + prefer_refined_angles: bool = True, + apply_projection_normalization: bool = True, + template_tensor: torch.Tensor | None = None, + output_mode: Literal["cross_correlation", "frc"] = "cross_correlation", + apply_template_dose_weighting: bool = False, + ) -> Path: + """Run per-frame peak inspection and write a self-describing ``.npz`` file. + + Parameters + ---------- + output_path : str | Path + Destination path for the ``.npz`` file (suffix appended if missing). + correlation_batch_size : int, optional + Number of orientation offsets processed per backend batch. + prefer_refined_angles : bool, optional + If True, use refined Euler angles from the particle stack when available. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + template_tensor : torch.Tensor | None, optional + Optional template volume override. + output_mode : Literal["cross_correlation", "frc"], optional + Score mode (CC maps or FRC spectra). + apply_template_dose_weighting : bool, optional + If True, apply cumulative dose filtering to the provided non-dose- + weighted template separately for each frame interval. + + Returns + ------- + Path + The path the result was written to (with ``.npz`` suffix). + """ + ctx = self._build_inspection_context( + correlation_batch_size=correlation_batch_size, + prefer_refined_angles=prefer_refined_angles, + apply_projection_normalization=apply_projection_normalization, + template_tensor=template_tensor, + output_mode=output_mode, + apply_template_dose_weighting=apply_template_dose_weighting, + ) + stacked_results = self._collect_stacked_results(ctx) + + # Move the frame axis behind the particle axis: (T, N, ...) -> (N, T, ...). + num_frames = ctx.movie.shape[0] + if output_mode == "frc": + stacked_frc, frequency_bins = stacked_results + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor] = ( + stacked_frc.transpose(0, 1), + frequency_bins, + ) + else: + assert isinstance(stacked_results, torch.Tensor) + result = stacked_results.transpose(0, 1) + + kwargs = ctx.frame_independent_kwargs + base_defocus = torch.stack( + [kwargs["defocus_u"], kwargs["defocus_v"], kwargs["defocus_angle"]], + dim=-1, + ) + df = self.particle_stack._df # pylint: disable=protected-access + particle_index = ( + df["particle_index"].to_numpy() if "particle_index" in df.columns else None + ) + + return save_inspection_result( + output_path, + result=result, + output_mode=output_mode, + euler_angle_offsets=kwargs["euler_angle_offsets"], + defocus_offsets=kwargs["defocus_offsets"], + pixel_size_offsets=kwargs["pixel_size_offsets"], + base_euler_angles=kwargs["euler_angles"], + base_defocus=base_defocus, + particle_index=particle_index, + frame_index=torch.arange(num_frames), + per_frame=True, + extra_metadata={ + "prefer_refined_angles": prefer_refined_angles, + "apply_projection_normalization": apply_projection_normalization, + "apply_template_dose_weighting": apply_template_dose_weighting, + "correlation_batch_size": correlation_batch_size, + "movie_path": self.movie_config.movie_path, + }, + ) 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..276e2d64 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 @@ -165,10 +170,14 @@ def make_backend_core_function_kwargs(self) -> dict[str, Any]: bandpass_filter = bp_config.calculate_bandpass_filter(image_dft.shape) # Calculate the cumulative filters for both the image and the template. + # NOTE: We don't want to do random fourier masking on the image, so skip the + # dropout mask for the image-side filter (without mutating the config). cumulative_filter_image = self.preprocessing_filters.get_combined_filter( ref_img_rfft=image_dft, output_shape=image_dft.shape, + apply_random_dropout=False, ) + # NOTE: Here, manually accounting for the RFFT in output shape since we have not # RFFT'd the template volume yet. Also, this is 2-dimensional, not 3-dimensional cumulative_filter_template = self.preprocessing_filters.get_combined_filter( @@ -221,7 +230,8 @@ 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, + compute_correlation_table: bool = True, ) -> None: """Runs the base match template in pytorch. @@ -233,7 +243,14 @@ 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. + compute_correlation_table : bool + If True, track cross-correlation values which surpass the correlation + table threshold during the search. If False, the `CorrelationTable` will be + empty. Default is True. Returns ------- @@ -245,11 +262,14 @@ def run_match_template( orientation_batch_size=orientation_batch_size, num_cuda_streams=self.computational_config.num_cpus, backend=self.computational_config.backend, + compute_correlation_table=compute_correlation_table, ) # 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 +281,8 @@ 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, + compute_correlation_table: bool = True, ) -> None: """Runs the base match template in a distributed, multi-node environment. @@ -279,7 +300,14 @@ 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. + compute_correlation_table : bool + If True, track cross-correlation values which surpass the correlation + table threshold during the search. If False, the `CorrelationTable` will be + empty. Default is True. Raises ------ @@ -313,6 +341,7 @@ def run_match_template_distributed( orientation_batch_size, self.computational_config.num_cpus, self.computational_config.backend, + compute_correlation_table=compute_correlation_table, **core_kwargs, ) @@ -320,6 +349,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 +358,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 +381,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 +499,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/managers/optimize_template_manager.py b/src/leopard_em/pydantic_models/managers/optimize_template_manager.py index a4d6c200..8616737b 100644 --- a/src/leopard_em/pydantic_models/managers/optimize_template_manager.py +++ b/src/leopard_em/pydantic_models/managers/optimize_template_manager.py @@ -15,7 +15,10 @@ PreprocessingFilters, ) from leopard_em.pydantic_models.custom_types import BaseModel2DTM, ExcludedTensor -from leopard_em.pydantic_models.data_structures import ParticleStack +from leopard_em.pydantic_models.data_structures import ( + ParticleStackCSV, + ParticleStackHDF5, +) from leopard_em.pydantic_models.formats import REFINED_DF_COLUMN_ORDER from leopard_em.utils.backend_setup import setup_particle_backend_kwargs @@ -25,8 +28,10 @@ class OptimizeTemplateManager(BaseModel2DTM): Attributes ---------- - particle_stack : ParticleStack - Particle stack object containing particle data. + particle_stack : ParticleStackCSV | ParticleStackHDF5 + Particle stack object containing particle data. Use ``ParticleStackCSV`` + for a CSV-backed particle table or ``ParticleStackHDF5`` for an HDF5-backed one. + Both expose the same in-memory API. pixel_size_coarse_search : PixelSizeSearchConfig Configuration for pixel size coarse search. pixel_size_fine_search : PixelSizeSearchConfig @@ -55,7 +60,7 @@ class OptimizeTemplateManager(BaseModel2DTM): model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True) - particle_stack: ParticleStack + particle_stack: ParticleStackCSV | ParticleStackHDF5 pixel_size_coarse_search: PixelSizeSearchConfig pixel_size_fine_search: PixelSizeSearchConfig preprocessing_filters: PreprocessingFilters @@ -446,7 +451,11 @@ def refine_result_to_dataframe( # pylint: disable=too-many-locals result: dict[str, np.ndarray], prefer_refined_angles: bool = True, ) -> None: - """Convert refine template result to dataframe. + """Convert refine template result to a dataframe and write it to CSV. + + NOTE: This always writes CSV, regardless of the input particle_stack's + back-end. It is only used to dump intermediate, per-pixel-size + diagnostic results during the pixel size search. Parameters ---------- diff --git a/src/leopard_em/pydantic_models/managers/peak_inspection_manager.py b/src/leopard_em/pydantic_models/managers/peak_inspection_manager.py new file mode 100644 index 00000000..cc46414c --- /dev/null +++ b/src/leopard_em/pydantic_models/managers/peak_inspection_manager.py @@ -0,0 +1,177 @@ +"""Pydantic model for running local peak inspection.""" + +from pathlib import Path +from typing import Any, Literal + +import torch + +from leopard_em.analysis.inspect_peaks import core_inspect_template +from leopard_em.analysis.inspect_peaks_result import save_inspection_result +from leopard_em.pydantic_models.managers.refine_template_manager import ( + RefineTemplateManager, +) + + +class PeakInspectionManager(RefineTemplateManager): + """Run refine-template search without best-peak reduction. + + This manager reuses the refine-template backend setup, but returns full local + score tensors for inspection rather than only the argmax result. + """ + + def get_peak_inspection_result( + self, + backend_kwargs: dict[str, Any], + correlation_batch_size: int = 32, + apply_projection_normalization: bool = True, + output_mode: Literal["cross_correlation", "frc"] = "cross_correlation", + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run the inspect backend and return scores for all local hypotheses. + + Parameters + ---------- + backend_kwargs : dict[str, Any] + Backend inputs from :meth:`make_backend_core_function_kwargs`. + correlation_batch_size : int, optional + Number of orientation offsets processed per backend batch. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + output_mode : Literal["cross_correlation", "frc"], optional + Score mode. ``"cross_correlation"`` returns local CC maps; ``"frc"`` + returns local FRC spectra. + + Returns + ------- + torch.Tensor | tuple[torch.Tensor, torch.Tensor] + - ``"cross_correlation"``: tensor with shape + ``(N, n_px, n_defocus, n_orient, H, W)``. + - ``"frc"``: ``(frc_tensor, frequency_bins)`` where + ``frc_tensor`` has shape ``(N, n_px, n_defocus, n_orient, n_freq)`` + and ``frequency_bins`` has shape ``(n_freq,)``. + """ + return core_inspect_template( + batch_size=correlation_batch_size, + num_cuda_streams=self.computational_config.num_cpus, + apply_projection_normalization=apply_projection_normalization, + output_mode=output_mode, + **backend_kwargs, + ) + + def run_peak_inspection( + self, + correlation_batch_size: int = 32, + prefer_refined_angles: bool = True, + apply_projection_normalization: bool = True, + template_tensor: torch.Tensor | None = None, + output_mode: Literal["cross_correlation", "frc"] = "cross_correlation", + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run peak inspection using configured data and optional template override. + + Parameters + ---------- + correlation_batch_size : int, optional + Number of orientation offsets processed per backend batch. + prefer_refined_angles : bool, optional + If True, use refined Euler angles from the particle stack when available. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + template_tensor : torch.Tensor | None, optional + Optional template volume override. + output_mode : Literal["cross_correlation", "frc"], optional + Score mode. ``"cross_correlation"`` returns local CC maps; ``"frc"`` + returns local FRC spectra. + + Returns + ------- + torch.Tensor | tuple[torch.Tensor, torch.Tensor] + Inspect output tensor (CC mode) or ``(frc_tensor, frequency_bins)`` + tuple (FRC mode). + """ + backend_kwargs = self.make_backend_core_function_kwargs( + prefer_refined_angles=prefer_refined_angles, + template_tensor=template_tensor, + ) + return self.get_peak_inspection_result( + backend_kwargs=backend_kwargs, + correlation_batch_size=correlation_batch_size, + apply_projection_normalization=apply_projection_normalization, + output_mode=output_mode, + ) + + def run_and_save_peak_inspection( + self, + output_path: str | Path, + correlation_batch_size: int = 32, + prefer_refined_angles: bool = True, + apply_projection_normalization: bool = True, + template_tensor: torch.Tensor | None = None, + output_mode: Literal["cross_correlation", "frc"] = "cross_correlation", + ) -> Path: + """Run peak inspection and write the score tensor to a ``.npz`` file. + + Parameters + ---------- + output_path : str | Path + Destination path for the ``.npz`` file (suffix appended if missing). + correlation_batch_size : int, optional + Number of orientation offsets processed per backend batch. + prefer_refined_angles : bool, optional + If True, use refined Euler angles from the particle stack when available. + apply_projection_normalization : bool, optional + Whether to normalize each projection before scoring. + template_tensor : torch.Tensor | None, optional + Optional template volume override. + output_mode : Literal["cross_correlation", "frc"], optional + Score mode. ``"cross_correlation"`` saves local CC maps; ``"frc"`` + saves local FRC spectra plus the frequency bins. + + Returns + ------- + Path + The path the result was written to (with ``.npz`` suffix). + """ + backend_kwargs = self.make_backend_core_function_kwargs( + prefer_refined_angles=prefer_refined_angles, + template_tensor=template_tensor, + ) + result = self.get_peak_inspection_result( + backend_kwargs=backend_kwargs, + correlation_batch_size=correlation_batch_size, + apply_projection_normalization=apply_projection_normalization, + output_mode=output_mode, + ) + + # Pull the particle ordering from the source dataframe when available so + # tensor rows can be mapped back to the original particle stack. + df = self.particle_stack._df # pylint: disable=protected-access + particle_index = ( + df["particle_index"].to_numpy() if "particle_index" in df.columns else None + ) + + # Per-particle base astigmatic defocus (defocus_u, defocus_v, defocus_angle) so + # absolute defocus can be reconstructed = base_defocus + defocus_offsets. + base_defocus = torch.stack( + [ + backend_kwargs["defocus_u"], + backend_kwargs["defocus_v"], + backend_kwargs["defocus_angle"], + ], + dim=-1, + ) + + return save_inspection_result( + output_path, + result=result, + output_mode=output_mode, + euler_angle_offsets=backend_kwargs["euler_angle_offsets"], + defocus_offsets=backend_kwargs["defocus_offsets"], + pixel_size_offsets=backend_kwargs["pixel_size_offsets"], + base_euler_angles=backend_kwargs["euler_angles"], + base_defocus=base_defocus, + particle_index=particle_index, + extra_metadata={ + "prefer_refined_angles": prefer_refined_angles, + "apply_projection_normalization": apply_projection_normalization, + "correlation_batch_size": correlation_batch_size, + }, + ) diff --git a/src/leopard_em/pydantic_models/managers/refine_template_manager.py b/src/leopard_em/pydantic_models/managers/refine_template_manager.py index b7f9061f..ae75fcaf 100644 --- a/src/leopard_em/pydantic_models/managers/refine_template_manager.py +++ b/src/leopard_em/pydantic_models/managers/refine_template_manager.py @@ -1,12 +1,12 @@ """Pydantic model for running the refine template program.""" -from typing import Any, ClassVar +from typing import Any, ClassVar, Literal import numpy as np import pandas as pd import torch from pydantic import ConfigDict -from torch_cubic_spline_grids import CubicCatmullRomGrid3d +from torch_motion_correction.deformation_field import DeformationField from leopard_em.backend.core_differentiable_refine import core_differentiable_refine from leopard_em.backend.core_refine_template import core_refine_template @@ -19,7 +19,11 @@ RefineOrientationConfig, ) from leopard_em.pydantic_models.custom_types import BaseModel2DTM, ExcludedTensor -from leopard_em.pydantic_models.data_structures import ParticleStack +from leopard_em.pydantic_models.data_structures import ( + ParticleStackCSV, + ParticleStackHDF5, + export_particle_stack, +) from leopard_em.pydantic_models.formats import REFINED_DF_COLUMN_ORDER from leopard_em.utils.backend_setup import setup_particle_backend_kwargs from leopard_em.utils.data_io import ( @@ -36,8 +40,10 @@ class RefineTemplateManager(BaseModel2DTM): ---------- template_volume_path : str Path to the template volume MRC file. - particle_stack : ParticleStack - Particle stack object containing particle data. + particle_stack : ParticleStackCSV | ParticleStackHDF5 + Particle stack object containing particle data. Use ``ParticleStackCSV`` for a + CSV-backed particle table or ``ParticleStackHDF5`` for an HDF5-backed one. Both + expose the same in-memory API. defocus_refinement_config : DefocusSearchConfig Configuration for defocus refinement. pixel_size_refinement_config : PixelSizeSearchConfig @@ -66,12 +72,18 @@ class RefineTemplateManager(BaseModel2DTM): Create the kwargs for the backend refine_template core function. run_refine_template(self, correlation_batch_size: int = 32) -> None Run the refine template program. + refine_result_to_dataframe(self, result, prefer_refined_angles: bool = True) + -> pd.DataFrame + Build the refined particle DataFrame from a backend result (no I/O). + export_results(self, output_dataframe_path: str, result, ...) -> None + Build the refined DataFrame and write it to disk, matching the input + particle_stack's back-end by default (override with `output_format`). """ model_config: ClassVar = ConfigDict(arbitrary_types_allowed=True) template_volume_path: str # In df per-particle, but ensure only one reference - particle_stack: ParticleStack + particle_stack: ParticleStackCSV | ParticleStackHDF5 defocus_refinement_config: DefocusSearchConfig pixel_size_refinement_config: PixelSizeSearchConfig orientation_refinement_config: RefineOrientationConfig @@ -91,7 +103,9 @@ def __init__(self, skip_mrc_preloads: bool = False, **data: Any): self.template_volume = load_mrc_volume(self.template_volume_path) def make_backend_core_function_kwargs( - self, prefer_refined_angles: bool = True + self, + prefer_refined_angles: bool = True, + template_tensor: torch.Tensor | None = None, ) -> dict[str, Any]: """Create the kwargs for the backend refine_template core function. @@ -100,12 +114,18 @@ def make_backend_core_function_kwargs( prefer_refined_angles : bool Whether to use the refined angles from the particle stack. Defaults to True. + template_tensor : torch.Tensor | None + Optional template volume override. If None, the configured template + volume/path is used. """ # Ensure the template is loaded in as a Tensor object - template = load_template_tensor( - template_volume=self.template_volume, - template_volume_path=self.template_volume_path, - ) + if template_tensor is None: + template = load_template_tensor( + template_volume=self.template_volume, + template_volume_path=self.template_volume_path, + ) + else: + template = load_template_tensor(template_volume=template_tensor) # The set of "best" euler angles from match template search # Check if refined angles exist, otherwise use the original angles @@ -144,8 +164,9 @@ def make_backend_core_function_kwargs( # Use deformation field if particle shifts not provided deformation_field_tensor = self.movie_config.deformation_field if deformation_field_tensor is not None: - deformation_field = CubicCatmullRomGrid3d.from_grid_data( - deformation_field_tensor + deformation_field = DeformationField( + data=deformation_field_tensor, + grid_type="catmull_rom", ) # Use the common utility function to set up the backend kwargs @@ -241,9 +262,13 @@ def make_differentiable_backend_kwargs( ) def run_refine_template( - self, output_dataframe_path: str, correlation_batch_size: int = 32 + self, + output_dataframe_path: str, + correlation_batch_size: int = 32, + output_format: Literal["csv", "hdf5"] | None = None, + allow_file_overwrite: bool = False, ) -> None: - """Run the refine template program and saves the resultant DataFrame to csv. + """Run the refine template program and export the resultant DataFrame. Parameters ---------- @@ -251,15 +276,26 @@ def run_refine_template( Path to save the refined particle data. correlation_batch_size : int Number of cross-correlations to process in one batch, defaults to 32. + output_format : Literal["csv", "hdf5"] | None + Output back-end to write. Defaults to None, which matches the back-end of + ``self.particle_stack`` (CSV in, CSV out; HDF5 in, HDF5 out). Pass "csv" or + "hdf5" to override. + allow_file_overwrite : bool + Whether to overwrite an existing file at ``output_dataframe_path``. Defaults + to False. """ backend_kwargs = self.make_backend_core_function_kwargs() result = self.get_refine_result(backend_kwargs, correlation_batch_size) - self.refine_result_to_dataframe( - output_dataframe_path=output_dataframe_path, result=result + self.export_results( + output_dataframe_path=output_dataframe_path, + result=result, + output_format=output_format, + allow_file_overwrite=allow_file_overwrite, ) + # pylint: disable=too-many-positional-arguments,too-many-arguments def run_differentiable_refine( self, output_dataframe_path: str, @@ -270,8 +306,10 @@ def run_differentiable_refine( template_tensor: torch.Tensor | None = None, correlation_batch_size: int = 32, images_are_particles: bool = False, + output_format: Literal["csv", "hdf5"] | None = None, + allow_file_overwrite: bool = False, ) -> None: - """Run the differentiable refine template program and saves DataFrame to csv. + """Run the differentiable refine template program and export the DataFrame. Parameters ---------- @@ -292,6 +330,13 @@ def run_differentiable_refine( Number of cross-correlations to process in one batch, defaults to 32. images_are_particles : bool Whether the images are particles or not. Defaults to False. + output_format : Literal["csv", "hdf5"] | None + Output back-end to write. Defaults to None, which matches the back-end of + ``self.particle_stack`` (CSV in, CSV out; HDF5 in, HDF5 out). Pass "csv" or + "hdf5" to override. + allow_file_overwrite : bool + Whether to overwrite an existing file at ``output_dataframe_path``. Defaults + to False. """ backend_kwargs = self.make_differentiable_backend_kwargs( @@ -307,8 +352,11 @@ def run_differentiable_refine( backend_kwargs, correlation_batch_size, use_differentiable=True ) - self.refine_result_to_dataframe( - output_dataframe_path=output_dataframe_path, result=result + self.export_results( + output_dataframe_path=output_dataframe_path, + result=result, + output_format=output_format, + allow_file_overwrite=allow_file_overwrite, ) def get_refine_result( @@ -355,21 +403,24 @@ def get_refine_result( # pylint: disable=too-many-locals def refine_result_to_dataframe( self, - output_dataframe_path: str, result: dict[str, np.ndarray | torch.Tensor], prefer_refined_angles: bool = True, - ) -> None: - """Convert refine template result to dataframe. + ) -> pd.DataFrame: + """Convert refine template result to a DataFrame. Parameters ---------- - output_dataframe_path : str - Path to save the refined particle data. result : dict[str, np.ndarray | torch.Tensor] The result of the refine template program. Can contain either np.ndarray (regular refine) or torch.Tensor (differentiable refine). prefer_refined_angles : bool Whether to use the refined angles or not. Defaults to True. + + Returns + ------- + pd.DataFrame + The refined particle data. Not written to disk; use ``export_results`` to do + both in one call. """ # pylint: disable=duplicate-code df_refined = self.particle_stack._df.copy() # pylint: disable=protected-access @@ -470,5 +521,49 @@ def refine_result_to_dataframe( # Reorder the columns df_refined = df_refined.reindex(columns=REFINED_DF_COLUMN_ORDER) - # Save the refined DataFrame to disk - df_refined.to_csv(output_dataframe_path) + return df_refined + + def export_results( + self, + output_dataframe_path: str, + result: dict[str, np.ndarray | torch.Tensor], + prefer_refined_angles: bool = True, + output_format: Literal["csv", "hdf5"] | None = None, + allow_file_overwrite: bool = False, + ) -> ParticleStackCSV | ParticleStackHDF5: + """Build the refined DataFrame and write it to disk. + + Parameters + ---------- + output_dataframe_path : str + Path to save the refined particle data. + result : dict[str, np.ndarray | torch.Tensor] + The result of the refine template program. Can contain either np.ndarray + (regular refine) or torch.Tensor (differentiable refine). + prefer_refined_angles : bool + Whether to use the refined angles or not. Defaults to True. + output_format : Literal["csv", "hdf5"] | None + Output back-end to write. Defaults to None, which matches the back-end of + ``self.particle_stack`` (CSV in, CSV out; HDF5 in, HDF5 out). Pass "csv" or + "hdf5" to override. + allow_file_overwrite : bool + Whether to overwrite an existing file at ``output_dataframe_path``. Defaults + to False. + + Returns + ------- + ParticleStackCSV | ParticleStackHDF5 + The refined particle stack, already written to ``output_dataframe_path``. + Reuse directly instead of re-reading from disk if feeding into another + program. + """ + df_refined = self.refine_result_to_dataframe( + result=result, prefer_refined_angles=prefer_refined_angles + ) + return export_particle_stack( + df=df_refined, + output_path=output_dataframe_path, + source_particle_stack=self.particle_stack, + output_format=output_format, + allow_file_overwrite=allow_file_overwrite, + ) 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..99a7bb1e --- /dev/null +++ b/src/leopard_em/pydantic_models/results/correlation_table.py @@ -0,0 +1,340 @@ +"""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 + + search_space = f.create_group("search_space") + search_space.create_dataset( + "defocus_offsets", + data=np.array(self.defocus_offsets, dtype=np.float32), + ) + search_space.create_dataset( + "phi_theta_angles", + data=np.array(self.phi_theta_angles, dtype=np.float32), + ) + search_space.create_dataset( + "psi_angles", + data=np.array(self.psi_angles, dtype=np.float32), + ) + + detections = f.create_group("detections") + detections.create_dataset( + "search_index", + data=np.array(self.search_index, dtype=np.int32), + ) + detections.create_dataset("x", data=np.array(self.x, dtype=np.int32)) + detections.create_dataset("y", data=np.array(self.y, dtype=np.int32)) + detections.create_dataset( + "correlation_value", + data=np.array(self.correlation_value, dtype=np.float32), + ) + detections.create_dataset( + "correlation_mean", + data=np.array(self.correlation_mean, dtype=np.float32), + ) + detections.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 + ) + + 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..27afa84a 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 @@ -31,7 +68,7 @@ def check_file_path_and_permissions(path: str, allow_overwrite: bool) -> None: # 2. Check write permissions if directory and not os.access(directory, os.W_OK): raise ValueError( - f"Directory '{directory}' does not permit writing." + f"Directory '{directory}' does not permit writing. " f"Will be unable to write results to '{path}'." ) @@ -39,94 +76,46 @@ def check_file_path_and_permissions(path: str, allow_overwrite: bool) -> None: if not allow_overwrite and os.path.exists(path): raise ValueError( f"File '{path}' already exists, but 'allow_file_overwrite' " - "is False. Set 'allow_file_overwrite' to True to permit. " + "is False. Set 'allow_file_overwrite' to True to permit " "overwriting.\n" "WARNING: Overwriting will delete the existing file(s)!" ) -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/src/leopard_em/utils/backend_setup.py b/src/leopard_em/utils/backend_setup.py index 2fe78350..df7b7764 100644 --- a/src/leopard_em/utils/backend_setup.py +++ b/src/leopard_em/utils/backend_setup.py @@ -4,7 +4,7 @@ import pandas as pd import torch -from torch_cubic_spline_grids import CubicCatmullRomGrid3d +from torch_motion_correction.deformation_field import DeformationField from leopard_em.utils.ctf_utils import _setup_ctf_kwargs_from_particle_stack from leopard_em.utils.image_processing import ( @@ -17,18 +17,25 @@ from leopard_em.pydantic_models.config.correlation_filters import ( PreprocessingFilters, ) - from leopard_em.pydantic_models.data_structures.particle_stack import ParticleStack + from leopard_em.pydantic_models.data_structures.particle_stack import ( + ParticleStackCSV, + ParticleStackHDF5, + ) # pylint: disable=too-many-locals # pylint: disable=too-many-arguments def _process_particle_images_for_filters( - particle_stack: "ParticleStack", + particle_stack: "ParticleStackCSV | ParticleStackHDF5", preprocessing_filters: "PreprocessingFilters", template: torch.Tensor, particle_images: torch.Tensor, apply_global_filtering: bool, - projective_filters: torch.Tensor | None, + fixed_image_filters: torch.Tensor | None, + fixed_projective_filters: torch.Tensor | None, + full_image_shape: tuple[int, int] | None = None, + extracted_box_shape: tuple[int, int] | None = None, + fixed_normalization_factor: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Process particle images and compute filters. @@ -36,7 +43,7 @@ def _process_particle_images_for_filters( Parameters ---------- - particle_stack : ParticleStack + particle_stack : ParticleStackCSV | ParticleStackHDF5 The particle stack containing images to process. preprocessing_filters : PreprocessingFilters Filters to apply to the particle images. @@ -46,8 +53,21 @@ def _process_particle_images_for_filters( The particle images to process. apply_global_filtering : bool Whether global filtering was applied. - projective_filters : torch.Tensor | None - Pre-computed projective filters (if global filtering was used). + fixed_image_filters : torch.Tensor | None + Optional per-particle Fourier filters matching ``particle_images`` RFFT shape, + applied when whitening particle images (same convention as + ``construct_image_filters(..., output_shape=images_dft.shape[-2:])``). + fixed_projective_filters : torch.Tensor | None + Optional per-particle filters at template RFFT resolution for the correlation + backend (same convention as historical ``construct_image_filters`` / + ``construct_projective_filters`` template ``output_shape``). + full_image_shape : tuple[int, int] | None + Optional full image shape used for normalization scaling. + extracted_box_shape : tuple[int, int] | None + Optional extracted box shape used for normalization scaling. + fixed_normalization_factor : torch.Tensor | None + Optional precomputed per-particle normalization factors to reuse for + each frame (old-style frame normalization). Returns ------- @@ -55,32 +75,69 @@ def _process_particle_images_for_filters( A tuple containing: - particle_images_dft: The particle images in Fourier space - template_dft: The Fourier transformed template - - projective_filters: Filters applied to the template + - projective_filters: Filters applied to template projections (backend) """ device = template.device box_h, box_w = particle_stack.extracted_box_size + filter_full_shape = full_image_shape or (box_h, box_w) + filter_extracted_shape = extracted_box_shape or (box_h, box_w) if not apply_global_filtering: particle_images_dft = torch.fft.rfftn(particle_images, dim=(-2, -1)) # pylint: disable=not-callable particle_images_dft[..., 0, 0] = 0.0 + 0.0j # Zero out DC component + raw_particle_dft = particle_images_dft.detach() + + image_filters = fixed_image_filters + if image_filters is None: + with torch.no_grad(): + image_filters = particle_stack.construct_image_filters( + preprocessing_filters, + output_shape=particle_images_dft.shape[-2:], + images_dft=raw_particle_dft, + ).to(device) - # Compute filters without gradient tracking (filters are just preprocessing) - with torch.no_grad(): - projective_filters = particle_stack.construct_image_filters( - preprocessing_filters, - output_shape=(template.shape[-2], template.shape[-1] // 2 + 1), - images_dft=particle_images_dft.detach(), - ).to(device) particle_images_dft = apply_image_filtering( particle_stack, preprocessing_filters, particle_images_dft, - full_image_shape=(box_h, box_w), - extracted_box_shape=(box_h, box_w), + full_image_shape=filter_full_shape, + extracted_box_shape=filter_extracted_shape, + precomputed_filter_stack=image_filters, + precomputed_normalization_factor=fixed_normalization_factor, ) + + projective_filters = fixed_projective_filters + if projective_filters is None: + with torch.no_grad(): + projective_filters = particle_stack.construct_image_filters( + preprocessing_filters, + output_shape=(template.shape[-2], template.shape[-1] // 2 + 1), + images_dft=raw_particle_dft, + ).to(device) else: particle_images_dft = torch.fft.rfftn(particle_images, dim=(-2, -1)) # pylint: disable=not-callable + particle_images_dft[..., 0, 0] = 0.0 + 0.0j + if fixed_image_filters is not None and fixed_normalization_factor is not None: + particle_images_dft = apply_image_filtering( + particle_stack, + preprocessing_filters, + particle_images_dft, + full_image_shape=filter_full_shape, + extracted_box_shape=filter_extracted_shape, + precomputed_filter_stack=fixed_image_filters, + precomputed_normalization_factor=fixed_normalization_factor, + ) + projective_filters = fixed_projective_filters + if projective_filters is None: + projective_filters = torch.ones( + ( + particle_images.shape[0], + template.shape[-2], + template.shape[-1] // 2 + 1, + ), + device=device, + ) template_dft = volume_to_rfft_fourier_slice(template) return ( @@ -90,16 +147,48 @@ def _process_particle_images_for_filters( ) +# pylint: disable=too-many-arguments +def setup_frame_filters_particle_stack( + particle_stack: "ParticleStackCSV | ParticleStackHDF5", + preprocessing_filters: "PreprocessingFilters", + template: torch.Tensor, + particle_images: torch.Tensor, + apply_global_filtering: bool = False, + fixed_image_filters: torch.Tensor | None = None, + fixed_projective_filters: torch.Tensor | None = None, + fixed_normalization_factor: torch.Tensor | None = None, + full_image_shape: tuple[int, int] | None = None, + extracted_box_shape: tuple[int, int] | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare backend image/template inputs for a per-frame particle stack. + + The frame-correlation path already extracts particles for a single movie frame, + so this helper only handles particle-local filtering and template DFT setup. + """ + return _process_particle_images_for_filters( + particle_stack=particle_stack, + preprocessing_filters=preprocessing_filters, + template=template, + particle_images=particle_images, + apply_global_filtering=apply_global_filtering, + fixed_image_filters=fixed_image_filters, + fixed_projective_filters=fixed_projective_filters, + full_image_shape=full_image_shape, + extracted_box_shape=extracted_box_shape, + fixed_normalization_factor=fixed_normalization_factor, + ) + + # pylint: disable=too-many-locals # pylint: disable=too-many-arguments # pylint: disable=too-many-positional-arguments def _setup_images_filters_from_micrographs( - particle_stack: "ParticleStack", + particle_stack: "ParticleStackCSV | ParticleStackHDF5", preprocessing_filters: "PreprocessingFilters", template: torch.Tensor, apply_global_filtering: bool, movie: torch.Tensor | None, - deformation_field: CubicCatmullRomGrid3d | None, + deformation_field: DeformationField | None, particle_shifts: torch.Tensor | None, pre_exposure: float, fluence_per_frame: float, @@ -112,7 +201,7 @@ def _setup_images_filters_from_micrographs( Parameters ---------- - particle_stack : ParticleStack + particle_stack : ParticleStackCSV | ParticleStackHDF5 The particle stack containing images to process. preprocessing_filters : PreprocessingFilters Filters to apply to the particle images. @@ -122,7 +211,7 @@ def _setup_images_filters_from_micrographs( If True, apply filtering to the full micrograph before particle extraction. movie: torch.Tensor | None The movie tensor. - deformation_field: CubicCatmullRomGrid3d | None + deformation_field: DeformationField | None The deformation field tensor. particle_shifts: torch.Tensor | None The particle shifts tensor. If provided, takes precedence over @@ -224,13 +313,14 @@ def _setup_images_filters_from_micrographs( template=template, particle_images=particle_images, apply_global_filtering=apply_global_filtering, - projective_filters=projective_filters, + fixed_image_filters=None, + fixed_projective_filters=projective_filters, ) # pylint: disable=too-many-arguments def _setup_images_filters_from_particles( - particle_stack: "ParticleStack", + particle_stack: "ParticleStackCSV | ParticleStackHDF5", preprocessing_filters: "PreprocessingFilters", template: torch.Tensor, apply_global_filtering: bool, @@ -242,7 +332,7 @@ def _setup_images_filters_from_particles( Parameters ---------- - particle_stack : ParticleStack + particle_stack : ParticleStackCSV | ParticleStackHDF5 The particle stack containing images to process. preprocessing_filters : PreprocessingFilters Filters to apply to the particle images. @@ -270,19 +360,20 @@ def _setup_images_filters_from_particles( template=template, particle_images=particle_images, apply_global_filtering=apply_global_filtering, - projective_filters=None, + fixed_image_filters=None, + fixed_projective_filters=None, ) # pylint: disable=too-many-arguments # pylint: disable=too-many-positional-arguments def setup_images_filters_particle_stack( - particle_stack: "ParticleStack", + particle_stack: "ParticleStackCSV | ParticleStackHDF5", preprocessing_filters: "PreprocessingFilters", template: torch.Tensor, apply_global_filtering: bool = True, movie: torch.Tensor | None = None, - deformation_field: CubicCatmullRomGrid3d | None = None, + deformation_field: DeformationField | None = None, particle_shifts: torch.Tensor | None = None, pre_exposure: float = 0.0, fluence_per_frame: float = 1.0, @@ -297,7 +388,7 @@ def setup_images_filters_particle_stack( Parameters ---------- - particle_stack : ParticleStack + particle_stack : ParticleStackCSV | ParticleStackHDF5 The particle stack containing images to process. preprocessing_filters : PreprocessingFilters Filters to apply to the particle images. @@ -309,7 +400,7 @@ def setup_images_filters_particle_stack( Default is True. movie: torch.Tensor | None The movie tensor. - deformation_field: CubicCatmullRomGrid3d | None + deformation_field: DeformationField | None The deformation field tensor. particle_shifts: torch.Tensor | None The particle shifts tensor. If provided, takes precedence over @@ -362,7 +453,7 @@ def setup_images_filters_particle_stack( # pylint: disable=too-many-arguments def _setup_correlation_stacks_from_micrographs( - particle_stack: "ParticleStack", + particle_stack: "ParticleStackCSV | ParticleStackHDF5", mean_stack: torch.Tensor | None, std_stack: torch.Tensor | None, particle_indices: list[pd.Index] | None, @@ -373,7 +464,7 @@ def _setup_correlation_stacks_from_micrographs( Parameters ---------- - particle_stack : ParticleStack + particle_stack : ParticleStackCSV | ParticleStackHDF5 The particle stack containing images to process. mean_stack : torch.Tensor | None Pre-loaded mean stack tensor. @@ -491,11 +582,125 @@ def _setup_correlation_stacks_from_particles( return mean_stack, corr_std_stack +def astigmatism_angle_tensor( + particle_stack: "ParticleStackCSV | ParticleStackHDF5", + device: torch.device, +) -> torch.Tensor: + """Return the per-particle astigmatism angle as a tensor on ``device``.""" + return torch.tensor(particle_stack["astigmatism_angle"], device=device) + + +# pylint: disable=too-many-locals +# pylint: disable=too-many-arguments +# pylint: disable=too-many-positional-arguments +def setup_static_particle_kwargs( + particle_stack: "ParticleStackCSV | ParticleStackHDF5", + template: torch.Tensor, + euler_angles: torch.Tensor, + euler_angle_offsets: torch.Tensor, + defocus_offsets: torch.Tensor, + pixel_size_offsets: torch.Tensor, + device_list: list, + mean_stack: torch.Tensor | None = None, + std_stack: torch.Tensor | None = None, + particle_indices: list[pd.Index] | None = None, + images_are_particles: bool = False, +) -> dict[str, Any]: + """Build the image-independent half of the backend kwargs. + + These are the per-particle inputs that do not depend on per-image filtering: + Euler angles/offsets, defocus values, correlation mean/std stacks, CTF kwargs, + and the magnification matrix. They are shared by the standard particle path + (:func:`setup_particle_backend_kwargs`) and the per-frame inspection path, which + recomputes image filters separately for each frame. + + Parameters + ---------- + particle_stack : ParticleStackCSV | ParticleStackHDF5 + The particle stack containing images to process. + template : torch.Tensor + The 3D template volume (used for device and CTF output shape). + euler_angles : torch.Tensor + The set of Euler angles to use. + euler_angle_offsets : torch.Tensor + The relative Euler angle offsets to search over. + defocus_offsets : torch.Tensor + The relative defocus values to search over. + pixel_size_offsets : torch.Tensor + The relative pixel size values to search over. + device_list : list + List of computational devices to use. + mean_stack : torch.Tensor | None + The mean stack tensor. + std_stack : torch.Tensor | None + The std stack tensor. + particle_indices : list[pd.Index] | None + The particle indices to process. + images_are_particles : bool + Whether the images are particles or not. Defaults to False. + + Returns + ------- + dict[str, Any] + Dictionary of the image-independent backend keyword arguments. + """ + device = template.device + h, w = particle_stack.original_template_size + box_h, box_w = particle_stack.extracted_box_size + extracted_box_size = (box_h - h + 1, box_w - w + 1) + + # Setup correlation stacks + if images_are_particles: + corr_mean_stack, corr_std_stack = _setup_correlation_stacks_from_particles( + mean_stack=mean_stack, + std_stack=std_stack, + particle_indices=particle_indices, + ) + else: + corr_mean_stack, corr_std_stack = _setup_correlation_stacks_from_micrographs( + particle_stack=particle_stack, + mean_stack=mean_stack, + std_stack=std_stack, + particle_indices=particle_indices, + extracted_box_size=extracted_box_size, + device=device, + ) + + # The best defocus values for each particle (+ astigmatism) + defocus_u, defocus_v = particle_stack.get_absolute_defocus() + defocus_u = defocus_u.to(device) + defocus_v = defocus_v.to(device) + defocus_angle = astigmatism_angle_tensor(particle_stack, device) + + ctf_kwargs = _setup_ctf_kwargs_from_particle_stack( + particle_stack, (template.shape[-2], template.shape[-1]) + ) + + # Extract mag_matrix from particle stack and convert to 2x2 tensor. + # All particles should have the same mag_matrix value. + mag_matrix_tensor = ctf_kwargs["mag_matrix"] + + return { + "euler_angles": euler_angles, + "euler_angle_offsets": euler_angle_offsets, + "defocus_u": defocus_u, + "defocus_v": defocus_v, + "defocus_angle": defocus_angle, + "defocus_offsets": defocus_offsets, + "pixel_size_offsets": pixel_size_offsets, + "corr_mean": corr_mean_stack, + "corr_std": corr_std_stack, + "ctf_kwargs": ctf_kwargs, + "device": device_list, + "mag_matrix": mag_matrix_tensor, + } + + # pylint: disable=too-many-locals # pylint: disable=too-many-arguments # pylint: disable=too-many-positional-arguments def setup_particle_backend_kwargs( - particle_stack: "ParticleStack", + particle_stack: "ParticleStackCSV | ParticleStackHDF5", template: torch.Tensor, preprocessing_filters: "PreprocessingFilters", euler_angles: torch.Tensor, @@ -505,7 +710,7 @@ def setup_particle_backend_kwargs( apply_global_filtering: bool, device_list: list, movie: torch.Tensor | None = None, - deformation_field: CubicCatmullRomGrid3d | None = None, + deformation_field: DeformationField | None = None, particle_shifts: torch.Tensor | None = None, pre_exposure: float = 0.0, fluence_per_frame: float = 1.0, @@ -522,7 +727,7 @@ def setup_particle_backend_kwargs( Parameters ---------- - particle_stack : ParticleStack + particle_stack : ParticleStackCSV | ParticleStackHDF5 The particle stack containing images to process. template : torch.Tensor The 3D template volume. @@ -543,7 +748,7 @@ def setup_particle_backend_kwargs( List of computational devices to use. movie: torch.Tensor | None The movie tensor. - deformation_field: CubicCatmullRomGrid3d | None + deformation_field: DeformationField | None The deformation field tensor. particle_shifts: torch.Tensor | None The particle shifts tensor. If provided, takes precedence over @@ -568,27 +773,19 @@ def setup_particle_backend_kwargs( dict[str, Any] Dictionary of keyword arguments for backend functions. """ - device = template.device - h, w = particle_stack.original_template_size - box_h, box_w = particle_stack.extracted_box_size - extracted_box_size = (box_h - h + 1, box_w - w + 1) - - # Setup correlation stacks - if images_are_particles: - corr_mean_stack, corr_std_stack = _setup_correlation_stacks_from_particles( - mean_stack=mean_stack, - std_stack=std_stack, - particle_indices=particle_indices, - ) - else: - corr_mean_stack, corr_std_stack = _setup_correlation_stacks_from_micrographs( - particle_stack=particle_stack, - mean_stack=mean_stack, - std_stack=std_stack, - particle_indices=particle_indices, - extracted_box_size=extracted_box_size, - device=device, - ) + static_kwargs = setup_static_particle_kwargs( + particle_stack=particle_stack, + template=template, + euler_angles=euler_angles, + euler_angle_offsets=euler_angle_offsets, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + device_list=device_list, + mean_stack=mean_stack, + std_stack=std_stack, + particle_indices=particle_indices, + images_are_particles=images_are_particles, + ) ( particle_images_dft, @@ -609,34 +806,9 @@ def setup_particle_backend_kwargs( images_are_particles=images_are_particles, ) - # The best defocus values for each particle (+ astigmatism) - defocus_u, defocus_v = particle_stack.get_absolute_defocus() - defocus_u = defocus_u.to(device) - defocus_v = defocus_v.to(device) - defocus_angle = torch.tensor(particle_stack["astigmatism_angle"], device=device) - - ctf_kwargs = _setup_ctf_kwargs_from_particle_stack( - particle_stack, (template.shape[-2], template.shape[-1]) - ) - - # Extract mag_matrix from particle stack and convert to 2x2 tensor - # All particles should have the same mag_matrix value - mag_matrix_tensor = ctf_kwargs["mag_matrix"] - return { + **static_kwargs, "particle_stack_dft": particle_images_dft, "template_dft": template_dft, - "euler_angles": euler_angles, - "euler_angle_offsets": euler_angle_offsets, - "defocus_u": defocus_u, - "defocus_v": defocus_v, - "defocus_angle": defocus_angle, - "defocus_offsets": defocus_offsets, - "pixel_size_offsets": pixel_size_offsets, - "corr_mean": corr_mean_stack, - "corr_std": corr_std_stack, - "ctf_kwargs": ctf_kwargs, "projective_filters": projective_filters, - "device": device_list, - "mag_matrix": mag_matrix_tensor, } diff --git a/src/leopard_em/utils/ctf_utils.py b/src/leopard_em/utils/ctf_utils.py index d7b216c3..83bb6465 100644 --- a/src/leopard_em/utils/ctf_utils.py +++ b/src/leopard_em/utils/ctf_utils.py @@ -13,7 +13,39 @@ # Using the TYPE_CHECKING statement to avoid circular imports if TYPE_CHECKING: from leopard_em.pydantic_models.data_structures.optics_group import OpticsGroup - from leopard_em.pydantic_models.data_structures.particle_stack import ParticleStack + from leopard_em.pydantic_models.data_structures.particle_stack import ( + ParticleStackCSV, + ParticleStackHDF5, + ) + + +def move_ctf_kwargs_tensors_to_device( + ctf_kwargs: dict[str, Any], device: torch.device +) -> dict[str, Any]: + """Copy CTF kwargs so tensor-valued entries live on ``device``. + + Used by multi-GPU refine/inspect workers: ``_setup_ctf_kwargs_from_particle_stack`` + builds ``mag_matrix`` and Zernike dicts on CPU; + ``calculate_ctf_filter_stack_full_args`` + must run with tensors on the same device as ``projective_filter`` / particle FFTs. + """ + kwargs = dict(ctf_kwargs) + mag = kwargs.get("mag_matrix") + if isinstance(mag, torch.Tensor): + kwargs["mag_matrix"] = mag.to(device=device) + ez = kwargs.get("even_zernikes") + if isinstance(ez, dict): + kwargs["even_zernikes"] = { + k: v.to(device=device) if isinstance(v, torch.Tensor) else v + for k, v in ez.items() + } + oz = kwargs.get("odd_zernikes") + if isinstance(oz, dict): + kwargs["odd_zernikes"] = { + k: v.to(device=device) if isinstance(v, torch.Tensor) else v + for k, v in oz.items() + } + return kwargs def calculate_ctf_filter_stack_full_args( @@ -200,13 +232,14 @@ def _parse_json_string_from_series_value(value: Any) -> dict | None: def _setup_ctf_kwargs_from_particle_stack( - particle_stack: "ParticleStack", template_shape: tuple[int, int] + particle_stack: "ParticleStackCSV | ParticleStackHDF5", + template_shape: tuple[int, int], ) -> dict[str, Any]: """Helper function for per-particle CTF kwargs. Parameters ---------- - particle_stack : ParticleStack + particle_stack : ParticleStackCSV | ParticleStackHDF5 The particle stack to extract the CTF parameters from. template_shape : tuple[int, int] The shape of the template to use for the CTF calculation. diff --git a/src/leopard_em/utils/image_processing.py b/src/leopard_em/utils/image_processing.py index 7b3fcf9a..e945208c 100644 --- a/src/leopard_em/utils/image_processing.py +++ b/src/leopard_em/utils/image_processing.py @@ -10,7 +10,10 @@ from leopard_em.pydantic_models.config.correlation_filters import ( PreprocessingFilters, ) - from leopard_em.pydantic_models.data_structures.particle_stack import ParticleStack + from leopard_em.pydantic_models.data_structures.particle_stack import ( + ParticleStackCSV, + ParticleStackHDF5, + ) def preprocess_image( @@ -20,7 +23,7 @@ def preprocess_image( full_image_shape: tuple[int, int], extracted_box_shape: tuple[int, int], ) -> torch.Tensor: - """Preprocesses and normalizes the image based on the given filters. + """Preprocess and normalize image FFTs with computed normalization factors. Parameters ---------- @@ -39,17 +42,53 @@ def preprocess_image( Returns ------- torch.Tensor - Preprocessed and normalized image in Fourier space + Preprocessed and normalized image in Fourier space. """ + normalization_factor = get_image_normalization_factor( + image_rfft=image_rfft, + cumulative_fourier_filters=cumulative_fourier_filters, + bandpass_filter=bandpass_filter, + full_image_shape=full_image_shape, + extracted_box_shape=extracted_box_shape, + ) image_rfft = image_rfft * cumulative_fourier_filters + return image_rfft * normalization_factor + - # Normalize the image after filtering - squared_image_rfft = torch.abs(image_rfft) ** 2 +def get_image_normalization_factor( + image_rfft: torch.Tensor, + cumulative_fourier_filters: torch.Tensor, + bandpass_filter: torch.Tensor, + full_image_shape: tuple[int, int], + extracted_box_shape: tuple[int, int], +) -> torch.Tensor: + """Compute per-image normalization factors for filtered Fourier images. + + Parameters + ---------- + image_rfft : torch.Tensor + Real Fourier-transformed images (unshifted), typically ``(..., H, W_rfft)``. + cumulative_fourier_filters : torch.Tensor + Combined Fourier-domain filtering stack applied to each image. + bandpass_filter : torch.Tensor + Bandpass filter used to estimate effective Fourier dimensionality. + full_image_shape : tuple[int, int] + Shape ``(H_full, W_full)`` of the source image used for extraction. + extracted_box_shape : tuple[int, int] + Shape ``(H_box, W_box)`` of extracted particle boxes. + + Returns + ------- + torch.Tensor + Multiplicative normalization factor broadcastable to ``image_rfft``. + """ + # Normalize image variance after filtering via Parseval-conjugate accounting. + filtered = image_rfft * cumulative_fourier_filters + squared_image_rfft = torch.abs(filtered) ** 2 squared_sum = torch.sum(squared_image_rfft, dim=(-2, -1), keepdim=True) squared_sum += torch.sum( squared_image_rfft[..., :, 1:-1], dim=(-2, -1), keepdim=True ) - image_rfft = image_rfft / torch.sqrt(squared_sum) # Non-in-place preserves gradient # NOTE: For two Gaussian random variables in d-dimensional space -- A and B -- # each with mean 0 and variance 1 their correlation will have on average a @@ -64,30 +103,32 @@ def preprocess_image( # Below, we calculate the dimensionality of our cross-correlation and divide # by the square root of that number to normalize the image. dimensionality = bandpass_filter.sum() + bandpass_filter[:, 1:-1].sum() - image_rfft = image_rfft * dimensionality**0.5 + normalization_factor = (dimensionality**0.5) / torch.sqrt(squared_sum) # NOTE: We need to rescale based on the relative area of the extracted box # to the full image. img_h, img_w = full_image_shape box_h, box_w = extracted_box_shape - image_rfft = image_rfft * ((img_h * img_w) / ((box_h) * (box_w))) ** 0.5 - - return image_rfft + normalization_factor = normalization_factor * ( + ((img_h * img_w) / ((box_h) * (box_w))) ** 0.5 + ) + return normalization_factor def apply_image_filtering( - particle_stack: "ParticleStack", + particle_stack: "ParticleStackCSV | ParticleStackHDF5", preprocessing_filters: "PreprocessingFilters", images_dft: torch.Tensor, full_image_shape: tuple[int, int], extracted_box_shape: tuple[int, int], + precomputed_filter_stack: torch.Tensor | None = None, + precomputed_normalization_factor: torch.Tensor | None = None, ) -> torch.Tensor: - """ - Apply filtering to a set of images. + """Apply filtering and normalization to a set of Fourier images. Parameters ---------- - particle_stack : ParticleStack + particle_stack : ParticleStackCSV | ParticleStackHDF5 The particle stack preprocessing_filters : PreprocessingFilters Filters to apply to the images. @@ -97,11 +138,17 @@ def apply_image_filtering( The shape of the full image. extracted_box_shape: tuple[int, int] The shape of the extracted box. + precomputed_filter_stack : torch.Tensor | None + Optional precomputed per-particle Fourier filter stack. If provided, + reuse these filters instead of recomputing from ``images_dft``. + precomputed_normalization_factor : torch.Tensor | None + Optional fixed per-particle normalization factors (old-style frame + normalization). If provided, apply these directly after filtering. Returns ------- torch.Tensor - The filtered images in Fourier space + Filtered images in Fourier space. """ device = images_dft.device @@ -112,11 +159,18 @@ def apply_image_filtering( images_dft.shape[-2:] ).to(device) ) - filter_stack = particle_stack.construct_image_filters( - preprocessing_filters, - output_shape=images_dft.shape[-2:], - images_dft=images_dft.detach(), - ).to(device) + if precomputed_filter_stack is None: + filter_stack = particle_stack.construct_image_filters( + preprocessing_filters, + output_shape=images_dft.shape[-2:], + images_dft=images_dft.detach(), + ).to(device) + else: + filter_stack = precomputed_filter_stack.to(device) + + if precomputed_normalization_factor is not None: + # Apply the precomputed filters and fixed normalization factor directly. + return images_dft * filter_stack * precomputed_normalization_factor.to(device) return preprocess_image( image_rfft=images_dft, diff --git a/tests/analysis/test_inspect_peaks_result.py b/tests/analysis/test_inspect_peaks_result.py new file mode 100644 index 00000000..faa77d5a --- /dev/null +++ b/tests/analysis/test_inspect_peaks_result.py @@ -0,0 +1,113 @@ +"""Round-trip tests for saving/loading self-describing inspection ``.npz`` files.""" + +import numpy as np +import torch + +from leopard_em.analysis.inspect_peaks_result import ( + CROSS_CORRELATION_AXES, + CROSS_CORRELATION_FRAME_AXES, + FRC_AXES, + FRC_FRAME_AXES, + load_inspection_result, + save_inspection_result, +) + + +def _common_kwargs(n_particles: int = 2, n_orient: int = 3): + """Build the shared base orientation/defocus/offset arrays for a tiny run.""" + return { + "euler_angle_offsets": torch.zeros((n_orient, 3)), + "defocus_offsets": torch.tensor([-100.0, 0.0, 100.0]), + "pixel_size_offsets": torch.tensor([0.0]), + "base_euler_angles": torch.arange(n_particles * 3, dtype=torch.float32).reshape( + n_particles, 3 + ), + "base_defocus": torch.tensor([[1.0, 2.0, 0.5], [3.0, 4.0, 1.5]]), + } + + +def test_save_load_cross_correlation_roundtrip(tmp_path): + """Spatial CC result round-trips with base orientation + astigmatic defocus.""" + # (N=2, n_px=1, n_def=3, n_orient=3, H=2, W=2) + scores = torch.randn((2, 1, 3, 3, 2, 2)) + kwargs = _common_kwargs() + + path = save_inspection_result( + tmp_path / "spatial", + result=scores, + output_mode="cross_correlation", + particle_index=torch.tensor([10, 11]), + **kwargs, + ) + assert path.suffix == ".npz" + + result = load_inspection_result(path) + assert result.output_mode == "cross_correlation" + assert result.axes == CROSS_CORRELATION_AXES + assert result.metadata["per_frame"] is False + assert result.scores.shape == (2, 1, 3, 3, 2, 2) + np.testing.assert_allclose(result.scores, scores.numpy()) + assert result.base_defocus.shape == (2, 3) + np.testing.assert_allclose(result.base_defocus, kwargs["base_defocus"].numpy()) + np.testing.assert_allclose( + result.base_euler_angles, kwargs["base_euler_angles"].numpy() + ) + np.testing.assert_array_equal(result.particle_index, np.array([10, 11])) + assert result.frequency_bins is None + assert result.frame_index is None + + +def test_save_load_per_frame_cross_correlation_inserts_frame_axis(tmp_path): + """Per-frame CC result carries a ``frame`` axis and a stored frame index.""" + # (N=2, T=4, n_px=1, n_def=3, n_orient=3, H=2, W=2) + scores = torch.randn((2, 4, 1, 3, 3, 2, 2)) + kwargs = _common_kwargs() + + path = save_inspection_result( + tmp_path / "per_frame", + result=scores, + output_mode="cross_correlation", + per_frame=True, + frame_index=torch.arange(4), + **kwargs, + ) + + result = load_inspection_result(path) + assert result.axes == CROSS_CORRELATION_FRAME_AXES + assert result.axes[:2] == ("particle", "frame") + assert result.metadata["per_frame"] is True + assert result.scores.shape == (2, 4, 1, 3, 3, 2, 2) + np.testing.assert_array_equal(result.frame_index, np.arange(4)) + + +def test_save_load_frc_modes(tmp_path): + """FRC results round-trip with frequency bins in spatial and per-frame layouts.""" + freq_bins = torch.linspace(0.0, 0.5, 5) + kwargs = _common_kwargs() + + # Spatial FRC: (N, n_px, n_def, n_orient, n_freq) + spatial = torch.randn((2, 1, 3, 3, 5)) + spatial_path = save_inspection_result( + tmp_path / "frc_spatial", + result=(spatial, freq_bins), + output_mode="frc", + **kwargs, + ) + spatial_result = load_inspection_result(spatial_path) + assert spatial_result.axes == FRC_AXES + np.testing.assert_allclose(spatial_result.frequency_bins, freq_bins.numpy()) + + # Per-frame FRC: (N, T, n_px, n_def, n_orient, n_freq) + per_frame = torch.randn((2, 4, 1, 3, 3, 5)) + per_frame_path = save_inspection_result( + tmp_path / "frc_per_frame", + result=(per_frame, freq_bins), + output_mode="frc", + per_frame=True, + frame_index=torch.arange(4), + **kwargs, + ) + per_frame_result = load_inspection_result(per_frame_path) + assert per_frame_result.axes == FRC_FRAME_AXES + assert per_frame_result.scores.shape == (2, 4, 1, 3, 3, 5) + np.testing.assert_allclose(per_frame_result.frequency_bins, freq_bins.numpy()) diff --git a/tests/backend/test_core_match_template.py b/tests/backend/test_core_match_template.py index efe1b8b2..aa71f3c4 100644 --- a/tests/backend/test_core_match_template.py +++ b/tests/backend/test_core_match_template.py @@ -33,8 +33,16 @@ ORIENTATION_BATCH_SIZE = 20 -def download_comparison_data() -> None: - """Downloads the example data from Zenodo.""" +def download_comparison_data(force_download: bool = False) -> None: + """Downloads the example data from Zenodo, skipping if already present. + + Parameters + ---------- + force_download : bool + If True, re-download the data even if it appears to already be present. + """ + if YAML_PATH.exists() and not force_download: + return subprocess.run(["zenodo_get", "--output-dir=tests/tmp", ZENODO_URL], check=True) @@ -63,8 +71,14 @@ def test_core_match_template(): 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", @@ -133,21 +147,25 @@ def get_diff_indices(path_a: str, path_b: str): theta_set = set(zip(diff_theta[0], diff_theta[1])) psi_set = set(zip(diff_psi[0], diff_psi[1])) - assert len(defocus_set.intersection(phi_set)) / len(defocus_set) - assert len(defocus_set.intersection(theta_set)) / len(defocus_set) - assert len(defocus_set.intersection(psi_set)) / len(defocus_set) - - assert len(phi_set.intersection(defocus_set)) / len(phi_set) - assert len(phi_set.intersection(theta_set)) / len(phi_set) - assert len(phi_set.intersection(psi_set)) / len(phi_set) - - assert len(theta_set.intersection(defocus_set)) / len(theta_set) - assert len(theta_set.intersection(phi_set)) / len(theta_set) - assert len(theta_set.intersection(psi_set)) / len(theta_set) - - assert len(psi_set.intersection(defocus_set)) / len(psi_set) - assert len(psi_set.intersection(phi_set)) / len(psi_set) - assert len(psi_set.intersection(theta_set)) / len(psi_set) + if len(defocus_set) > 0: + assert len(defocus_set.intersection(phi_set)) / len(defocus_set) + assert len(defocus_set.intersection(theta_set)) / len(defocus_set) + assert len(defocus_set.intersection(psi_set)) / len(defocus_set) + + if len(phi_set) > 0: + assert len(phi_set.intersection(defocus_set)) / len(phi_set) + assert len(phi_set.intersection(theta_set)) / len(phi_set) + assert len(phi_set.intersection(psi_set)) / len(phi_set) + + if len(theta_set) > 0: + assert len(theta_set.intersection(defocus_set)) / len(theta_set) + assert len(theta_set.intersection(phi_set)) / len(theta_set) + assert len(theta_set.intersection(psi_set)) / len(theta_set) + + if len(psi_set) > 0: + assert len(psi_set.intersection(defocus_set)) / len(psi_set) + assert len(psi_set.intersection(phi_set)) / len(psi_set) + assert len(psi_set.intersection(theta_set)) / len(psi_set) if __name__ == "__main__": diff --git a/tests/backend/test_core_refine_template.py b/tests/backend/test_core_refine_template.py new file mode 100644 index 00000000..33fa033d --- /dev/null +++ b/tests/backend/test_core_refine_template.py @@ -0,0 +1,83 @@ +"""Characterization tests for the refine-template per-particle reducer. + +These lock the behavior of ``_reduce_refine_best_zscore`` on deterministic synthetic +correlation batches so that backend refactors (e.g. moving z-score computation +out of the batch generator and into the reducer) provably do not change the +refined statistics. Runs on CPU; no GPU or downloaded data required. +""" + +import torch + +from leopard_em.backend.core_refine_template import _reduce_refine_best_zscore + +# Shapes for the synthetic search space. +N_PX, N_DEF, CROP_H, CROP_W = 2, 3, 4, 5 +BATCH_SIZES = (3, 2) # two orientation batches, 5 offsets total + + +def _build_synthetic_batches() -> tuple[list[tuple], torch.Tensor, torch.Tensor]: + """Build deterministic correlation batches in the post-refactor tuple format. + + Returns the batch list ``(start_idx, angle_offsets, cross_correlation, + crop_h, crop_w)`` (no pre-computed z-score) plus the ``corr_mean`` and + ``corr_std`` maps. The RNG draw order is fixed so the result is reproducible. + """ + torch.manual_seed(0) + corr_mean = torch.randn(CROP_H, CROP_W) + corr_std = torch.rand(CROP_H, CROP_W) + 0.5 # strictly positive + euler_offsets = torch.randn(sum(BATCH_SIZES), 3) + + batches = [] + start = 0 + for batch_size in BATCH_SIZES: + cross_correlation = torch.randn(N_PX, N_DEF, batch_size, CROP_H, CROP_W) + batches.append( + ( + start, + euler_offsets[start : start + batch_size], + cross_correlation, + CROP_H, + CROP_W, + ) + ) + start += batch_size + + return batches, corr_mean, corr_std + + +def test_reduce_refine_best_zscore_matches_snapshot(): + """The reducer reproduces the pre-refactor refined statistics exactly. + + Expected values were captured from the original implementation (which + pre-computed z-score inside the batch generator) on the same synthetic data. + """ + batches, corr_mean, corr_std = _build_synthetic_batches() + defocus_offsets = torch.tensor([-10.0, 0.0, 10.0]) # len N_DEF + pixel_size_offsets = torch.tensor([-0.01, 0.01]) # len N_PX + + result = _reduce_refine_best_zscore( + iter(batches), + corr_mean=corr_mean, + corr_std=corr_std, + defocus_offsets=defocus_offsets, + pixel_size_offsets=pixel_size_offsets, + ) + + def _val(key: str) -> float: + value = result[key] + return value.item() if torch.is_tensor(value) else value + + assert _val("max_cc") == 3.4028263092041016 + assert _val("max_z_score") == 5.6347527503967285 + assert _val("refined_phi_offset") == -2.2187793254852295 + assert _val("refined_theta_offset") == 0.2589845359325409 + assert _val("refined_psi_offset") == -1.0297021865844727 + assert _val("refined_defocus_offset") == 0.0 + assert _val("refined_pixel_size_offset") == 0.009999999776482582 + assert _val("refined_pos_y") == 1 + assert _val("refined_pos_x") == 0 + assert _val("angle_idx") == 2 + + +if __name__ == "__main__": + test_reduce_refine_best_zscore_matches_snapshot() 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_frame_inspection_manager.py b/tests/pydantic_models/test_frame_inspection_manager.py new file mode 100644 index 00000000..16d17076 --- /dev/null +++ b/tests/pydantic_models/test_frame_inspection_manager.py @@ -0,0 +1,142 @@ +"""Unit tests for the per-frame peak inspection manager.""" + +import pandas as pd +import pytest +import torch + +from leopard_em.pydantic_models.config import ( + ComputationalConfigRefine, + DefocusSearchConfig, + MovieConfig, + PixelSizeSearchConfig, + PreprocessingFilters, + RefineOrientationConfig, +) +from leopard_em.pydantic_models.data_structures.particle_stack import ParticleStack +from leopard_em.pydantic_models.managers.frame_inspection_manager import ( + FrameInspectionManager, +) + + +def make_frame_inspection_manager() -> FrameInspectionManager: + """Construct a minimal frame inspection manager for method-level tests.""" + particle_stack = ParticleStack( + df_path="", + extracted_box_size=(2, 2), + original_template_size=(2, 2), + skip_df_load=True, + ) + return FrameInspectionManager.model_construct( + template_volume_path="", + particle_stack=particle_stack, + defocus_refinement_config=DefocusSearchConfig(enabled=False), + pixel_size_refinement_config=PixelSizeSearchConfig( + enabled=False, + pixel_size_step=1.0, + ), + orientation_refinement_config=RefineOrientationConfig(enabled=False), + preprocessing_filters=PreprocessingFilters(), + computational_config=ComputationalConfigRefine(gpu_ids="cpu"), + movie_config=MovieConfig(enabled=False), + apply_global_filtering=True, + template_volume=None, + ) + + +# --------------------------------------------------------------------------- +# _stack_frame_results: pure static stacking of per-frame backend outputs. +# --------------------------------------------------------------------------- + + +def test_stack_frame_results_cross_correlation(): + """Cross-correlation results stack with frame as the leading axis, in order.""" + frame_results = [ + torch.full((2, 1, 1, 1, 2, 2), float(i), dtype=torch.float32) for i in range(3) + ] + stacked = FrameInspectionManager._stack_frame_results( + frame_results, output_mode="cross_correlation" + ) + assert isinstance(stacked, torch.Tensor) + assert stacked.shape == (3, 2, 1, 1, 1, 2, 2) + assert torch.all(stacked[0] == 0) + assert torch.all(stacked[2] == 2) + + +def test_stack_frame_results_frc_returns_tensor_and_shared_bins(): + """FRC results stack the spectra and carry the frequency bins through once.""" + freq_bins = torch.linspace(0.0, 0.5, 4) + frame_results = [ + (torch.full((2, 1, 1, 1, 4), float(i)), freq_bins) for i in range(3) + ] + stacked, bins = FrameInspectionManager._stack_frame_results( + frame_results, output_mode="frc" + ) + assert stacked.shape == (3, 2, 1, 1, 1, 4) + assert torch.equal(bins, freq_bins) + + +def test_stack_frame_results_empty_raises(): + """An empty result list is an error in either output mode.""" + with pytest.raises(ValueError, match="No frame results"): + FrameInspectionManager._stack_frame_results([], output_mode="cross_correlation") + + +# --------------------------------------------------------------------------- +# _frame_dose_template: per-frame template selection / dose weighting. +# --------------------------------------------------------------------------- + + +def test_frame_dose_template_disabled_returns_same_template(): + """With dose weighting off, the shared template is returned unchanged.""" + manager = make_frame_inspection_manager() + template = torch.randn((4, 4, 4)) + result = manager._frame_dose_template( + template, frame_idx=2, apply_template_dose_weighting=False + ) + assert result is template + + +def test_frame_dose_template_applies_distinct_dose_per_frame(): + """Each frame's template is dose-weighted over its own exposure interval.""" + manager = make_frame_inspection_manager() + manager.particle_stack._df = pd.DataFrame( + { + "particle_index": [0, 1], + "pixel_size": [1.0, 1.0], + "refined_pixel_size": [1.0, 1.0], + } + ) + manager.movie_config = MovieConfig( + enabled=True, + movie_path="movie.mrc", + pre_exposure=0.0, + fluence_per_frame=10.0, + ) + + torch.manual_seed(0) + template = torch.randn((6, 6, 6)) + + frame0 = manager._frame_dose_template( + template, frame_idx=0, apply_template_dose_weighting=True + ) + frame1 = manager._frame_dose_template( + template, frame_idx=1, apply_template_dose_weighting=True + ) + + # Dose weighting actually attenuates the template ... + assert not torch.allclose(frame0, template) + # ... and later frames (higher cumulative exposure) differ from earlier ones. + assert not torch.allclose(frame0, frame1) + assert frame0.shape == template.shape + + +# --------------------------------------------------------------------------- +# run_peak_inspection_per_frame: real input-validation guard (no backend). +# --------------------------------------------------------------------------- + + +def test_run_peak_inspection_per_frame_requires_movie(): + """Per-frame inspection fails fast when the movie config is disabled.""" + manager = make_frame_inspection_manager() + with pytest.raises(ValueError, match="requires movie_config.enabled"): + manager.run_peak_inspection_per_frame(template_tensor=torch.zeros((2, 2, 2))) 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/pydantic_models/test_particle_stack.py b/tests/pydantic_models/test_particle_stack.py index d3f017af..70c1d3af 100644 --- a/tests/pydantic_models/test_particle_stack.py +++ b/tests/pydantic_models/test_particle_stack.py @@ -4,8 +4,15 @@ import numpy as np import pandas as pd import pytest +import torch -from leopard_em.pydantic_models.data_structures.particle_stack import ParticleStack +from leopard_em.pydantic_models.data_structures.particle_stack import ( + ParticleStack, + ParticleStackHDF5, +) + +# Tests construct minimal ParticleStack instances by assigning their backing frames. +# pylint: disable=protected-access REQUIRED_COLUMNS = [ "particle_index", @@ -374,3 +381,219 @@ def test_particle_stack_top_left_and_center_self_consistency(): assert np.allclose(extracted_mips_tl[1], mip2_tl) assert np.allclose(extracted_mips_center[0], mip1_center) assert np.allclose(extracted_mips_center[1], mip2_center) + + +def test_construct_particle_movie_stack_assumes_aligned_without_motion_source(): + """Per-frame movie extraction should support already aligned movies.""" + df = make_minimal_df(num_rows=1) + df["pos_x"] = [1] + df["pos_y"] = [1] + df["pixel_size"] = [1.0] + + particle_stack = ParticleStack( + df_path="", + extracted_box_size=(2, 2), + original_template_size=(2, 2), + skip_df_load=True, + ) + particle_stack._df = df + movie = torch.arange(2 * 4 * 4, dtype=torch.float32).reshape(2, 4, 4) + + with pytest.warns(UserWarning, match="movie is already aligned"): + particle_movie = particle_stack.construct_particle_movie_stack( + movie=movie, + use_gradient_checkpointing=False, + ) + + mean_zero_movie = movie - torch.mean(movie, dim=(-2, -1), keepdim=True) + expected = mean_zero_movie[:, None, 1:3, 1:3] + assert particle_movie.shape == (2, 1, 2, 2) + assert torch.allclose(particle_movie, expected) + + +def test_construct_image_stack_from_movie_preserves_summed_output_shape(monkeypatch): + """The existing movie stack API should still return one image per particle.""" + df = make_minimal_df(num_rows=1) + df["pos_x"] = [1] + df["pos_y"] = [1] + df["pixel_size"] = [1.0] + df["voltage"] = [300.0] + + particle_stack = ParticleStack( + df_path="", + extracted_box_size=(2, 2), + original_template_size=(2, 2), + skip_df_load=True, + ) + particle_stack._df = df + movie = torch.arange(2 * 4 * 4, dtype=torch.float32).reshape(2, 4, 4) + particle_shifts = torch.zeros((2, 1, 2), dtype=torch.float32) + dose_weight_calls = [] + + def fake_dose_weight_movie_to_micrograph( + movie_fft, pixel_size, pre_exposure, fluence_per_frame, voltage + ): + dose_weight_calls.append((pixel_size, pre_exposure, fluence_per_frame, voltage)) + return torch.fft.irfftn( # pylint: disable=not-callable + movie_fft, s=(2, 2), dim=(-2, -1) + ).sum(dim=0) + + monkeypatch.setattr( + "leopard_em.pydantic_models.data_structures.particle_stack." + "dose_weight_movie_to_micrograph", + fake_dose_weight_movie_to_micrograph, + ) + + image_stack = particle_stack.construct_image_stack_from_movie( + movie=movie, + particle_shifts=particle_shifts, + use_gradient_checkpointing=False, + ) + + assert image_stack.shape == (1, 2, 2) + assert particle_stack.image_stack.shape == (1, 2, 2) + assert len(dose_weight_calls) == 1 + + +def test_get_local_stat_maps_does_not_mutate_image_stack(): + """Ensure get_local_stat_maps doesn't clobber image_stack like construct does.""" + df, _r1, _r2, mip1, mip2 = make_reference_example_df() + h, w = (32, 32) + box_h, box_w = (34, 34) + + ps = ParticleStack( + df_path="", + extracted_box_size=(box_h, box_w), + original_template_size=(h, w), + skip_df_load=True, + ) + ps._df = df + + sentinel = torch.full((2, box_h, box_w), -1.0) + ps.image_stack = sentinel + + stat_maps = ps.get_local_stat_maps(columns=["mip_path"]) + + # image_stack must be untouched + assert ps.image_stack is sentinel + assert torch.equal(ps.image_stack, sentinel) + + mip1_ground_truth = np.array([[0, 0, 0], [0, mip1, 0], [0, 0, 0]], dtype=np.float32) + mip2_ground_truth = np.array([[0, 0, 0], [0, mip2, 0], [0, 0, 0]], dtype=np.float32) + + assert stat_maps["mip_path"].shape == (2, 3, 3) + assert np.allclose(stat_maps["mip_path"][0].numpy(), mip1_ground_truth) + assert np.allclose(stat_maps["mip_path"][1].numpy(), mip2_ground_truth) + + +def test_hdf5_get_local_stat_maps_returns_stored_values(tmp_path): + """When present in local_stats, stored tensors are returned as-is, not re-derived. + + The referenced path columns point at files that do not exist; if + ``get_local_stat_maps`` fell back to loading from disk this would raise. + Also covers a non-correlation column (``mip_path``) to show storage isn't + limited to correlation average/variance. + """ + df = make_minimal_df(num_rows=2) + df["correlation_average_path"] = "/nonexistent/does_not_exist.mrc" + df["correlation_variance_path"] = "/nonexistent/does_not_exist.mrc" + df["mip_path"] = "/nonexistent/does_not_exist.mrc" + df.index = pd.Index(["particle_00000", "particle_00001"], name="particle_id") + + stored_avg = torch.arange(2 * 3 * 3, dtype=torch.float32).reshape(2, 3, 3) + stored_var = torch.ones(2, 3, 3) + stored_mip = torch.full((2, 3, 3), 7.0) + + ps = ParticleStackHDF5( + hdf5_path=str(tmp_path / "particles.h5"), + extracted_box_size=(34, 34), + original_template_size=(32, 32), + local_stats_stored=True, + local_stats={ + "correlation_average_path": stored_avg, + "correlation_variance_path": stored_var, + "mip_path": stored_mip, + }, + skip_df_load=True, + ) + ps._df = df + + stat_maps = ps.get_local_stat_maps( + columns=["correlation_average_path", "correlation_variance_path", "mip_path"] + ) + + assert torch.equal(stat_maps["correlation_average_path"], stored_avg) + assert torch.equal(stat_maps["correlation_variance_path"], stored_var) + assert torch.equal(stat_maps["mip_path"], stored_mip) + + +def test_hdf5_to_hdf5_local_stats_requires_non_empty(tmp_path): + """to_hdf5(include_local_stats=True) raises when local_stats is empty.""" + df = make_minimal_df(num_rows=2) + df.index = pd.Index(["particle_00000", "particle_00001"], name="particle_id") + + ps = ParticleStackHDF5( + hdf5_path=str(tmp_path / "particles.h5"), + extracted_box_size=(34, 34), + original_template_size=(32, 32), + skip_df_load=True, + ) + ps._df = df + + with pytest.raises(ValueError, match="local_stats is empty"): + ps.to_hdf5(include_local_stats=True) + + +def test_hdf5_local_stats_roundtrip_stores_all_columns(tmp_path): + """to_hdf5/from_hdf5 round-trips every entry in local_stats, not just avg/var.""" + df = make_minimal_df(num_rows=2) + df.index = pd.Index(["particle_00000", "particle_00001"], name="particle_id") + + local_stats = { + "correlation_average_path": torch.zeros(2, 3, 3), + "correlation_variance_path": torch.ones(2, 3, 3), + "mip_path": torch.full((2, 3, 3), 2.0), + "psi_path": torch.full((2, 3, 3), 3.0), + "theta_path": torch.full((2, 3, 3), 4.0), + "phi_path": torch.full((2, 3, 3), 5.0), + } + + hdf5_path = str(tmp_path / "particles.h5") + ps = ParticleStackHDF5( + hdf5_path=hdf5_path, + extracted_box_size=(34, 34), + original_template_size=(32, 32), + local_stats=dict(local_stats), + skip_df_load=True, + ) + ps._df = df + ps.to_hdf5(include_local_stats=True) + + loaded = ParticleStackHDF5.from_hdf5(hdf5_path) + + assert loaded.local_stats_stored is True + assert set(loaded.local_stats.keys()) == set(local_stats.keys()) + for column, tensor in local_stats.items(): + assert torch.equal(loaded.local_stats[column], tensor) + + +def test_hdf5_get_local_stat_maps_falls_back_when_not_stored(tmp_path): + """Without local_stats_stored, columns are still loaded/cropped from disk.""" + df, _r1, _r2, mip1, mip2 = make_reference_example_df() + df.index = pd.Index(["particle_00000", "particle_00001"], name="particle_id") + + ps = ParticleStackHDF5( + hdf5_path=str(tmp_path / "particles.h5"), + extracted_box_size=(34, 34), + original_template_size=(32, 32), + skip_df_load=True, + ) + ps._df = df + + stat_maps = ps.get_local_stat_maps(columns=["mip_path"]) + + mip1_ground_truth = np.array([[0, 0, 0], [0, mip1, 0], [0, 0, 0]], dtype=np.float32) + mip2_ground_truth = np.array([[0, 0, 0], [0, mip2, 0], [0, 0, 0]], dtype=np.float32) + + assert np.allclose(stat_maps["mip_path"][0].numpy(), mip1_ground_truth) + assert np.allclose(stat_maps["mip_path"][1].numpy(), mip2_ground_truth) diff --git a/tests/self_consistency/test_backend_cross_correlate.py b/tests/self_consistency/test_backend_cross_correlate.py index e1f0fbae..0f4a81e3 100644 --- a/tests/self_consistency/test_backend_cross_correlate.py +++ b/tests/self_consistency/test_backend_cross_correlate.py @@ -25,16 +25,18 @@ from torch_fourier_filter.envelopes import b_envelope from leopard_em.backend.cross_correlation import ( + ZIPFFT_AVAILABLE, 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 +135,62 @@ 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") +@pytest.mark.skipif(not ZIPFFT_AVAILABLE, reason="zipfft package not installed") +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() + + 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" + ) + + +@pytest.mark.skipif(ZIPFFT_AVAILABLE, reason="zipfft package is installed") +def test_batched_zipfft_cross_correlate_raises_import_error_when_unavailable(): + """When zipfft isn't installed, calling the zipfft backend should raise clearly.""" + dummy = torch.zeros(1) + with pytest.raises(ImportError, match="zipfft"): + do_batched_orientation_cross_correlate_zipfft( + image_dft=dummy, + template_dft=dummy, + rotation_matrices=dummy, + projective_filters=dummy, + ) diff --git a/tests/utils/test_crop_extraction.py b/tests/utils/test_crop_extraction.py index 6bc6889d..ea33580d 100644 --- a/tests/utils/test_crop_extraction.py +++ b/tests/utils/test_crop_extraction.py @@ -150,7 +150,7 @@ def test_get_cropped_image_regions_torch_random_nonoverlapping(): y = np.random.randint(0, image_size[0] - box_size[0] + 1) x = np.random.randint(0, image_size[1] - box_size[1] + 1) if all( - not (y <= py < y + box_size[0] and x <= px < x + box_size[1]) + not (abs(y - py) < box_size[0] and abs(x - px) < box_size[1]) for py, px in positions ): positions.append((y, x)) diff --git a/zensical.toml b/zensical.toml new file mode 100644 index 00000000..efe3d9ab --- /dev/null +++ b/zensical.toml @@ -0,0 +1,126 @@ +[project] +site_name = "Leopard-EM: Two-Dimensional Template Matching in Python" + +# MathJax for LaTeX rendering in tutorials +extra_javascript = [ + "https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.4/MathJax.js?config=TeX-AMS-MML_HTMLorMML", +] + +nav = [ + { "Home" = "index.md" }, + { "Tutorials" = [ + { "2DTM Introduction" = "tutorials/match_template_intro.md" }, + { "Batch Processing" = "tutorials/batch_processing.md" }, + { "Distributed Computing" = "tutorials/distributed_match_template.md" }, + ]}, + { "Examples" = [ + { "Match Template Config" = "examples/01_basic_configuration.ipynb" }, + { "Peaks to DataFrame" = "examples/02_extract_peak_info.ipynb" }, + { "Comparing Scoring Metrics" = "examples/03_compare_scoring_metrics.ipynb" }, + { "Plotting match template results" = "examples/04_plotting_2dtm_results.ipynb" }, + { "Structure re-projection plot" = "examples/05_structure_reprojection.ipynb" }, + ]}, + { "Programs" = [ + { "Overview" = "programs/overview.md" }, + { "Match Template details" = "programs/match_template.md" }, + { "Refine Template details" = "programs/refine_template.md" }, + { "Optimize Template details" = "programs/optimize_template.md" }, + { "Constrained Search details" = "programs/constrained_search.md" }, + { "Peak & Frame Inspection" = "programs/inspect_peaks.md" }, + ]}, + { "Program Output Formats" = [ + { "Overview" = "data_formats.md" }, + { "Match Template Results (MRC vs HDF5)" = "data_formats/match_template_results.md" }, + { "Particle Stacks (CSV vs HDF5)" = "data_formats/particle_stack.md" }, + ]}, + { "Benchmarks" = "benchmarks.md" }, + { "Common Issues" = "common_issues.md" }, + { "API Reference" = [ + { "Overview" = "api/index.md" }, + { "Analysis" = "api/analysis.md" }, + { "Backend" = "api/backend.md" }, + { "Pydantic Models" = "api/pydantic_models.md" }, + { "Config Models" = "api/pydantic_models/config.md" }, + { "Data Structures" = "api/pydantic_models/data_structures.md" }, + { "Managers" = "api/pydantic_models/managers.md" }, + { "Results" = "api/pydantic_models/results.md" }, + { "Utilities" = "api/utils.md" }, + ]}, + { "Contributing" = "contributing.md" }, +] + +# ---------------------------------------------------------------------------- +# Theme +# ---------------------------------------------------------------------------- + +[project.theme] +features = [ + "content.code.copy", +] + +[[project.theme.palette]] +scheme = "slate" +primary = "purple" +toggle.icon = "lucide/moon" +toggle.name = "Switch to light mode" + +[[project.theme.palette]] +scheme = "default" +primary = "purple" +toggle.icon = "lucide/sun" +toggle.name = "Switch to dark mode" + + +# ---------------------------------------------------------------------------- +# Plugins +# NOTE: "search" is intentionally omitted — Zensical includes its Disco search +# engine built-in. The inline-select-svg plugin may not yet be Zensical- +# compatible; remove it if the build fails and it complains about unknown plugin. +# ---------------------------------------------------------------------------- + +[project.plugins.autorefs] + +[project.plugins."mkdocs-jupyter"] + +[project.plugins."inline-select-svg"] + +[project.plugins.mkdocstrings] +default_handler = "python" + +[project.plugins.mkdocstrings.handlers.python] +paths = ["src"] + +[project.plugins.mkdocstrings.handlers.python.options] +show_source = false +preload_modules = ["pydantic"] +show_symbol_type_toc = true +show_symbol_type_heading = true +show_root_full_path = true +show_root_heading = true +group_by_category = true +show_labels = false +inherited_members = false +members_order = "source" +docstring_section_style = "list" +merge_init_into_class = false +docstring_style = "numpy" +show_signature_annotations = true +signature_crossrefs = false +separate_signature = true +line_length = 60 + +# ---------------------------------------------------------------------------- +# Markdown extensions +# ---------------------------------------------------------------------------- + +[project.markdown_extensions.admonition] + +[project.markdown_extensions.footnotes] + +[project.markdown_extensions.tables] + +[project.markdown_extensions.mdx_math] + +[project.markdown_extensions."pymdownx.details"] + +[project.markdown_extensions."pymdownx.superfences"]