pytme 0.3.1.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.2.dev0__cp311-cp311-macosx_15_0_arm64.whl

pytme-0.3.2.dev0.data/scripts/estimate_ram_usage.py

@@ -0,0 +1,97 @@
+#!python
+""" Estimate RAM requirements for template matching jobs.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    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.matching_exhaustive import MATCHING_EXHAUSTIVE_REGISTER
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Estimate RAM usage for template matching."
+    )
+    parser.add_argument(
+        "-m",
+        "--target",
+        dest="target",
+        type=str,
+        required=True,
+        help="Path to a target in CCP4/MRC format.",
+    )
+    parser.add_argument(
+        "-i",
+        "--template",
+        dest="template",
+        type=str,
+        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",
+        dest="score",
+        type=str,
+        default="FLCSphericalMask",
+        help="Template matching scoring function.",
+        choices=MATCHING_EXHAUSTIVE_REGISTER.keys(),
+    )
+    parser.add_argument(
+        "--ncores", type=int, help="Number of cores for parallelization.", required=True
+    )
+    parser.add_argument(
+        "--no_edge_padding",
+        dest="no_edge_padding",
+        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.",
+    )
+    args = parser.parse_args()
+    return 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)
+
+    template_box = template.shape
+    if args.no_fourier_padding:
+        template_box = np.ones(len(template_box), dtype=int)
+
+    result = estimate_ram_usage(
+        shape1=target_box,
+        shape2=template_box,
+        matching_method=args.score,
+        ncores=args.ncores,
+        analyzer_method="MaxScoreOverRotations",
+    )
+    print(result)
+
+
+if __name__ == "__main__":
+    main()