pytme 0.2.9.post1__tar.gz → 0.3b0__tar.gz

{pytme-0.2.9.post1 → pytme-0.3b0}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: pytme
-Version: 0.2.9.post1
+Version: 0.3b0
 Summary: Python Template Matching Engine
 Author: Valentin Maurer
 Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>
-License: Proprietary
+License-Expression: GPL-2.0-only
 Project-URL: Homepage, https://github.com/KosinskiLab/pyTME
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
@@ -12,7 +12,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: mrcfile>=1.4.3
-Requires-Dist: numpy<2.0,>=1.22.2
+Requires-Dist: numpy>=1.22.2
 Requires-Dist: scipy>=1.9.1
 Requires-Dist: pyfftw>=0.13.1
 Requires-Dist: pytest>=6.2.5
@@ -25,19 +25,20 @@ Requires-Dist: h5py
 Requires-Dist: importlib_resources
 Requires-Dist: joblib
 Provides-Extra: cupy
-Requires-Dist: cupy-cuda12x; extra == "cupy"
-Requires-Dist: voltools==0.6.0; extra == "cupy"
+Requires-Dist: cupy-cuda12x>=13.0.0; extra == "cupy"
 Provides-Extra: pytorch
 Requires-Dist: torch; extra == "pytorch"
 Requires-Dist: torchvision; extra == "pytorch"
 Provides-Extra: jax
-Requires-Dist: jax; extra == "jax"
+Requires-Dist: jax[cuda12]; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"
+Provides-Extra: jax-cpu
+Requires-Dist: jax; extra == "jax-cpu"
+Requires-Dist: jaxlib; extra == "jax-cpu"
 Provides-Extra: all
-Requires-Dist: cupy; extra == "all"
-Requires-Dist: voltools==0.6.0; extra == "all"
+Requires-Dist: cupy-cuda12x>=13.0.0; extra == "all"
 Requires-Dist: torch; extra == "all"
-Requires-Dist: jax; extra == "all"
+Requires-Dist: jax[cuda12]; extra == "all"
 Requires-Dist: jaxlib; extra == "all"
 Requires-Dist: tifffile; extra == "all"
 Dynamic: author

{pytme-0.2.9.post1 → pytme-0.3b0}/pyproject.toml

@@ -7,13 +7,13 @@ name="pytme"
 authors = [
   { name = "Valentin Maurer", email = "valentin.maurer@embl-hamburg.de" },
 ]
-version="0.2.9.post1"
+version="0.3.b0"
 description="Python Template Matching Engine"
 readme="README.md"
 requires-python = ">=3.11"
 dependencies=[
     "mrcfile>=1.4.3",
-    "numpy>=1.22.2,<2.0",
+    "numpy>=1.22.2",
     "scipy>=1.9.1",
     "pyfftw>=0.13.1",
     "pytest>=6.2.5",
@@ -26,28 +26,28 @@ dependencies=[
     "importlib_resources",
     "joblib"
 ]
-license = {text = "Proprietary"}
+license = "GPL-2.0-only"
 classifiers = [
     "Programming Language :: Python :: 3",
     "Operating System :: OS Independent",
 ]
 
 [project.optional-dependencies]
-cupy = ["cupy-cuda12x", "voltools==0.6.0"]
+cupy = ["cupy-cuda12x>=13.0.0"]
 pytorch = ["torch", "torchvision"]
-jax = ["jax", "jaxlib"]
+jax = ["jax[cuda12]", "jaxlib"]
+jax_cpu = ["jax", "jaxlib"]
 all = [
-    "cupy",
-    "voltools==0.6.0",
+    "cupy-cuda12x>=13.0.0",
     "torch",
-    "jax",
+    "jax[cuda12]",
     "jaxlib",
     "tifffile"
 ]
 
 [project.scripts]
 match_template = "scripts:match_template.main"
-estimate_ram_usage = "scripts:estimate_ram_usage.main"
+estimate_memory_usage = "scripts:estimate_memory_usage.main"
 preprocessor_gui = "scripts:preprocessor_gui.main"
 preprocess = "scripts:preprocess.main"
 postprocess = "scripts:postprocess.main"

{pytme-0.2.9.post1 → pytme-0.3b0}/pytme.egg-info/SOURCES.txt

@@ -4,11 +4,13 @@ README.md
 pyproject.toml
 setup.py
 scripts/__init__.py
-scripts/estimate_ram_usage.py
+scripts/estimate_memory_usage.py
+scripts/extract_candidates.py
 scripts/match_template.py
 scripts/postprocess.py
 scripts/preprocess.py
 scripts/preprocessor_gui.py
+scripts/refine_matches.py
 tests/__init__.py
 tests/test_analyzer.py
 tests/test_backends.py
@@ -34,10 +36,8 @@ tests/data/Blurring/localgaussian_sigma0510.npy
 tests/data/Blurring/mean_size5.npy
 tests/data/Blurring/ntree_sigma0510.npy
 tests/data/Blurring/rank_rank3.npy
-tests/data/Maps/.DS_Store
 tests/data/Maps/emd_8621.mrc.gz
 tests/data/Raw/em_map.map
-tests/data/Structures/.DS_Store
 tests/data/Structures/1pdj.cif
 tests/data/Structures/1pdj.pdb
 tests/data/Structures/5khe.cif
@@ -51,6 +51,7 @@ tests/preprocessing/test_preprocessor.py
 tests/preprocessing/test_utils.py
 tme/__init__.py
 tme/__version__.py
+tme/cli.py
 tme/density.py
 tme/matching_data.py
 tme/matching_exhaustive.py
@@ -67,7 +68,9 @@ tme/types.py
 tme/analyzer/__init__.py
 tme/analyzer/_utils.py
 tme/analyzer/aggregation.py
+tme/analyzer/base.py
 tme/analyzer/peaks.py
+tme/analyzer/proxy.py
 tme/backends/__init__.py
 tme/backends/_cupy_utils.py
 tme/backends/_jax_utils.py

pytme-0.2.9.post1/scripts/estimate_ram_usage.py → pytme-0.3b0/scripts/estimate_memory_usage.py

@@ -1,20 +1,21 @@
 #!python3
-""" Estimate RAM requirements for template matching jobs.
+"""Estimate memory requirements for template matching jobs.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 import numpy as np
 import argparse
 from tme import Density
-from tme.matching_utils import estimate_ram_usage
+from tme.memory import estimate_memory_usage
 from tme.matching_exhaustive import MATCHING_EXHAUSTIVE_REGISTER
 
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description="Estimate RAM usage for template matching."
+        description="Estimate memory usage for template matching.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
     parser.add_argument(
         "-m",
@@ -32,14 +33,9 @@ def parse_args():
         required=True,
         help="Path to a template in PDB/MMCIF or CCP4/MRC format.",
     )
-    parser.add_argument(
-        "--matching_method",
-        required=False,
-        default=None,
-        help="Analyzer method to use.",
-    )
     parser.add_argument(
         "-s",
+        "--score",
         dest="score",
         type=str,
         default="FLCSphericalMask",
@@ -50,21 +46,12 @@ def parse_args():
         "--ncores", type=int, help="Number of cores for parallelization.", required=True
     )
     parser.add_argument(
-        "--no_edge_padding",
-        dest="no_edge_padding",
+        "--pad_edges",
+        dest="pad_edges",
         action="store_true",
         default=False,
-        help="Whether to pad the edges of the target. This is useful, if the target"
-        " has a well defined bounding box, e.g. a density map.",
-    )
-    parser.add_argument(
-        "--no_fourier_padding",
-        dest="no_fourier_padding",
-        action="store_true",
-        default=False,
-        help="Whether input arrays should be zero-padded to the full convolution shape"
-        " for numerical stability. When working with very large targets such as"
-        " tomograms it is safe to use this flag and benefit from the performance gain.",
+        help="Whether to pad the edges of the target. Useful if the target does not "
+        "a well-defined bounding box. Defaults to True if splitting is required.",
     )
     args = parser.parse_args()
     return args
@@ -72,19 +59,15 @@ def parse_args():
 
 def main():
     args = parse_args()
-    target = Density.from_file(args.target)
-    template = Density.from_file(args.template)
-
-    target_box = target.shape
-    if not args.no_edge_padding:
-        target_box = np.add(target_box, template.shape)
+    target = Density.from_file(args.target, use_memmap=True)
+    template = Density.from_file(args.template, use_memmap=True)
 
     template_box = template.shape
-    if args.no_fourier_padding:
+    if not args.pad_edges:
         template_box = np.ones(len(template_box), dtype=int)
 
-    result = estimate_ram_usage(
-        shape1=target_box,
+    result = estimate_memory_usage(
+        shape1=target.shape,
         shape2=template_box,
         matching_method=args.score,
         ncores=args.ncores,

pytme-0.3b0/scripts/extract_candidates.py

@@ -0,0 +1,224 @@
+#!python3
+""" Prepare orientations stack for refinement.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+import argparse
+from os.path import splitext
+
+import numpy as np
+
+from tme import Density, Orientations
+from tme.matching_utils import (
+    generate_tempfile_name,
+    rotation_aligning_vectors,
+    euler_from_rotationmatrix,
+    euler_to_rotationmatrix,
+)
+
+
+class ProgressBar:
+    """
+    ASCII progress bar.
+    """
+
+    def __init__(self, message: str, nchars: int, total: int):
+        self._size = nchars - len(message) - (len(str(total))+2) * 2
+        self._message = message
+        self._total = total
+
+    def update(self, cur):
+        x = int(cur * self._size / self._total)
+        print(
+            "%s[%s%s] %i/%i\r"
+            % (self._message, "#" * x, "." * (self._size - x), cur, self._total),
+            end="",
+        )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Extract matching candidates for further refinement.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    io_group = parser.add_argument_group("Input / Output")
+    io_group.add_argument(
+        "--target",
+        required=True,
+        type=str,
+        help="Extract candidates from this target.",
+    )
+    io_group.add_argument(
+        "--orientations",
+        required=True,
+        type=str,
+        help="Path to file generated by postprocess.py using output_format orientations.",
+    )
+    io_group.add_argument(
+        "--orientations_sampling",
+        required=False,
+        type=float,
+        default=1.0,
+        help="Factor to map candidate coordinates onto the target. Only relevant if "
+        "target sampling rate differs from candidate orientation sampling rate.",
+    )
+    io_group.add_argument(
+        "-o",
+        "--output_file",
+        required=True,
+        type=str,
+        help="Path to write output H5 file.",
+    )
+
+    alignment_group = parser.add_argument_group("Alignment")
+    alignment_group.add_argument(
+        "--align_orientations",
+        action="store_true",
+        required=False,
+        help="Whether to align extracted orientations based on their angles. Allows "
+        "for efficient subsequent sampling of cone angles.",
+    )
+    alignment_group.add_argument(
+        "--angles_are_vector",
+        action="store_true",
+        required=False,
+        help="Considers euler_z euler_y, euler_x as vector that will be rotated to align "
+        "with the z-axis (1,0,0). Only considered when --align_orientations is set.",
+    )
+    alignment_group.add_argument(
+        "--interpolation_order",
+        dest="interpolation_order",
+        required=False,
+        type=int,
+        default=1,
+        help="Interpolation order for alignment, less than zero is no interpolation.",
+    )
+
+    extraction_group = parser.add_argument_group("Extraction")
+    extraction_group.add_argument(
+        "--box_size",
+        required=False,
+        type=int,
+        help="Box size for extraction, defaults to two times the template.",
+    )
+    extraction_group.add_argument(
+        "--translation_uncertainty",
+        required=False,
+        type=int,
+        help="Sets box size for extraction to template box plus this value.",
+    )
+    extraction_group.add_argument(
+        "--keep_out_of_box",
+        action="store_true",
+        required=False,
+        help="Whether to keep orientations that fall outside the box. If the "
+        "orientations are sensible, it is safe to pass this flag.",
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+def main():
+    args = parse_args()
+    orientations = Orientations.from_file(args.orientations)
+    orientations.translations = np.divide(
+        orientations.translations, args.orientations_sampling
+    )
+
+    target = Density.from_file(args.target, use_memmap=True)
+
+    box_size = np.array(args.box_size)
+    box_size = np.repeat(box_size, target.data.ndim // box_size.size).astype(int)
+
+    extraction_shape = np.copy(box_size)
+    if args.align_orientations:
+        extraction_shape[:] = int(np.linalg.norm(box_size) + 1)
+
+    orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
+        target_shape=target.shape,
+        extraction_shape=extraction_shape,
+        drop_out_of_box=not args.keep_out_of_box,
+        return_orientations=True,
+    )
+
+    if args.align_orientations:
+        for index in range(orientations.rotations.shape[0]):
+            rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
+            rotation_matrix = np.linalg.inv(rotation_matrix)
+            if args.angles_are_vector:
+                rotation_matrix = rotation_aligning_vectors(
+                    orientations.rotations[index], target_vector=(1, 0, 0)
+                )
+            orientations.rotations[index] = euler_from_rotationmatrix(rotation_matrix)
+
+    filename = generate_tempfile_name()
+    output_dtype = target.data.dtype
+    if args.align_orientations is not None:
+        output_dtype = np.float32
+
+    target.data = target.data.astype(output_dtype)
+
+    dens = Density(
+        np.memmap(
+            filename,
+            mode="w+",
+            shape=(len(obs_slices), *box_size),
+            dtype=output_dtype,
+        ),
+        sampling_rate=(1, *target.sampling_rate),
+        origin=(0, *target.origin),
+    )
+    dens.data[:] = target.metadata["mean"]
+
+    data_subset = np.zeros(extraction_shape, dtype=target.data.dtype)
+    pbar = ProgressBar(message="Aligning ", nchars=80, total=len(obs_slices))
+    for index, (obs_slice, cand_slice) in enumerate(zip(obs_slices, cand_slices)):
+        pbar.update(index + 1)
+
+        data_subset.fill(0)
+        data_subset[cand_slice] = target.data[obs_slice]
+        target_subset = Density(
+            data_subset,
+            sampling_rate=target.sampling_rate,
+            origin=target.origin,
+        )
+
+        if args.align_orientations:
+            rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
+            target_subset = target_subset.rigid_transform(
+                rotation_matrix=rotation_matrix,
+                use_geometric_center=True,
+                order=args.interpolation_order,
+            )
+        target_subset.pad(box_size, center=True)
+
+        # target_value = target.data[tuple(orientations.translations[index].astype(int))]
+        # center = np.divide(target_subset.data.shape, 2).astype(int)
+        # print(np.where(target_subset.data == target_value), center)
+        # print(target_subset.data[tuple(center.astype(int))],
+        # target_value,
+        # target_subset.data[tuple(center.astype(int))] == target_value
+        # )
+
+        dens.data[index] = target_subset.data
+    print("")
+
+    target_meta = {
+        k: v for k, v in target.metadata.items() if k in ("mean", "max", "min", "std")
+    }
+    dens.metadata.update(target_meta)
+    dens.metadata["batch_dimension"] = (0,)
+    dens.metadata["normals"] = orientations.rotations
+
+    dens.to_file(args.output_file)
+    orientations.to_file(
+        f"{splitext(args.output_file)[0]}_aligned.tsv", file_format="text"
+    )
+
+if __name__ == "__main__":
+    main()