pytme 0.2.0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.1__cp311-cp311-macosx_14_0_arm64.whl

{pytme-0.2.0.dist-info → pytme-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytme
-Version: 0.2.0
+Version: 0.2.1
 Summary: Python Template Matching Engine
 Author: Valentin Maurer
 Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>

pytme-0.2.1.dist-info/RECORD

@@ -0,0 +1,73 @@
+pytme-0.2.1.data/scripts/estimate_ram_usage.py,sha256=R1NDpFajcF-MonJ4a43SfDlA-nxBYwK7D2quzCdsVFM,2767
+pytme-0.2.1.data/scripts/match_template.py,sha256=d9c_014dy4X9jFtulHkypuRqcDTCJsddA5-tt2ZUw6g,39833
+pytme-0.2.1.data/scripts/postprocess.py,sha256=dg2TQpoTiQFDm1Q-l-WH3g0kIqwBEvKk-hX1QhhgjoI,23132
+pytme-0.2.1.data/scripts/preprocess.py,sha256=zog-l2Je-GeouJ6SnamOMuHgTn7fFPiGnO5X03y5qSY,2527
+pytme-0.2.1.data/scripts/preprocessor_gui.py,sha256=-lYdd9--xvgZghhuPE43JGycventc3tPoG4JF1wkXo0,35254
+scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+scripts/estimate_ram_usage.py,sha256=rN7haobnHg3YcgGJIp81FNiCzy8-saJGeEurQlmQmNQ,2768
+scripts/extract_candidates.py,sha256=j8SZNw4-F73_7o0_NorioQcyH8LatrbNX4pwfe-XC9A,8139
+scripts/match_template.py,sha256=yTFtX0cZqk6Q38wtlDDTzHRRJsGLNnoVTXQeXWBH8Rc,39834
+scripts/match_template_filters.py,sha256=-QnSZsGgLP5k6nH5DXINO2Yd7c1zEq_qyIDwfGB0w8o,40272
+scripts/postprocess.py,sha256=L67WF7A9RIWEjr04KzHjHDg4LUpbn6h19JFKh_OyonI,23133
+scripts/preprocess.py,sha256=ebJVLxbRlB6TI5YHNr0VavZ4lmaRdf8QVafyiDhh_oU,2528
+scripts/preprocessor_gui.py,sha256=2-IKFVwztnvOa_G6fZ4SYsk1g6rFb3_-LMDmWSHqCno,35255
+scripts/refine_matches.py,sha256=BlQqZFYdatAUAPly14xobDgqJOWvwEidjTz4pW4QBlg,14628
+tme/__init__.py,sha256=MGVeRX1f5k6RwExBcffZtCM7igNNtpoIr2G2yJZ4OHI,251
+tme/__version__.py,sha256=HfjVOrpTnmZ-xVFCYSVmX50EXaBQeJteUHG-PD6iQs8,22
+tme/analyzer.py,sha256=esp23U15jsctbtDy1QMcUwsiBCwXbapiRfH3ZEVvwOg,61500
+tme/density.py,sha256=cMfHfZSY8Ff9w7HfE-Sqigo4XKT_naht3jMtdWg4sKk,88846
+tme/extensions.cpython-311-darwin.so,sha256=nIqjOFVGoKxod55SSK7hNvHxk6uTnTQ-oqkYdSsmn6Y,411792
+tme/helpers.py,sha256=TMtBuJoZk6q4M_rfkr8yPpGzJD74ycqyN3yXMuU9mr4,23625
+tme/matching_constrained.py,sha256=NAJnZpusjQs5rKabHiNyrclXSPm27K0V8HVFjTQ2zB0,6725
+tme/matching_data.py,sha256=QeAhwgqiwCpQIghlj9wN0jCbW1UpZH55DOZWFudtXSw,24869
+tme/matching_exhaustive.py,sha256=YidMx0RGnV-YdAK-tEdG9aHg06sSE3QgV-D2KMTKcZE,63676
+tme/matching_memory.py,sha256=bmCAUYyXWEet-1XXhldtc0irio2ytMSsAzWYyFI5LNM,11273
+tme/matching_optimization.py,sha256=wil0-5DYrqLzgFMhi-uClkt5mzp15AE_--zmmhJbdkw,43530
+tme/matching_utils.py,sha256=qLcRcfMCURsA5vFvJDUBa7xIAHsICWn_PB_jOYK3Pb0,42520
+tme/orientations.py,sha256=3SuKZhscTGnclOfNJCz2Rq1KiePHFytPdjuHiQi7m5Q,23929
+tme/parser.py,sha256=tA9ABeV95cZn8lJCukVUaocQ9RguR6ZZzQsMXf_-ud0,13887
+tme/preprocessor.py,sha256=ylVO-D1o9lpvvApXEymkmtdRrvaFd2HrJSn9nsWcMTs,51225
+tme/structure.py,sha256=mIhDhGksMf8cSARr0ofNhJQUpQIdpKZtJsX63Law2z0,52503
+tme/types.py,sha256=2Tyh_xnMLxIWYb3aJDAUb6GWpaL6gcYMUm2YNbJlAPI,295
+tme/backends/__init__.py,sha256=xNX4Vx182F57AwZjb0DwLJA-sps6DixnDRWPcMlOvWw,4397
+tme/backends/cupy_backend.py,sha256=1Wsnoo0HiIELRGb8wT9VZJ5spVEU5-DVbyyY092RixI,13230
+tme/backends/jax_backend.py,sha256=uq_Vk_ofb5JG9IGg3SfHn3YqcBD9c4wWypgdCu9HTno,7372
+tme/backends/matching_backend.py,sha256=eFoowia-uzc1UYAq9AnDN3CMN31rWAT1ilCskH6sruw,29876
+tme/backends/mlx_backend.py,sha256=pq4U3m3-dAkSC5y6P0R1CayvVKwi5rZjTVoeG4XOQ8g,8569
+tme/backends/npfftw_backend.py,sha256=_PzGKO6fuo20RofChRqq7WlfqVm-6Ne6tMPEyxB3fvE,27856
+tme/backends/pytorch_backend.py,sha256=cY3KSQdLy80NVvKSMy5Jj6HzD1UDsLRQTeC2WA70Abo,18694
+tme/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tme/data/c48n309.npy,sha256=NwH64mOEbm3tStq5c98o81fY1vMOoq4nvXDAh7Z7iZg,296768
+tme/data/c48n527.npy,sha256=saSUMTa1R0MisPvgFL02a7IHQSwEZ-mJu0v3qJjg5AU,506048
+tme/data/c48n9.npy,sha256=bDVLV6mWjZHSQfeDc-MOCKKarfc1jaNeVvpoe2xMUy4,8768
+tme/data/c48u1.npy,sha256=JeXMFzFITs2ezdc3x5lp3jo1cHHHHVADSA1Tpf77kXs,1088
+tme/data/c48u1153.npy,sha256=ECiEximtYDWtIux3Fwe_EJlyn08gUqP85DN9gjkT9_k,1107008
+tme/data/c48u1201.npy,sha256=aceC_Jeienz_81X4520nPpZcg5tnRhbW795EqbpWkrg,1153088
+tme/data/c48u1641.npy,sha256=p4LwW3LzdTjrUUpA7H53RfNWxYfPX0XjeSwZ39Ac78Q,1575488
+tme/data/c48u181.npy,sha256=mLYXrv1YHLH6DsBp5MkxHkxlxgMnj1mw_KKI0udH-FY,173888
+tme/data/c48u2219.npy,sha256=p8TQeX8YHu4pdxnwJjEAlQWAPa66W7kpK96iZKZr9JE,2130368
+tme/data/c48u27.npy,sha256=k03ZNEsoPwBKCy8IeIa5G0WRZqjGZMtX6Ibu7EpJHvU,26048
+tme/data/c48u2947.npy,sha256=icI97ED6ct66y7FIaJAugmjzrIWk7CINCxtO3wDTnrU,2829248
+tme/data/c48u3733.npy,sha256=tla-__Pf-hpN6h04vtFIfkkFdCLple11VO06kr1dXkM,3583808
+tme/data/c48u4749.npy,sha256=tItOA4oV7SiqCCREwz3fyEpZoxM0lCq_jfEo5_-fp2s,4559168
+tme/data/c48u5879.npy,sha256=bFk89MllIFCX_sLXTYWFquSyN1NuahH4wwnEsPJLxzA,5643968
+tme/data/c48u7111.npy,sha256=CMy9kI2edH-q9eTIVdgUtXurplYNI7Uqp4dXfkkVdf8,6826688
+tme/data/c48u815.npy,sha256=bCuJxLtm0Sjg3GGxtyjGzRYZ1G0Gz79XHI-71GvqQnI,782528
+tme/data/c48u83.npy,sha256=7ODJYnsiuDjGbgd9GFopsyIW2IjrYI0J2X2f-cK868U,79808
+tme/data/c48u8649.npy,sha256=-IPlpR4zrPQZWhhSPu4zEulFdrCEVgTMFffCB5d-huE,8303168
+tme/data/c600v.npy,sha256=JqSu3ALoL1A9iguehc0YGUMFPsh2fprHHp76VXeFXIw,2528
+tme/data/c600vc.npy,sha256=Yht-GFXDSjjGvsjFBvyxxEZAI-ODADPd5gEgFNZQVTA,14528
+tme/data/metadata.yaml,sha256=fAgX-mEzB0QMHTEtYDG4cSMbJhYxBbDJH3sdvJvL7a8,750
+tme/data/quat_to_numpy.py,sha256=-gkDZb10fKBxwfYrSLCUWvMB76TzZWELCeKsYProwws,1333
+tme/preprocessing/__init__.py,sha256=7O3vDzJcIfxovJkf7avWSPtzaIVlTbmsW7egQFukC_s,98
+tme/preprocessing/_utils.py,sha256=jj5VSRLewO4uk0cHoEhpw2rTU9taT8iuZCijvbFcrw0,5190
+tme/preprocessing/composable_filter.py,sha256=sp3bN8JeFB2r384cEIgvN6yXjC53GCoPhBmGI0S0kbI,781
+tme/preprocessing/compose.py,sha256=tbQMo3LaKRBnOFKoqZ64as3i0t6U6312s76Csz0FPvk,1324
+tme/preprocessing/frequency_filters.py,sha256=k3MvXTBeoETTp5WINpqxEcmhiHSFh90QyAJBXe-fGxs,12614
+tme/preprocessing/tilt_series.py,sha256=OaYMuN4e0F5Z7tuNVwG_kn69E4zPAU6_xIoA5B-IQUw,34747
+pytme-0.2.1.dist-info/LICENSE,sha256=K1IUNSVAz8BXbpH5EA8y5FpaHdvFXnAF2zeK95Lr2bY,18467
+pytme-0.2.1.dist-info/METADATA,sha256=pTq0YVA2oWM7Zo8p-lk6ynkuZsEHAPnDDEbzcbEon_4,2165
+pytme-0.2.1.dist-info/WHEEL,sha256=BDgKu9_KNfDn85ptly0T56JpX4avXH07X_ZCqAJnQwY,110
+pytme-0.2.1.dist-info/entry_points.txt,sha256=ff3LQL3FCWfCYOwFiP9zatm7laUbnwCkuPELkQVyUO4,241
+pytme-0.2.1.dist-info/top_level.txt,sha256=J8FUkazOb2fZ0n_KexnqCGyNOtie2bwisFSUBiM5-0w,12
+pytme-0.2.1.dist-info/RECORD,,

scripts/extract_candidates.py

@@ -1,5 +1,5 @@
 #!python3
-""" CLI to simplify analysing the output of match_template.py.
+""" Prepare orientations stack for refinement.
 
     Copyright (c) 2023 European Molecular Biology Laboratory
 
@@ -12,13 +12,33 @@ import numpy as np
 
 from tme import Density, Orientations
 from tme.matching_utils import (
-    create_mask,
+    load_pickle,
     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)
+        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."
@@ -52,57 +72,61 @@ def parse_args():
     io_group.add_argument(
         "-o",
         "--output_file",
-        required=False,
+        required=True,
         type=str,
         help="Path to output HDF5 file.",
     )
 
-    extraction_group = parser.add_argument_group("Extraction")
-    extraction_group.add_argument(
-        "--box_size",
+    alignment_group = parser.add_argument_group("Alignment")
+    alignment_group.add_argument(
+        "--align_orientations",
+        action="store_true",
         required=False,
-        type=int,
-        help="Box size for extraction, defaults to two times the template box.",
+        help="Whether to align extracted orientations based on their angles. Allows "
+        "for efficient subsequent sampling of cone angles.",
     )
-    extraction_group.add_argument(
-        "--alignment_vector",
+    alignment_group.add_argument(
+        "--angles_are_vector",
+        action="store_true",
         required=False,
-        type=str,
-        help="Vector to align extracted coordinate to e.g. '1,0,0'. If provided, euler_z"
-        "euler_y and euler_x are assumed to be a vector that will be rotated to match "
-        "alignment_vector.",
+        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."
     )
-    extraction_group.add_argument(
+    alignment_group.add_argument(
         "--interpolation_order",
         dest="interpolation_order",
         required=False,
         type=int,
         default=1,
-        help="Spline interpolation used for template rotations. If less than zero "
-        "no interpolation is performed.",
+        help="Interpolation order for alignment, less than zero is no interpolation."
     )
+
+    extraction_group = parser.add_argument_group("Extraction")
     extraction_group.add_argument(
-        "--translation_uncertainty",
-        dest="translation_uncertainty",
+        "--box_size",
         required=False,
         type=int,
-        default=None,
-        help="Creates a cented spherical target mask with given radius in voxel.",
+        help="Box size for extraction, defaults to two times the template.",
     )
     extraction_group.add_argument(
-        "--mask_path",
-        dest="mask_path",
+        "--keep_out_of_box",
+        action="store_true",
         required=False,
-        type=str,
-        default="target_mask.h5",
-        help="Path to write spherical mask to, defaults to target_mask.h5.",
+        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()
 
-    if args.alignment_vector is not None:
-        args.alignment_vector = [float(x) for x in args.alignment_vector.split(",")]
-        if len(args.alignment_vector) != 3:
-            raise ValueError("Only three-dimensional vectors are supported")
+    data_present = args.target is not None and args.template is not None
+    if args.input_file is None and not data_present:
+        raise ValueError(
+            "Either --input_file or --target and --template need to be specified."
+        )
+    elif args.input_file is not None and data_present:
+        raise ValueError(
+            "Please specific either --input_file or --target and --template."
+        )
 
     return args
 
@@ -111,57 +135,59 @@ def main():
     args = parse_args()
     orientations = Orientations.from_file(args.orientations)
 
+    if args.input_file is not None:
+        data = load_pickle(args.input_file)
+        target_origin, _, sampling_rate, cli_args = data[-1]
+        args.target, args.template  = cli_args.target, cli_args.template
+
     target = Density.from_file(args.target, use_memmap=True)
-    template = Density.from_file(args.template)
 
-    box_size = template.shape
-    if args.box_size is not None:
-        args.box_size = args.box_size // 2
-        box_size = np.full(template.data.ndim, fill_value=args.box_size, dtype=int)
+    try:
+        template = Density.from_file(args.template)
+    except Exception:
+        template = Density.from_structure(args.template, sampling_rate = target.sampling_rate)
+
+    box_size = args.box_size
+    if box_size is None:
+        box_size = np.multiply(template.shape, 2)
+    box_size = np.array(box_size)
+    box_size = np.repeat(box_size, template.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)
 
-    subtomo_shape = np.multiply(box_size, 2).astype(int)
-    extraction_scaling = 3 if args.alignment_vector is not None else 2
-    extraction_shape = np.multiply(box_size, extraction_scaling).astype(int)
     orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
         target_shape=target.shape,
         extraction_shape=extraction_shape,
-        drop_out_of_box=True,
+        drop_out_of_box=not args.keep_out_of_box,
         return_orientations=True,
     )
 
+    if args.align_orientations:
+        orientations.rotations = orientations.rotations.astype(np.float32)
+        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.alignment_vector is not None:
+    if args.align_orientations is not None:
         output_dtype = np.float32
 
-    if args.translation_uncertainty is not None:
-        dens = Density(
-            np.memmap(
-                args.mask_path,
-                mode="w+",
-                shape=(len(obs_slices), *subtomo_shape),
-                dtype=output_dtype,
-            ),
-            sampling_rate=(1, *target.sampling_rate),
-            origin=(0, *target.origin),
-        )
-        dens.data[:] = 0
-        mask = create_mask(
-            mask_type="ellipse",
-            center=template.shape,
-            radius=args.translation_uncertainty,
-            shape=subtomo_shape,
-        )
-        dens.data[:] = mask
-        dens.to_file(args.mask_path)
-
     target.data = target.data.astype(output_dtype)
 
     dens = Density(
         np.memmap(
             filename,
             mode="w+",
-            shape=(len(obs_slices), *subtomo_shape),
+            shape=(len(obs_slices), *box_size),
             dtype=output_dtype,
         ),
         sampling_rate=(1, *target.sampling_rate),
@@ -169,50 +195,56 @@ def main():
     )
     dens.data[:] = target.metadata["mean"]
 
-    target_meta = {
-        k: v for k, v in target.metadata.items() if k in ("mean", "max", "min", "std")
-    }
-
-    total_slices = len(obs_slices)
-    magnitude = len(str(total_slices))
-    for index, obs_slice in enumerate(obs_slices):
-        print(f"Processing {index:{magnitude}} / {total_slices:{magnitude}}", end="\r")
+    print(target.data.shape)
+    # There appears to be an isseu with the stack creation. Trace this further
+    data_subset = np.zeros(extraction_shape, dtype = target.data.dtype)
+    pbar = ProgressBar(message = "Orientation ", 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(
-            target.data[obs_slice],
+            data_subset,
             sampling_rate=target.sampling_rate,
             origin=target.origin,
         )
 
-        rotation_matrix = np.eye(target_subset.data.ndim)
-        if args.alignment_vector is not None:
-            normal = orientations.rotations[index]
-            # To visualize only normals were picked we would need to do the opposite
-            # However, this allows us to align everything to the template axis
-            rotation_matrix = rotation_aligning_vectors(
-                normal, target_vector=args.alignment_vector
-            )
-
+        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(subtomo_shape, center=True)
+        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 ) + np.mod(target_subset.shape, 2)
+        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
-        orientations.rotations[index] = euler_from_rotationmatrix(rotation_matrix)
+    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.to_file(args.output_file)
     orientations.to_file(
-        f"{splitext(args.output_file)}_aligned.tsv",
+        f"{splitext(args.output_file)[0]}_aligned.tsv",
         file_format="text"
     )
     orientations.to_file(
-        f"{splitext(args.output_file)}_aligned.star",
+        f"{splitext(args.output_file)[0]}_aligned.star",
         file_format="relion"
     )
 
-
 if __name__ == "__main__":
     main()