torch-reconstruct-tomogram 0.6.0__tar.gz

torch_reconstruct_tomogram-0.6.0/.gitignore

@@ -0,0 +1,115 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+.DS_Store
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Files downloaded for unit tests
+**/tests/tmp/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# dotenv
+.env
+
+# virtualenv
+.venv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# ruff
+.ruff_cache/
+
+# IDEs
+.idea/
+.vscode/
+
+# Mojo extension compile cache (experimental torch-fourier-slice kernels)
+__mojocache__/
+*.mojopkg
+# experimental demo outputs
+packages/primitives/torch-fourier-slice/examples/*.npz
+packages/primitives/torch-fourier-slice/examples/*.png
+packages/primitives/torch-fourier-slice/examples/*.gif
+# Personal local notes (not for commit)
+notes/*.local.md
+
+lightning_logs/

torch_reconstruct_tomogram-0.6.0/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2020, TeamTomo
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

torch_reconstruct_tomogram-0.6.0/PKG-INFO

@@ -0,0 +1,71 @@
+Metadata-Version: 2.5
+Name: torch-reconstruct-tomogram
+Version: 0.6.0
+Summary: (sub-)tomogram reconstruction and subtilt extraction for cryoET.
+Project-URL: homepage, https://github.com/teamtomo/teamtomo
+Project-URL: repository, https://github.com/teamtomo/teamtomo
+Project-URL: documentation, https://github.com/teamtomo/teamtomo#readme
+Project-URL: changelog, https://github.com/teamtomo/teamtomo#changelog
+Project-URL: Bug Tracker, https://github.com/teamtomo/teamtomo/issues
+Project-URL: Source Code, https://github.com/teamtomo/teamtomo
+Author-email: Marten Chaillet <martenchaillet@gmail.com>, Davide Torre <davidetorre99@gmail.com>
+License: BSD-3-Clause
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: einops
+Requires-Dist: torch
+Requires-Dist: torch-ctf
+Requires-Dist: torch-fourier-rescale
+Requires-Dist: torch-fourier-slice
+Requires-Dist: torch-grid-utils>=0.0.8
+Requires-Dist: torch-subpixel-crop
+Requires-Dist: torch-tilt-series[io]
+Description-Content-Type: text/markdown
+
+# torch-reconstruct-tomogram
+
+[![License](https://img.shields.io/pypi/l/torch-reconstruct-tomogram.svg?color=green)](https://github.com/teamtomo/torch-reconstruct-tomogram/raw/main/LICENSE)
+[![PyPI](https://img.shields.io/pypi/v/torch-reconstruct-tomogram.svg?color=green)](https://pypi.org/project/torch-reconstruct-tomogram)
+[![Python Version](https://img.shields.io/pypi/pyversions/torch-reconstruct-tomogram.svg?color=green)](https://python.org)
+[![CI](https://github.com/teamtomo/torch-reconstruct-tomogram/actions/workflows/ci.yml/badge.svg)](https://github.com/teamtomo/torch-reconstruct-tomogram/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/teamtomo/torch-reconstruct-tomogram/branch/main/graph/badge.svg)](https://codecov.io/gh/teamtomo/torch-reconstruct-tomogram)
+
+(sub-)Tomogram reconstruction and subtilt extraction for cryo-ET.
+
+## Overview
+
+This package provides (sub-)tomogram reconstruction and subtilt extraction driven entirely from a [`torch-tilt-series`](https://github.com/teamtomo/torch-tilt-series) `TiltSeries`. It supports
+
+* `extract_particle_tilt_series()`: extract a subtilt-series at 3D location(s) in the sample
+* `reconstruct_subvolume()`: rank-polymorphic reconstruction of 3D patch(es) at location(s) in the sample
+* `reconstruct_tomogram()`: full volume reconstruction by tiling reconstructed patches in 3D
+
+`TiltSeries` holds alignment geometry (in Angstroms) plus `image_path`/`image_indices`/`pixel_spacing` metadata. The functions above take a `TiltSeries`, and load and (by default) preprocess the matching raw images internally via `torch_tilt_series.load_tilt_series_images()` / `preprocess_tilt_series_images()` (by default: plane subtraction, a DC-excluding bandpass with no low-pass, i.e. up to Nyquist, and central-crop normalization); pass `**preprocessing_kwargs` to override any of its defaults (`low=`, `high=`, `falloff=`, `bandpass_padding=`, `subtract_background=`, `normalize=`) - see `preprocess_tilt_series_images()`'s docstring for details.  `output_pixel_spacing` lets both local (`reconstruct_subvolume`) and global (`reconstruct_tomogram`) reconstruction target an arbitrary output voxel size. Reconstruction happens at the input pixel spacing and is Fourier-rescaled to the requested output size. Reconstruction is performed in Fourier space using central slice insertion. Positions are in `zyx` coordinates, in Angstroms, relative to the tomogram center.
+
+## Installation
+
+```bash
+pip install torch-reconstruct-tomogram
+```
+
+To load a tilt series from AreTomo or ETOMO output, also install the IO dependencies for [`torch-tilt-series`](https://github.com/teamtomo/torch-tilt-series):
+
+```bash
+pip install torch-tilt-series[io]
+```
+
+## Examples
+
+See the [`examples/`](examples/) folder for scripts showing how to load a tilt series, reconstruct subvolumes and tomograms, and save the result.
+
+## License
+
+This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.

torch_reconstruct_tomogram-0.6.0/README.md

@@ -0,0 +1,39 @@
+# torch-reconstruct-tomogram
+
+[![License](https://img.shields.io/pypi/l/torch-reconstruct-tomogram.svg?color=green)](https://github.com/teamtomo/torch-reconstruct-tomogram/raw/main/LICENSE)
+[![PyPI](https://img.shields.io/pypi/v/torch-reconstruct-tomogram.svg?color=green)](https://pypi.org/project/torch-reconstruct-tomogram)
+[![Python Version](https://img.shields.io/pypi/pyversions/torch-reconstruct-tomogram.svg?color=green)](https://python.org)
+[![CI](https://github.com/teamtomo/torch-reconstruct-tomogram/actions/workflows/ci.yml/badge.svg)](https://github.com/teamtomo/torch-reconstruct-tomogram/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/teamtomo/torch-reconstruct-tomogram/branch/main/graph/badge.svg)](https://codecov.io/gh/teamtomo/torch-reconstruct-tomogram)
+
+(sub-)Tomogram reconstruction and subtilt extraction for cryo-ET.
+
+## Overview
+
+This package provides (sub-)tomogram reconstruction and subtilt extraction driven entirely from a [`torch-tilt-series`](https://github.com/teamtomo/torch-tilt-series) `TiltSeries`. It supports
+
+* `extract_particle_tilt_series()`: extract a subtilt-series at 3D location(s) in the sample
+* `reconstruct_subvolume()`: rank-polymorphic reconstruction of 3D patch(es) at location(s) in the sample
+* `reconstruct_tomogram()`: full volume reconstruction by tiling reconstructed patches in 3D
+
+`TiltSeries` holds alignment geometry (in Angstroms) plus `image_path`/`image_indices`/`pixel_spacing` metadata. The functions above take a `TiltSeries`, and load and (by default) preprocess the matching raw images internally via `torch_tilt_series.load_tilt_series_images()` / `preprocess_tilt_series_images()` (by default: plane subtraction, a DC-excluding bandpass with no low-pass, i.e. up to Nyquist, and central-crop normalization); pass `**preprocessing_kwargs` to override any of its defaults (`low=`, `high=`, `falloff=`, `bandpass_padding=`, `subtract_background=`, `normalize=`) - see `preprocess_tilt_series_images()`'s docstring for details.  `output_pixel_spacing` lets both local (`reconstruct_subvolume`) and global (`reconstruct_tomogram`) reconstruction target an arbitrary output voxel size. Reconstruction happens at the input pixel spacing and is Fourier-rescaled to the requested output size. Reconstruction is performed in Fourier space using central slice insertion. Positions are in `zyx` coordinates, in Angstroms, relative to the tomogram center.
+
+## Installation
+
+```bash
+pip install torch-reconstruct-tomogram
+```
+
+To load a tilt series from AreTomo or ETOMO output, also install the IO dependencies for [`torch-tilt-series`](https://github.com/teamtomo/torch-tilt-series):
+
+```bash
+pip install torch-tilt-series[io]
+```
+
+## Examples
+
+See the [`examples/`](examples/) folder for scripts showing how to load a tilt series, reconstruct subvolumes and tomograms, and save the result.
+
+## License
+
+This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.

torch_reconstruct_tomogram-0.6.0/pyproject.toml

@@ -0,0 +1,195 @@
+# https://peps.python.org/pep-0517/
+[build-system]
+requires = ["hatchling", "hatch-vcs"]
+build-backend = "hatchling.build"
+
+# https://hatch.pypa.io/latest/config/metadata/
+[tool.hatch.version]
+source = "vcs"
+tag-pattern = "^torch-reconstruct-tomogram@v(?P<version>.+)$"
+fallback-version = "0.5.0"
+
+[tool.hatch.version.raw-options]
+search_parent_directories = true
+tag_regex = "^torch-reconstruct-tomogram@v(?P<version>\\d+\\.\\d+\\.\\d+.*)$"
+# Constrain git-describe so it only considers TeamTomo's own tags, not other workspace tags.
+# See https://github.com/ofek/hatch-vcs/issues/71
+git_describe_command = "git describe --dirty --tags --long --match 'torch-reconstruct-tomogram@v[0-9]*.[0-9]*.[0-9]*'"
+
+# read more about configuring hatch at:
+# https://hatch.pypa.io/latest/config/build/
+[tool.hatch.build.targets.wheel]
+only-include = ["src"]
+sources = ["src"]
+
+# https://peps.python.org/pep-0621/
+[project]
+name = "torch-reconstruct-tomogram"
+dynamic = ["version"]
+description = "(sub-)tomogram reconstruction and subtilt extraction for cryoET."
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "BSD-3-Clause" }
+authors = [
+    { name = "Marten Chaillet", email = "martenchaillet@gmail.com" },
+    { name = "Davide Torre", email = "davidetorre99@gmail.com" },
+]
+# https://pypi.org/classifiers/
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: BSD License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Typing :: Typed",
+]
+# add your package dependencies here
+dependencies = [
+    "torch",
+    "einops",
+    "torch-fourier-slice",
+    "torch-fourier-rescale",
+    "torch-ctf",
+    "torch-grid-utils>=0.0.8",
+    "torch-subpixel-crop",
+    "torch-tilt-series[io]",
+]
+
+# https://peps.python.org/pep-0621/#dependencies-optional-dependencies
+# add dependencies for "extra" features here.  Not dev dependencies.
+# [project.optional-dependencies]
+# name = ["dependency"]
+
+[project.urls]
+homepage = "https://github.com/teamtomo/teamtomo"
+repository = "https://github.com/teamtomo/teamtomo"
+documentation = "https://github.com/teamtomo/teamtomo#readme"
+changelog = "https://github.com/teamtomo/teamtomo#changelog"
+"Bug Tracker" = "https://github.com/teamtomo/teamtomo/issues"
+"Source Code" = "https://github.com/teamtomo/teamtomo"
+# Entry points
+# https://peps.python.org/pep-0621/#entry-points
+# same as console_scripts entry point
+# [project.scripts]
+# torch-reconstruct-tomogram-cli = "torch_reconstruct_tomogram:main_cli"
+
+# [project.entry-points."some.group"]
+# tomatoes = "torch_reconstruct_tomogram:main_tomatoes"
+
+# https://peps.python.org/pep-0735/
+# setup with `uv sync` or `pip install -e . --group dev`
+[dependency-groups]
+test = ["pytest", "pytest-cov"]
+dev = [
+    { include-group = "test" },
+    "ipython",
+    "mypy",
+    "pdbpp",  # https://github.com/pdbpp/pdbpp
+    "pre-commit-uv",
+    "rich",  # https://github.com/Textualize/rich
+    "ruff",
+]
+
+[tool.uv.sources]
+torch-tilt-series = { workspace = true }
+
+# https://docs.astral.sh/ruff
+[tool.ruff]
+line-length = 88
+target-version = "py311"
+src = ["src"]
+fix = true
+# unsafe-fixes = true
+
+[tool.ruff.lint]
+pydocstyle = { convention = "numpy" }
+select = [
+    "E",    # style errors
+    "W",    # style warnings
+    "F",    # flakes
+    "D",    # pydocstyle
+    "D417", # Missing argument descriptions in Docstrings
+    "I",    # isort
+    "UP",   # pyupgrade
+    "C4",   # flake8-comprehensions
+    "B",    # flake8-bugbear
+    "A001", # flake8-builtins
+    "RUF",  # ruff-specific rules
+    "TC",   # flake8-type-checking
+    "TID",  # flake8-tidy-imports
+]
+ignore = [
+    "D401", # First line should be in imperative mood (remove to opt in)
+]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*.py" = ["D", "S"]
+"docs/*.py" = ["D", "A"]
+
+# https://docs.astral.sh/ruff/formatter/
+[tool.ruff.format]
+docstring-code-format = true
+skip-magic-trailing-comma = false  # default is false
+
+# https://docs.pytest.org/
+[tool.pytest.ini_options]
+minversion = "7.0"
+addopts = ["--color=yes"]
+testpaths = ["tests"]
+filterwarnings = [
+    "error",
+    # mrcfile 1.5.x triggers this when reading MRC headers with NumPy >= 2.5
+    "ignore:Setting the dtype on a NumPy array has been deprecated:DeprecationWarning",
+]
+
+# https://mypy.readthedocs.io/en/stable/config_file.html
+[tool.mypy]
+files = "src/**/"
+strict = true
+disallow_any_generics = false
+disallow_subclassing_any = false
+show_error_codes = true
+pretty = true
+# plugins = ["pydantic.mypy"]
+
+# # module specific overrides
+# [[tool.mypy.overrides]]
+# module = ["numpy.*",]
+# ignore_errors = true
+
+# https://coverage.readthedocs.io/
+[tool.coverage.report]
+show_missing = true
+exclude_lines = [
+    "pragma: no cover",
+    "if TYPE_CHECKING:",
+    "@overload",
+    "except ImportError",
+    "raise AssertionError",
+    "\\.\\.\\.",
+    "raise NotImplementedError()",
+    "pass",
+]
+
+[tool.coverage.run]
+source = ["torch_reconstruct_tomogram"]
+
+# https://github.com/mgedmin/check-manifest#configuration
+# add files that you want check-manifest to explicitly ignore here
+# (files that are in the repo but shouldn't go in the package)
+[tool.check-manifest]
+ignore = [
+    ".pre-commit-config.yaml",
+    ".ruff_cache/**/*",
+    "tests/**/*",
+    "uv.lock",
+]
+
+# https://github.com/crate-ci/typos/blob/master/docs/reference.md
+[tool.typos.default]
+extend-ignore-identifiers-re = []
+
+[tool.typos.default.extend-words]
+writeable = "writeable"

torch_reconstruct_tomogram-0.6.0/src/torch_reconstruct_tomogram/__init__.py

@@ -0,0 +1,26 @@
+"""(sub-)tomogram reconstruction and subtilt extraction for cryoET."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("torch-reconstruct-tomogram")
+except PackageNotFoundError:
+    __version__ = "uninstalled"
+__author__ = "Marten Chaillet, Davide Torre"
+__email__ = "martenchaillet@gmail.com, davidetorre99@gmail.com"
+
+from torch_reconstruct_tomogram.projection import (
+    extract_particle_tilt_series,
+    project_points,
+)
+from torch_reconstruct_tomogram.reconstruct import (
+    reconstruct_subvolume,
+    reconstruct_tomogram,
+)
+
+__all__ = [
+    "extract_particle_tilt_series",
+    "project_points",
+    "reconstruct_subvolume",
+    "reconstruct_tomogram",
+]

torch_reconstruct_tomogram-0.6.0/src/torch_reconstruct_tomogram/projection.py

@@ -0,0 +1,77 @@
+"""Project 3D points into tilt images and crop patches for reconstruction."""
+
+from typing import Any
+
+import torch
+from torch_grid_utils import dft_center
+from torch_subpixel_crop import subpixel_crop_2d
+from torch_tilt_series import (
+    TiltSeries,
+    load_tilt_series_images,
+    preprocess_tilt_series_images,
+)
+
+
+def project_points(tilt_series: TiltSeries, points_zyx: torch.Tensor) -> torch.Tensor:
+    """Project 3D points to 2D image pixel coordinates.
+
+    - points are 3D zyx coordinates in Angstroms, relative to the tomogram center
+    - tilt_series supplies the projection geometry (`tilt_series.project_points`
+      works in Angstroms) and `tilt_series.pixel_spacing` (raises if unset),
+      used to convert the projected Angstrom positions to pixels
+    - projected 2D points are in pixels, relative to the center of each image
+    """
+    return tilt_series.project_points(points_zyx) / tilt_series.pixel_spacing
+
+
+def _extract_particle_tilt_series(
+    tilt_series: TiltSeries,
+    images: torch.Tensor,
+    points_zyx: torch.Tensor,
+    sidelength: int,
+    return_rfft: bool = True,
+) -> torch.Tensor:
+    """Extract a subtilt-series given already-loaded images."""
+    projected_yx = project_points(tilt_series, points_zyx)
+    projected_yx = projected_yx + dft_center(
+        images.shape[-2:], rfft=False, fftshift=True, device=images.device
+    )
+    return subpixel_crop_2d(
+        image=images,
+        positions=projected_yx,
+        sidelength=sidelength,
+        return_rfft=return_rfft,
+        decenter=return_rfft,
+    )
+
+
+def extract_particle_tilt_series(
+    tilt_series: TiltSeries,
+    points_zyx: torch.Tensor,
+    sidelength: int,
+    return_rfft: bool = True,
+    preprocess: bool = True,
+    **preprocessing_kwargs: Any,
+) -> torch.Tensor:
+    """Extract a subtilt-series at 3D location(s) in the sample.
+
+    Loads (and, by default, preprocesses) the raw tilt images matching
+    `tilt_series` via `tilt_series.image_path`/`image_indices`. Preprocessing
+    (see `torch_tilt_series.preprocess_tilt_series_images`) by default applies
+    plane subtraction, a DC-excluding bandpass with no low-pass (i.e. up to
+    Nyquist), and central-crop normalization. `**preprocessing_kwargs` are
+    forwarded to `preprocess_tilt_series_images`, overriding any of its
+    defaults (`low`, `high`, `falloff`, `bandpass_padding`,
+    `subtract_background`, `normalize`) - see that function's docstring for
+    details.
+    """
+    images = load_tilt_series_images(tilt_series)
+    if preprocess:
+        images = preprocess_tilt_series_images(images, **preprocessing_kwargs)
+    return _extract_particle_tilt_series(
+        tilt_series,
+        images,
+        points_zyx,
+        sidelength,
+        return_rfft=return_rfft,
+    )

torch_reconstruct_tomogram-0.6.0/src/torch_reconstruct_tomogram/py.typed

@@ -0,0 +1,5 @@
+You may remove this file if you don't intend to add types to your package
+
+Details at:
+
+https://mypy.readthedocs.io/en/stable/installed_packages.html#creating-pep-561-compatible-packages

torch_reconstruct_tomogram-0.6.0/src/torch_reconstruct_tomogram/reconstruct.py

@@ -0,0 +1,294 @@
+"""(sub-)tomogram reconstruction in pytorch."""
+
+from typing import Any
+
+import einops
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch_fourier_rescale import fourier_rescale_rfft_2d
+from torch_fourier_slice import insert_central_slices_rfft_3d_multichannel
+from torch_grid_utils import fftfreq_grid
+from torch_tilt_series import (
+    TiltSeries,
+    load_tilt_series_images,
+    preprocess_tilt_series_images,
+)
+
+from torch_reconstruct_tomogram.projection import _extract_particle_tilt_series
+
+_PAD_FACTOR = 2.0
+
+
+def _writable(data):
+    if isinstance(data, np.ndarray) and not data.flags.writeable:
+        data = data.copy()
+    return data
+
+
+def _reconstruct_subvolume(
+    tilt_series: TiltSeries,
+    images: torch.Tensor,
+    points_zyx: torch.Tensor,
+    sidelength: int,
+    output_pixel_spacing: float | None = None,
+) -> torch.Tensor:
+    """Reconstruct subvolume(s), given already-loaded images."""
+    device = images.device
+    input_pixel_spacing = tilt_series.pixel_spacing  # raises if unset
+    if output_pixel_spacing is None:
+        output_pixel_spacing = input_pixel_spacing
+
+    points_zyx = torch.as_tensor(_writable(points_zyx), device=device).float()
+    points_zyx, ps = einops.pack([points_zyx], "* zyx")
+
+    # tomogram -> detector rotation: projection_matrices is sample -> detector
+    # only, so compose tomo2sample's rotation in first. Every patch's
+    # Fourier-insertion rotation must be expressed relative to the
+    # tomogram frame.
+    rotation_matrices = (
+        tilt_series.projection_matrices[:, :3, :3] @ tilt_series.tomo2sample[:3, :3]
+    )
+    rotation_matrices = torch.linalg.pinv(rotation_matrices)
+
+    sidelength_padded_output = int(_PAD_FACTOR * sidelength)
+    sidelength_padded_native = max(
+        1,
+        round(sidelength_padded_output * output_pixel_spacing / input_pixel_spacing),
+    )
+
+    particle_tilt_series_rfft = _extract_particle_tilt_series(
+        tilt_series,
+        images,
+        points_zyx,
+        sidelength=sidelength_padded_native,
+        return_rfft=True,
+    )
+
+    particle_tilt_series_rfft = torch.fft.fftshift(particle_tilt_series_rfft, dim=(-2,))
+
+    particle_tilt_series_rfft = fourier_rescale_rfft_2d(
+        dft=particle_tilt_series_rfft,
+        image_shape=(sidelength_padded_native, sidelength_padded_native),
+        target_shape=(sidelength_padded_output, sidelength_padded_output),
+    )
+
+    particle_tilt_series_rfft = einops.rearrange(
+        particle_tilt_series_rfft,
+        "n_positions n_tilts h w_rfft -> n_tilts n_positions h w_rfft",
+    )
+
+    patches_rfft, weights = insert_central_slices_rfft_3d_multichannel(
+        image_rfft=particle_tilt_series_rfft,
+        volume_shape=(sidelength_padded_output,) * 3,
+        rotation_matrices=rotation_matrices,
+        zyx_matrices=True,
+        fftfreq_max=0.5,
+    )
+
+    valid_weights = weights > 1e-3
+    patches_rfft[:, valid_weights] /= weights[valid_weights]
+
+    patches_rfft = torch.fft.ifftshift(patches_rfft, dim=(-3, -2))
+
+    patches = torch.fft.irfftn(
+        patches_rfft,
+        s=(sidelength_padded_output,) * 3,
+        dim=(-3, -2, -1),
+    )
+
+    patches = torch.fft.ifftshift(patches, dim=(-3, -2, -1))
+
+    grid = fftfreq_grid(
+        image_shape=(sidelength_padded_output,) * 3,
+        rfft=False,
+        fftshift=True,
+        norm=True,
+        device=device,
+    )
+    patches = patches / torch.sinc(grid) ** 2
+
+    p = (sidelength_padded_output - sidelength) // 2
+    patches = F.pad(patches, [-p] * 6)
+
+    [patches] = einops.unpack(patches, ps, "* d h w")
+
+    return patches
+
+
+def reconstruct_subvolume(
+    tilt_series: TiltSeries,
+    points_zyx: torch.Tensor,
+    sidelength: int,
+    output_pixel_spacing: float | None = None,
+    preprocess: bool = True,
+    **preprocessing_kwargs: Any,
+) -> torch.Tensor:
+    """Reconstruct 3D patch(es) at location(s) in the sample.
+
+    Rank-polymorphic: input (..., 3) -> output (..., d, h, w)
+
+    - tilt_series supplies the projection geometry,
+    - points_zyx are zyx coordinates in Angstroms, relative to the tomogram center
+    - sidelength is the output subvolume size in voxels
+    - output_pixel_spacing is the voxel size of the output in Angstroms
+      (defaults to `tilt_series.pixel_spacing`); the per-tilt 2D crops are
+      Fourier-rescaled to this pixel size before 3D reconstruction, so local
+      (subvolume) and global (tomogram) reconstructions can each target an
+      arbitrary output pixel size independent of the raw data's
+    - preprocess, if True (default), applies
+      `torch_tilt_series.preprocess_tilt_series_images` to the loaded images
+      before reconstruction - by default plane subtraction, a DC-excluding
+      bandpass with no low-pass, i.e. up to Nyquist, and central-crop
+      normalization
+    - `**preprocessing_kwargs` are forwarded to `preprocess_tilt_series_images`,
+      overriding any of its defaults (`low`, `high`, `falloff`,
+      `bandpass_padding`, `subtract_background`, `normalize`) - see that
+      function's docstring for details
+    """
+    images = load_tilt_series_images(tilt_series)
+    if preprocess:
+        images = preprocess_tilt_series_images(images, **preprocessing_kwargs)
+    return _reconstruct_subvolume(
+        tilt_series,
+        images,
+        points_zyx,
+        sidelength,
+        output_pixel_spacing=output_pixel_spacing,
+    )
+
+
+def _cosine_taper_window(core_length: int, margin: int, device) -> torch.Tensor:
+    """1D cosine-taper window, flat in the middle, tapered at the edges.
+
+    1.0 over the central `core_length` samples, cosine-tapered from 0 up to 1
+    (and back down to 0) over `margin` samples on each side. Total length is
+    core_length + 2 * margin.
+    """
+    if margin == 0:
+        return torch.ones(core_length, device=device)
+    ramp = 0.5 * (1 - torch.cos(torch.linspace(0, torch.pi, margin, device=device)))
+    core = torch.ones(core_length, device=device)
+    return torch.cat([ramp, core, ramp.flip(0)])
+
+
+def reconstruct_tomogram(
+    tilt_series: TiltSeries,
+    volume_shape: tuple[int, int, int],
+    sidelength: int,
+    batch_size: int | None = None,
+    output_pixel_spacing: float | None = None,
+    preprocess: bool = True,
+    blend_margin: int | None = None,
+    **preprocessing_kwargs: Any,
+) -> torch.Tensor:
+    """Reconstruct the full tomogram by tiling reconstructed patches in 3D.
+
+    - tilt_series supplies the projection geometry, plus `image_path`/
+      `image_indices` used to load the raw images
+    - volume_shape is the (d, h, w) shape of the output tomogram, in voxels
+    - sidelength is the spacing between patch centers, in voxels; patches
+      are reconstructed on a grid tiling `volume_shape`
+    - batch_size, if set, reconstructs at most this many patches per chunk
+      (to bound memory usage); defaults to reconstructing all patches at once
+    - output_pixel_spacing is the voxel size of the output in Angstroms
+      (defaults to `tilt_series.pixel_spacing`)
+    - preprocess, if True (default), applies
+      `torch_tilt_series.preprocess_tilt_series_images` to the loaded images
+      before reconstruction - by default plane subtraction, a DC-excluding
+      bandpass with no low-pass, i.e. up to Nyquist, and central-crop
+      normalization
+    - blend_margin is the extra margin, in voxels, added around each patch
+      (total reconstructed patch size is `sidelength + 2 * blend_margin`);
+      overlapping patches are cosine-tapered and blended together over this
+      margin to avoid seams at patch boundaries. Defaults to
+      `sidelength // 4`
+    - `**preprocessing_kwargs` are forwarded to `preprocess_tilt_series_images`,
+      overriding any of its defaults (`low`, `high`, `falloff`,
+      `bandpass_padding`, `subtract_background`, `normalize`) - see that
+      function's docstring for details
+    """
+    images = load_tilt_series_images(tilt_series)
+    if preprocess:
+        images = preprocess_tilt_series_images(images, **preprocessing_kwargs)
+
+    pixel_spacing = tilt_series.pixel_spacing  # raises if unset
+    if output_pixel_spacing is None:
+        output_pixel_spacing = pixel_spacing
+
+    if blend_margin is None:
+        blend_margin = sidelength // 4
+    patch_sidelength = sidelength + 2 * blend_margin
+    half = patch_sidelength // 2
+
+    d, h, w = volume_shape
+    r = sidelength // 2
+    device = images.device
+
+    z_centers = torch.arange(start=r, end=d + r, step=sidelength, device=device)
+    y_centers = torch.arange(start=r, end=h + r, step=sidelength, device=device)
+    x_centers = torch.arange(start=r, end=w + r, step=sidelength, device=device)
+    # absolute 0-indexed voxel coordinates of each patch center
+    centers_voxel = torch.stack(
+        torch.meshgrid(z_centers, y_centers, x_centers, indexing="ij"), dim=-1
+    )
+
+    volume_center = torch.tensor([d, h, w], device=device) // 2
+    centers_zyx_ang = (centers_voxel - volume_center) * output_pixel_spacing
+
+    window_1d = _cosine_taper_window(sidelength, blend_margin, device="cpu")
+    window_3d = (
+        window_1d[:, None, None] * window_1d[None, :, None] * window_1d[None, None, :]
+    )
+
+    tomogram_sum = torch.zeros(volume_shape, dtype=torch.float32)
+    weight_sum = torch.zeros(volume_shape, dtype=torch.float32)
+
+    centers_flat, _ = einops.pack([centers_voxel], "* zyx")
+    centers_ang_flat, _ = einops.pack([centers_zyx_ang], "* zyx")
+    chunk_size = batch_size or len(centers_flat)
+
+    for start in range(0, len(centers_flat), chunk_size):
+        chunk_centers = centers_flat[start : start + chunk_size]
+        chunk_centers_ang = centers_ang_flat[start : start + chunk_size]
+
+        patches_batch = _reconstruct_subvolume(
+            tilt_series,
+            images,
+            chunk_centers_ang,
+            patch_sidelength,
+            output_pixel_spacing=output_pixel_spacing,
+        ).cpu()
+
+        for j in range(len(patches_batch)):
+            cz, cy, cx = chunk_centers[j].tolist()
+            z0, y0, x0 = cz - half, cy - half, cx - half
+            z1, y1, x1 = (
+                z0 + patch_sidelength,
+                y0 + patch_sidelength,
+                x0 + patch_sidelength,
+            )
+
+            # clip the patch's placement to the volume bounds
+            cz0, cy0, cx0 = max(z0, 0), max(y0, 0), max(x0, 0)
+            cz1, cy1, cx1 = min(z1, d), min(y1, h), min(x1, w)
+            if cz0 >= cz1 or cy0 >= cy1 or cx0 >= cx1:
+                continue
+
+            src = (
+                slice(cz0 - z0, cz1 - z0),
+                slice(cy0 - y0, cy1 - y0),
+                slice(cx0 - x0, cx1 - x0),
+            )
+            dst = (slice(cz0, cz1), slice(cy0, cy1), slice(cx0, cx1))
+            weight_block = window_3d[src]
+            tomogram_sum[dst] += patches_batch[j][src] * weight_block
+            weight_sum[dst] += weight_block
+
+        del patches_batch
+        if device.type != "cpu":
+            torch.cuda.empty_cache()
+
+    tomogram = tomogram_sum / weight_sum.clamp_min(1e-6)
+
+    return tomogram.to(device)

torch_reconstruct_tomogram-0.6.0/tests/test_torch_reconstruct_tomogram.py

@@ -0,0 +1,242 @@
+import numpy as np
+import pytest
+import torch
+from torch_tilt_series import TiltSeries
+
+import torch_reconstruct_tomogram
+from torch_reconstruct_tomogram import (
+    project_points,
+    reconstruct_subvolume,
+    reconstruct_tomogram,
+)
+
+DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
+
+
+def make_tilt_series(tmp_path, device="cpu", size=32):
+    mrcfile = pytest.importorskip("mrcfile")
+    tilt_angles = torch.tensor([-30.0, 0.0, 30.0])
+    images = np.zeros((3, size, size), dtype=np.float32)
+    c = size // 2
+    images[:, c - 2 : c + 2, c - 2 : c + 2] = 1.0
+    image_path = tmp_path / "images.mrc"
+    mrcfile.write(image_path, images, overwrite=True)
+    return TiltSeries(
+        tilt_angles=tilt_angles,
+        tilt_axis_angle=torch.tensor(0.0),
+        sample_translations=torch.zeros((3, 2)),
+        image_path=image_path,
+        pixel_spacing=1.0,
+        device=device,
+    )
+
+
+def test_imports_with_version():
+    assert isinstance(torch_reconstruct_tomogram.__version__, str)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_subvolume(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    point_zyx = torch.tensor([0.0, 0.0, 0.0], device=device)
+    subvolume = reconstruct_subvolume(tilt_series, point_zyx, sidelength=8)
+    assert subvolume.shape == (8, 8, 8)
+    assert subvolume.dtype == torch.float32
+    assert device in str(subvolume.device)
+    assert torch.isfinite(subvolume).all()
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_subvolume_rank_polymorphic(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+
+    # single point (3,) -> (d, h, w)
+    point = torch.tensor([0.0, 0.0, 0.0], device=device)
+    assert reconstruct_subvolume(tilt_series, point, sidelength=8).shape == (8, 8, 8)
+
+    # batch (N, 3) -> (N, d, h, w)
+    points = torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], device=device)
+    assert reconstruct_subvolume(tilt_series, points, sidelength=8).shape == (
+        2,
+        8,
+        8,
+        8,
+    )
+
+    # 2D grid (a, b, 3) -> (a, b, d, h, w)
+    grid_2d = torch.zeros(2, 3, 3, device=device)
+    assert reconstruct_subvolume(tilt_series, grid_2d, sidelength=8).shape == (
+        2,
+        3,
+        8,
+        8,
+        8,
+    )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_subvolume_output_pixel_spacing(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    point = torch.tensor([0.0, 0.0, 0.0], device=device)
+
+    subvolume_default = reconstruct_subvolume(tilt_series, point, sidelength=8)
+    subvolume_explicit = reconstruct_subvolume(
+        tilt_series, point, sidelength=8, output_pixel_spacing=1.0
+    )
+    assert torch.allclose(subvolume_default, subvolume_explicit)
+
+    subvolume_coarse = reconstruct_subvolume(
+        tilt_series, point, sidelength=8, output_pixel_spacing=2.0
+    )
+    assert subvolume_coarse.shape == (8, 8, 8)
+    assert torch.isfinite(subvolume_coarse).all()
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_subvolume_local_shifts(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    point = torch.tensor([0.0, 0.0, 0.0], device=device)
+
+    def zero_local_shifts(projected_yx):
+        return torch.zeros_like(projected_yx)
+
+    subvolume_default = reconstruct_subvolume(tilt_series, point, sidelength=8)
+    tilt_series.local_shifts_2d = zero_local_shifts
+    subvolume_with_hook = reconstruct_subvolume(tilt_series, point, sidelength=8)
+    assert torch.allclose(subvolume_default, subvolume_with_hook)
+
+
+def test_project_points_local_shifts_are_angstroms_not_pixels():
+    # pixel_spacing != 1 so an Angstrom-space shift and a pixel-space shift
+    # would disagree if local_shifts were (still) being applied in pixels.
+    shift_ang = torch.tensor([5.0, -3.0])
+
+    def shift_fn(projected_yx_ang):
+        return shift_ang.expand_as(projected_yx_ang)
+
+    tilt_series_shifted = TiltSeries(
+        tilt_angles=torch.tensor([0.0]),
+        tilt_axis_angle=torch.tensor(0.0),
+        sample_translations=torch.zeros((1, 2)),
+        pixel_spacing=2.0,
+        local_shifts_2d=shift_fn,
+    )
+    tilt_series = TiltSeries(
+        tilt_angles=torch.tensor([0.0]),
+        tilt_axis_angle=torch.tensor(0.0),
+        sample_translations=torch.zeros((1, 2)),
+        pixel_spacing=2.0,
+    )
+    point = torch.tensor([[0.0, 0.0, 0.0]])
+
+    shifted_px = project_points(tilt_series_shifted, point)
+    unshifted_px = project_points(tilt_series, point)
+    # shift is applied in Angstroms, then the whole result is divided by
+    # pixel_spacing so the pixel-space delta is shift_ang / pixel_spacing,
+    # not shift_ang itself.
+    expected_delta_px = shift_ang / tilt_series.pixel_spacing
+    assert torch.allclose(
+        (shifted_px - unshifted_px)[0, 0], expected_delta_px, atol=1e-5
+    )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_subvolume_preprocess_toggle(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    point = torch.tensor([0.0, 0.0, 0.0], device=device)
+
+    preprocessed = reconstruct_subvolume(
+        tilt_series, point, sidelength=8, preprocess=True
+    )
+    raw = reconstruct_subvolume(tilt_series, point, sidelength=8, preprocess=False)
+    assert not torch.allclose(preprocessed, raw)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_subvolume_preprocessing_kwargs_are_forwarded(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    point = torch.tensor([0.0, 0.0, 0.0], device=device)
+
+    default = reconstruct_subvolume(tilt_series, point, sidelength=8)
+    custom = reconstruct_subvolume(
+        tilt_series, point, sidelength=8, high=0.3, bandpass_padding=4
+    )
+    assert not torch.allclose(default, custom)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_tomogram_output_pixel_spacing(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    volume = reconstruct_tomogram(
+        tilt_series, (16, 16, 16), sidelength=8, output_pixel_spacing=2.0
+    )
+    assert volume.shape == (16, 16, 16)
+    assert torch.isfinite(volume).all()
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_tomogram(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    volume = reconstruct_tomogram(tilt_series, (16, 16, 16), sidelength=8)
+    assert volume.shape == (16, 16, 16)
+    assert volume.dtype == torch.float32
+    assert device in str(volume.device)
+    assert torch.isfinite(volume).all()
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_tomogram_non_cubic(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    # shape not divisible by sidelength is still cropped to the requested shape
+    volume = reconstruct_tomogram(tilt_series, (8, 24, 20), sidelength=8)
+    assert volume.shape == (8, 24, 20)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_reconstruct_tomogram_batch_size(device, tmp_path):
+    tilt_series = make_tilt_series(tmp_path, device)
+    recon_no_batch = reconstruct_tomogram(tilt_series, (16, 16, 16), sidelength=8)
+    recon_with_batch = reconstruct_tomogram(
+        tilt_series, (16, 16, 16), sidelength=8, batch_size=2
+    )
+    assert recon_no_batch.shape == (16, 16, 16)
+    assert recon_with_batch.shape == (16, 16, 16)
+    diff = torch.abs(recon_no_batch - recon_with_batch.to(recon_no_batch.device)).max()
+    assert diff == 0.0
+
+
+def test_reconstruct_subvolume_rotation_includes_tomo2sample(tmp_path, monkeypatch):
+    """The Fourier-insertion rotation must be the tomogram -> detector
+    rotation (projection_matrices' rotation composed with tomo2sample's),
+    not just projection_matrices' sample -> detector rotation alone.
+    Otherwise every patch is reconstructed oriented to sample space while
+    being tiled into an axis-aligned tomogram-space grid, tearing apart at
+    patch boundaries whenever sample2tomo carries a rotation. No-op when
+    sample2tomo is the default identity, which is why this needs its own
+    explicit test rather than relying on the (identity-only) tests above.
+    """
+    from torch_affine_utils.transforms_3d import Rz
+
+    import torch_reconstruct_tomogram.reconstruct as reconstruct_module
+
+    tilt_series = make_tilt_series(tmp_path)
+    tilt_series.levelled2tomo = Rz(torch.tensor(37.0), zyx=True, device="cpu")
+
+    captured = {}
+    real_insert = reconstruct_module.insert_central_slices_rfft_3d_multichannel
+
+    def spy(*args, **kwargs):
+        captured["rotation_matrices"] = kwargs["rotation_matrices"].clone()
+        return real_insert(*args, **kwargs)
+
+    monkeypatch.setattr(
+        reconstruct_module, "insert_central_slices_rfft_3d_multichannel", spy
+    )
+
+    reconstruct_subvolume(tilt_series, torch.tensor([0.0, 0.0, 0.0]), sidelength=8)
+
+    expected_forward = (
+        tilt_series.projection_matrices[:, :3, :3] @ tilt_series.tomo2sample[:3, :3]
+    )
+    expected = torch.linalg.pinv(expected_forward)
+    assert torch.allclose(captured["rotation_matrices"], expected, atol=1e-5)