torch-2dtm 0.5.0__tar.gz

torch_2dtm-0.5.0/.gitignore

@@ -0,0 +1,105 @@
+# 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/
+
+lightning_logs/

torch_2dtm-0.5.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_2dtm-0.5.0/PKG-INFO

@@ -0,0 +1,96 @@
+Metadata-Version: 2.4
+Name: torch-2dtm
+Version: 0.5.0
+Summary: 2D template matching in pytorch
+Project-URL: homepage, https://github.com/teamtomo/teamtomo
+Project-URL: repository, https://github.com/teamtomo/teamtomo
+Author-email: Josh Dickerson <jdickerson@berkeley.edu>, Matthew Giammar <matthew_giammar@berkeley.edu>, Alister Burt <alisterburt@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: setuptools
+Requires-Dist: torch
+Requires-Dist: torch-fourier-slice
+Description-Content-Type: text/markdown
+
+# torch-2dtm
+
+[![License](https://img.shields.io/pypi/l/torch-2dtm.svg?color=green)](https://github.com/teamtomo/torch-2dtm/raw/main/LICENSE)
+[![PyPI](https://img.shields.io/pypi/v/torch-2dtm.svg?color=green)](https://pypi.org/project/torch-2dtm)
+[![Python Version](https://img.shields.io/pypi/pyversions/torch-2dtm.svg?color=green)](https://python.org)
+[![CI](https://github.com/teamtomo/torch-2dtm/actions/workflows/ci.yml/badge.svg)](https://github.com/teamtomo/torch-2dtm/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/teamtomo/torch-2dtm/branch/main/graph/badge.svg)](https://codecov.io/gh/teamtomo/torch-2dtm)
+
+## Overview
+
+torch-2dtm is a Python package for efficient templating matching of 
+2D projections of a 3D template with a 2D image in PyTorch.
+
+This is implemented for cryo-EM applications, see 
+[Rickgauer et al. 2017 eLife](https://doi.org/10.7554/eLife.25648) for details.
+
+## Features
+
+- Fast 2D template matching using Fourier transforms
+- Batch processing over orientations
+- Batch processing over Fourier space filters (e.g. for defocus sweeps)
+- GPU acceleration through PyTorch
+
+Projections are calculated on-the-fly using 
+[*torch-fourier-slice*](https://github.com/teamtomo/torch-fourier-slice).
+
+## Installation
+
+```bash
+pip install torch-2dtm
+```
+
+## Basic Usage
+
+```python
+import torch
+import torch_2dtm
+from scipy.stats import special_ortho_group
+
+# Create random test data
+# 1. Create a random image and compute its FFT
+image_size = (128, 128)
+image = torch.randn(*image_size, dtype=torch.float32)
+image_dft = torch.fft.rfftn(image, dim=(0, 1))  # Shape: (128, 65)
+
+# 2. Create a random 3D template and compute its FFT
+template_size = (64, 64, 64)
+template = torch.randn(*template_size, dtype=torch.float32)
+template_dft = torch.fft.rfftn(template, dim=(0, 1, 2))  # Shape: (64, 64, 33)
+
+# 3. Create a batch of random rotation matrices with shape (b, 3, 3)
+num_orientations = 10
+rotation_matrices = torch.tensor(special_ortho_group.rvs(size=num_orientations, dim=3), dtype=torch.float32)
+
+# 4. Create an arbitrary stack of Fourier space filters (identity filter in this example)
+# These filters operate on rffts of the 2D projection images
+# Filter shape: (..., h, w // 2 + 1)
+filters = torch.ones(template_size[0], template_size[1] // 2 + 1, dtype=torch.complex64)
+
+# Perform template matching
+cross_correlation = torch_2dtm.match_template_dft_2d(
+    image_dft=image_dft,
+    template_dft=template_dft,
+    rotation_matrices=rotation_matrices,
+    filters=filters
+)
+# The result has shape (..., num_orientations, image_height, image_width)
+```
+
+## License
+
+This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.

torch_2dtm-0.5.0/README.md

@@ -0,0 +1,72 @@
+# torch-2dtm
+
+[![License](https://img.shields.io/pypi/l/torch-2dtm.svg?color=green)](https://github.com/teamtomo/torch-2dtm/raw/main/LICENSE)
+[![PyPI](https://img.shields.io/pypi/v/torch-2dtm.svg?color=green)](https://pypi.org/project/torch-2dtm)
+[![Python Version](https://img.shields.io/pypi/pyversions/torch-2dtm.svg?color=green)](https://python.org)
+[![CI](https://github.com/teamtomo/torch-2dtm/actions/workflows/ci.yml/badge.svg)](https://github.com/teamtomo/torch-2dtm/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/teamtomo/torch-2dtm/branch/main/graph/badge.svg)](https://codecov.io/gh/teamtomo/torch-2dtm)
+
+## Overview
+
+torch-2dtm is a Python package for efficient templating matching of 
+2D projections of a 3D template with a 2D image in PyTorch.
+
+This is implemented for cryo-EM applications, see 
+[Rickgauer et al. 2017 eLife](https://doi.org/10.7554/eLife.25648) for details.
+
+## Features
+
+- Fast 2D template matching using Fourier transforms
+- Batch processing over orientations
+- Batch processing over Fourier space filters (e.g. for defocus sweeps)
+- GPU acceleration through PyTorch
+
+Projections are calculated on-the-fly using 
+[*torch-fourier-slice*](https://github.com/teamtomo/torch-fourier-slice).
+
+## Installation
+
+```bash
+pip install torch-2dtm
+```
+
+## Basic Usage
+
+```python
+import torch
+import torch_2dtm
+from scipy.stats import special_ortho_group
+
+# Create random test data
+# 1. Create a random image and compute its FFT
+image_size = (128, 128)
+image = torch.randn(*image_size, dtype=torch.float32)
+image_dft = torch.fft.rfftn(image, dim=(0, 1))  # Shape: (128, 65)
+
+# 2. Create a random 3D template and compute its FFT
+template_size = (64, 64, 64)
+template = torch.randn(*template_size, dtype=torch.float32)
+template_dft = torch.fft.rfftn(template, dim=(0, 1, 2))  # Shape: (64, 64, 33)
+
+# 3. Create a batch of random rotation matrices with shape (b, 3, 3)
+num_orientations = 10
+rotation_matrices = torch.tensor(special_ortho_group.rvs(size=num_orientations, dim=3), dtype=torch.float32)
+
+# 4. Create an arbitrary stack of Fourier space filters (identity filter in this example)
+# These filters operate on rffts of the 2D projection images
+# Filter shape: (..., h, w // 2 + 1)
+filters = torch.ones(template_size[0], template_size[1] // 2 + 1, dtype=torch.complex64)
+
+# Perform template matching
+cross_correlation = torch_2dtm.match_template_dft_2d(
+    image_dft=image_dft,
+    template_dft=template_dft,
+    rotation_matrices=rotation_matrices,
+    filters=filters
+)
+# The result has shape (..., num_orientations, image_height, image_width)
+```
+
+## License
+
+This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.

torch_2dtm-0.5.0/pyproject.toml

@@ -0,0 +1,112 @@
+# 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-2dtm@v(?P<version>.+)$"
+fallback-version = "0.5.0"
+
+[tool.hatch.version.raw-options]
+search_parent_directories = true
+# Parse tags of the form: <package-name>@v<semver>
+tag_regex = "^torch-2dtm@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-2dtm@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-2dtm"
+dynamic = ["version"]
+description = "2D template matching in pytorch"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "BSD-3-Clause" }
+authors = [
+    { name = "Josh Dickerson", email = "jdickerson@berkeley.edu" },
+    { name = "Matthew Giammar", email = "matthew_giammar@berkeley.edu" },
+    { name = "Alister Burt", email = "alisterburt@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",
+    "setuptools",          # needed for torch compilation? lol
+]
+
+
+# https://peps.python.org/pep-0621/#dependencies-optional-dependencies
+# "extras" (e.g. for `pip install .[test]`)
+[dependency-groups]
+# add dependencies used for testing here
+test = ["pytest", "pytest-cov", "scipy"]
+# add anything else you like to have in your dev environment here
+dev = [
+    { include-group = "test" },
+    "ipython",
+    "pdbpp",                    # https://github.com/pdbpp/pdbpp
+    "rich",                     # https://github.com/Textualize/rich
+]
+
+[project.urls]
+homepage = "https://github.com/teamtomo/teamtomo"
+repository = "https://github.com/teamtomo/teamtomo"
+
+# Entry points
+# https://peps.python.org/pep-0621/#entry-points
+# same as console_scripts entry point
+# [project.scripts]
+# torch-2dtm-cli = "torch_2dtm:main_cli"
+
+# [project.entry-points."some.group"]
+# tomatoes = "torch_2dtm:main_tomatoes"
+
+
+# https://docs.pytest.org/
+[tool.pytest.ini_options]
+minversion = "7.0"
+testpaths = ["tests"]
+filterwarnings = ["error"]
+
+# https://coverage.readthedocs.io/
+[tool.coverage.report]
+show_missing = true
+exclude_lines = [
+    "pragma: no cover",
+    "if TYPE_CHECKING:",
+    "@overload",
+    "except ImportError",
+    "\\.\\.\\.",
+    "raise NotImplementedError()",
+    "pass",
+]
+
+[tool.coverage.run]
+source = ["torch_2dtm"]
+
+# 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/**/*"]

torch_2dtm-0.5.0/src/torch_2dtm/__init__.py

@@ -0,0 +1,9 @@
+"""2D template matching in pytorch"""
+
+__version__ = "0.1.0"
+__author__ = "Josh Dickerson"
+__email__ = "jdickerson@berkeley.edu"
+
+from .cross_correlate import match_template_dft_2d
+
+__all__ = ["match_template_dft_2d"]

torch_2dtm-0.5.0/src/torch_2dtm/cross_correlate.py

@@ -0,0 +1,90 @@
+"""Cross-correlation functions."""
+
+import torch
+import platform
+import einops
+from einops._torch_specific import allow_ops_in_compiled_graph
+from torch_fourier_slice import extract_central_slices_rfft_3d
+
+from torch_2dtm.utils import normalize_template_projection
+
+# compile normalization utility function
+allow_ops_in_compiled_graph()
+if platform.system() == "Linux":
+    COMPILE_BACKEND = "aot_eager"  # More stable than inductor on Linux
+else:
+    COMPILE_BACKEND = "inductor"  # inductor for macOS
+
+normalize_template_projection_compiled = torch.compile(
+    normalize_template_projection, backend=COMPILE_BACKEND
+)
+
+
+def match_template_dft_2d(
+    image_dft: torch.Tensor,
+    template_dft: torch.Tensor,
+    rotation_matrices: torch.Tensor,
+    filters: torch.Tensor,
+) -> torch.Tensor:
+    """Batched projection and cross-correlation with a set of filters.
+
+    Note that this function returns a cross-correlation image which is the
+    same size as the input image prior to FFT calculation.
+
+    Parameters
+    ----------
+    image_dft : torch.Tensor
+        `(h_im, w_im // 2 + 1)` fourier transform (rfft) of the real space image.
+        Any filters etc are assumed to have already been applied to this image.
+    template_dft : torch.Tensor
+       `(d, h, w // 2 + 1)` fftshifted fourier transform (rfft) of the real valued template volume to take Fourier
+        slices from.
+    rotation_matrices : torch.Tensor
+        `(b, 3, 3)` batched rotation matrices to rotate slices sampled from the template fourier transform.
+    filters : torch.Tensor
+        `(..., h, w // 2 + 1)` filters applied to FFT slices which are fftshifted results of a rfft.
+
+    Returns
+    -------
+    torch.Tensor
+        Cross-correlation of the image with the template volume for each
+        orientation and defocus value. Will have shape
+        (orientations, defocus_batch, H, W).
+    """
+    # Grab relevant dimensions
+    _, h, w = template_dft.shape
+    h_im, w_im = image_dft.shape
+    w_im = 2 * (w_im - 1)
+    w = 2 * (w - 1)
+
+    # Extract central slice(s) from the template volume
+    fourier_slices = extract_central_slices_rfft_3d(
+        volume_rfft=template_dft,
+        rotation_matrices=rotation_matrices,
+    )  # (b, h, w)
+    fourier_slices = torch.fft.ifftshift(fourier_slices, dim=(-2,))
+    fourier_slices[..., 0, 0] = 0 + 0j  # zero out the DC component (mean zero)
+    fourier_slices *= -1  # flip contrast
+
+    # Apply the projective filters with broadcasting
+    filters = einops.rearrange(filters, "... h w -> ... 1 h w")
+    fourier_slices = fourier_slices * filters  # (..., b, h, w)
+
+    # Inverse Fourier transform into real space and normalize
+    projections = torch.fft.irfftn(fourier_slices, dim=(-2, -1))
+    projections = torch.fft.ifftshift(projections, dim=(-2, -1))
+    projections = normalize_template_projection_compiled(
+        projections, (h, w), (h_im, w_im)
+    )
+
+    # Padded forward Fourier transform for cross-correlation
+    projections_dft = torch.fft.rfftn(projections, dim=(-2, -1), s=(h_im, w_im))
+
+    # Zero the DC component (set mean zero)
+    projections_dft[..., 0, 0] = 0 + 0j
+
+    # Cross correlation step by element-wise multiplication
+    projections_dft = image_dft * torch.conj(projections_dft)
+    cross_correlation = torch.fft.irfftn(projections_dft, dim=(-2, -1))
+
+    return cross_correlation  # (..., h_im, w_im)

torch_2dtm-0.5.0/src/torch_2dtm/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_2dtm-0.5.0/src/torch_2dtm/utils.py

@@ -0,0 +1,88 @@
+"""Utility functions associated with backend functions."""
+
+import torch
+import einops
+
+
+def normalize_template_projection(
+    projections: torch.Tensor,
+    small_shape: tuple[int, int],
+    large_shape: tuple[int, int],
+) -> torch.Tensor:
+    r"""Subtract mean of edge values and set variance to 1 (in large shape).
+
+    This function uses the fact that variance of a sequence, Var(X), is scaled by the
+    relative size of the small (unpadded) and large (padded with zeros) space. Some
+    negligible error is introduced into the variance (~1e-4) due to this routine.
+
+    Let $X$ be the large, zero-padded projection and $x$ the small projection each
+    with sizes $(H, W)$ and $(h, w)$, respectively. The mean of the zero-padded
+    projection in terms of the small projection is:
+    .. math::
+        \begin{align}
+            \mu(X) &= \frac{1}{H \cdot W} \sum_{i=1}^{H} \sum_{j=1}^{W} X_{ij} \\
+            \mu(X) &= \frac{1}{H \cdot W} \sum_{i=1}^{h} \sum_{j=1}^{w} X_{ij} + 0 \\
+            \mu(X) &= \frac{h \cdot w}{H \cdot W} \mu(x)
+        \end{align}
+    The variance of the zero-padded projection in terms of the small projection can be
+    obtained by:
+    .. math::
+        \begin{align}
+            Var(X) &= \frac{1}{H \cdot W} \sum_{i=1}^{H} \sum_{j=1}^{W} (X_{ij} -
+                \mu(X))^2 \\
+            Var(X) &= \frac{1}{H \cdot W} \left(\sum_{i=1}^{h}
+                \sum_{j=1}^{w} (X_{ij} - \mu(X))^2 +
+                \sum_{i=h+1}^{H}\sum_{i=w+1}^{W} \mu(X)^2 \right) \\
+            Var(X) &= \frac{1}{H \cdot W} \sum_{i=1}^{h} \sum_{j=1}^{w} (X_{ij} -
+                \mu(X))^2 + (H-h)(W-w)\mu(X)^2
+        \end{align}
+
+    Parameters
+    ----------
+    projections : torch.Tensor
+        `(..., h, w)` real-space projections of the template (in small space).
+    small_shape : tuple[int, int]
+        `(h, w)` shape of the template (in real space).
+    large_shape : tuple[int, int]
+        `(h_im, w_im)` shape of the image (in real space).
+
+    Returns
+    -------
+    projections: torch.Tensor
+        `(..., h, w)` edge-mean subtracted projections
+         normalized so variance of zero-padded projection would be 1.
+    """
+    h, w = small_shape
+    h_im, w_im = large_shape
+
+    # Extract edges while preserving batch dimensions
+    top_edge = projections[..., 0, :]  # shape: (..., w)
+    bottom_edge = projections[..., -1, :]  # shape: (..., w)
+    left_edge = projections[..., 1:-1, 0]  # shape: (..., h-2)
+    right_edge = projections[..., 1:-1, -1]  # shape: (..., h-2)
+    edge_pixels = torch.concatenate(
+        [top_edge, bottom_edge, left_edge, right_edge], dim=-1
+    )  # shape: (..., w + w + h-2 + h-2)
+
+    # Subtract the edge pixel mean and calculate variance of small, unpadded projection
+    edge_mean = einops.reduce(edge_pixels, "... b -> ...", reduction="mean")
+    edge_mean = einops.rearrange(edge_mean, "... -> ... 1 1")
+    projections -= edge_mean
+
+    # Fast calculation of mean/var using Torch + appropriate scaling.
+    relative_size = h * w / (h_im * w_im)
+    per_image_mean = einops.reduce(projections, "... h w -> ...", reduction="mean")
+    per_image_mean *= relative_size**2
+
+    # First term of the variance calculation
+    diff = projections - einops.rearrange(per_image_mean, "... -> ... 1 1")
+    per_image_variance = einops.reduce(
+        diff**2, pattern="... h w -> ...", reduction="sum"
+    )
+
+    # Add the second term of the variance calculation
+    per_image_variance += (h_im - h) * (w_im - w) * per_image_mean**2
+    per_image_variance /= h_im * w_im
+    per_image_variance = einops.rearrange(per_image_variance, "... -> ... 1 1")
+
+    return projections / torch.sqrt(per_image_variance)

torch_2dtm-0.5.0/tests/test_torch_2dtm.py

@@ -0,0 +1,40 @@
+import torch
+import torch_2dtm
+
+from scipy.stats import special_ortho_group
+
+
+def test_template_match_dft_2d():
+    # Create random test data
+    # 1. Create a random image and compute its FFT
+    image_size = (128, 128)
+    image = torch.randn(*image_size, dtype=torch.float32)
+    image_dft = torch.fft.rfftn(image, dim=(0, 1))  # Shape: (128, 65)
+
+    # 2. Create a random 3D template and compute its FFT
+    template_size = (64, 64, 64)
+    template = torch.randn(*template_size, dtype=torch.float32)
+    template_dft = torch.fft.rfftn(template, dim=(0, 1, 2))  # Shape: (64, 64, 33)
+
+    # 3. Create a batch of random rotation matrices with shape (b, 3, 3)
+    num_orientations = 10
+    rotation_matrices = torch.tensor(
+        special_ortho_group.rvs(size=num_orientations, dim=3), dtype=torch.float32
+    )
+
+    # 4. Create an arbitrary stack of Fourier space filters (identity filter in this example)
+    # These filters operate on rffts of the 2D projection images
+    # Filter shape: (..., h, w // 2 + 1)
+    filters_shape = (5, 4, 3, template_size[0], template_size[1] // 2 + 1)
+    filters = torch.ones(filters_shape, dtype=torch.complex64)
+
+    # Perform template matching
+    cross_correlation = torch_2dtm.match_template_dft_2d(
+        image_dft=image_dft,
+        template_dft=template_dft,
+        rotation_matrices=rotation_matrices,
+        filters=filters,
+    )
+
+    # correct output shape is (..., num_orientations, h, w)
+    assert cross_correlation.shape == (5, 4, 3, num_orientations, *image_size)