pytme 0.3.1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1.post1__cp311-cp311-macosx_15_0_arm64.whl

{pytme-0.3.1.data → pytme-0.3.1.post1.data}/scripts/postprocess.py

@@ -188,7 +188,7 @@ def parse_args():
     )
     additional_group.add_argument(
         "--n-false-positives",
-        type=int,
+        type=float,
         default=None,
         required=False,
         help="Number of accepted false-positives picks to determine minimum score.",
@@ -318,11 +318,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
         data = load_matching_output(foreground)
         scores, _, rotations, rotation_mapping, *_ = data
 
-        # We could normalize to unit sdev, but that might lead to unexpected
-        # results for flat background distributions
-        # scores -= scores.mean()
         indices = tuple(slice(0, x) for x in scores.shape)
-
         indices_update = scores > scores_out[indices]
         scores_out[indices][indices_update] = scores[indices_update]
 
@@ -369,9 +365,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
     scores_norm = np.full(out_shape_norm, fill_value=0, dtype=np.float32)
     for background in backgrounds:
         data_norm = load_matching_output(background)
-
-        scores = data_norm[0]
-        # scores -= scores.mean()
+        scores, _, rotations, rotation_mapping, *_ = data_norm
 
         indices = tuple(slice(0, x) for x in scores.shape)
         indices_update = scores > scores_norm[indices]
@@ -381,6 +375,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
     update = tuple(slice(0, int(x)) for x in np.minimum(out_shape, scores.shape))
     scores_out = np.full(out_shape, fill_value=0, dtype=np.float32)
     scores_out[update] = data[0][update] - scores_norm[update]
+    scores_out[update] += scores_norm[update].mean()
 
     # scores_out[update] = np.divide(scores_out[update], 1 - scores_norm[update])
     scores_out = np.fmax(scores_out, 0, out=scores_out)
@@ -485,8 +480,11 @@ def main():
     if orientations is None:
         translations, rotations, scores, details = [], [], [], []
 
-        # Data processed by normalize_input is guaranteed to have this shape
-        scores, offset, rotation_array, rotation_mapping, meta = data
+        var = None
+        # Data processed by normalize_input is guaranteed to have this shape)
+        scores, _, rotation_array, rotation_mapping, *_ = data
+        if len(data) == 6:
+            scores, _, rotation_array, rotation_mapping, var, *_ = data
 
         cropped_shape = np.subtract(
             scores.shape, np.multiply(args.min_boundary_distance, 2)
@@ -509,13 +507,16 @@ def main():
             )
             args.n_false_positives = max(args.n_false_positives, 1)
             n_correlations = np.size(scores[cropped_slice]) * len(rotation_mapping)
+            std = np.std(scores[cropped_slice])
+            if var is not None:
+                std = np.asarray(np.sqrt(var)).reshape(())
+
             minimum_score = np.multiply(
                 erfcinv(2 * args.n_false_positives / n_correlations),
-                np.sqrt(2) * np.std(scores[cropped_slice]),
+                np.sqrt(2) * std,
             )
-            print(f"Determined minimum score cutoff: {minimum_score}.")
-            minimum_score = max(minimum_score, 0)
-            args.min_score = minimum_score
+            print(f"Determined cutoff --min-score {minimum_score}.")
+            args.min_score = max(minimum_score, 0)
 
         args.batch_dims = None
         if hasattr(cli_args, "batch_dims"):

{pytme-0.3.1.dist-info → pytme-0.3.1.post1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytme
-Version: 0.3.1
+Version: 0.3.1.post1
 Summary: Python Template Matching Engine
 Author: Valentin Maurer
 Author-email: Valentin Maurer <valentin.maurer@embl-hamburg.de>
@@ -26,9 +26,7 @@ Requires-Dist: importlib_resources
 Requires-Dist: joblib
 Provides-Extra: cupy
 Requires-Dist: cupy-cuda12x>13.0.0; extra == "cupy"
-Provides-Extra: cupy-voltools
-Requires-Dist: cupy-cuda12x>13.0.0; extra == "cupy-voltools"
-Requires-Dist: voltools; extra == "cupy-voltools"
+Requires-Dist: voltools; extra == "cupy"
 Provides-Extra: pytorch
 Requires-Dist: torch; extra == "pytorch"
 Requires-Dist: torchvision; extra == "pytorch"

{pytme-0.3.1.dist-info → pytme-0.3.1.post1.dist-info}/RECORD

@@ -1,26 +1,26 @@
-pytme-0.3.1.data/scripts/estimate_memory_usage.py,sha256=Ry46LXUv3SZ0g41g3RDUg9UH6hiSnnG3mHTyaGletXE,2114
-pytme-0.3.1.data/scripts/match_template.py,sha256=X0Sa4dBuE3fvKer--JgY7kPSyZ7H6fa5PKi2lxRoyzM,37563
-pytme-0.3.1.data/scripts/postprocess.py,sha256=0QOUzqa76Wrg7JK51qEtjm7XH78UFdMuDia50HmlXwU,27638
-pytme-0.3.1.data/scripts/preprocess.py,sha256=eq-67cuj3WFx5jAdS56yQSVv_sCt7SRW-dqzQRbLnVE,6328
-pytme-0.3.1.data/scripts/preprocessor_gui.py,sha256=R73N-UTSiAplTqyRw-SLGcGusnJVv2dlAFLuqmUrl-8,44153
-pytme-0.3.1.data/scripts/pytme_runner.py,sha256=dqCd60puAOOOSvuCxrJC1MbfdsRS-ctMso5YIwM-JkI,40356
-pytme-0.3.1.dist-info/licenses/LICENSE,sha256=gXf5dRMhNSbfLPYYTY_5hsZ1r7UU1OaKQEAQUhuIBkM,18092
+pytme-0.3.1.post1.data/scripts/estimate_memory_usage.py,sha256=Ry46LXUv3SZ0g41g3RDUg9UH6hiSnnG3mHTyaGletXE,2114
+pytme-0.3.1.post1.data/scripts/match_template.py,sha256=X0Sa4dBuE3fvKer--JgY7kPSyZ7H6fa5PKi2lxRoyzM,37563
+pytme-0.3.1.post1.data/scripts/postprocess.py,sha256=GVlavaLWPdbTmMeEczsrG0cGhHXW4wZOZ5zLIV0sx4o,27708
+pytme-0.3.1.post1.data/scripts/preprocess.py,sha256=eq-67cuj3WFx5jAdS56yQSVv_sCt7SRW-dqzQRbLnVE,6328
+pytme-0.3.1.post1.data/scripts/preprocessor_gui.py,sha256=R73N-UTSiAplTqyRw-SLGcGusnJVv2dlAFLuqmUrl-8,44153
+pytme-0.3.1.post1.data/scripts/pytme_runner.py,sha256=dqCd60puAOOOSvuCxrJC1MbfdsRS-ctMso5YIwM-JkI,40356
+pytme-0.3.1.post1.dist-info/licenses/LICENSE,sha256=gXf5dRMhNSbfLPYYTY_5hsZ1r7UU1OaKQEAQUhuIBkM,18092
 scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 scripts/estimate_memory_usage.py,sha256=UiaX30o_H59vzTQZaXISxgzWj5jwgFbMoH1E5ydVHcw,2115
 scripts/eval.py,sha256=ebJVLxbRlB6TI5YHNr0VavZ4lmaRdf8QVafyiDhh_oU,2528
 scripts/extract_candidates.py,sha256=B2O4Xm0eVJzBZOblfkH5za8fTxLIdBRDl89Qwkq4Kjk,8097
 scripts/match_template.py,sha256=LsoS7gPh6tx3c5uyF3v5XbATjb8jJ_jh91zmO2IruZ0,37564
 scripts/match_template_filters.py,sha256=Gj4a1b_S5NWp_dfFEPFn0D7jGf-qYgBbnTvZZ4bwqOQ,42036
-scripts/postprocess.py,sha256=YNSySz76axtNvz7R5eCvS90eJjtKQQbWMLSexFTkEbg,27639
+scripts/postprocess.py,sha256=ysI1hBFqRly8noeOgL2ZONm27hSkPoVsbeG1kwI53G4,27709
 scripts/preprocess.py,sha256=PrtO0aWGbTvSMqCdNYuzW02FoGvSsvQzJab52DTssQ4,6329
 scripts/preprocessor_gui.py,sha256=eozUE6D_vJOO11Prardwm7zWrjvTXLHFaAMmY7yfNIc,44154
 scripts/pytme_runner.py,sha256=BD5u45vtzJo_d6ZAM2N_CKnGpfjtNg4v7C-31vbTGnU,40357
 scripts/refine_matches.py,sha256=1AG-E2PuTauflR_qVC1CaZKg7ev9xxGu3CIYeyJ2-WQ,12777
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/test_analyzer.py,sha256=HNZf12dZbuMTJ1oYhdDd1nn9NSGXV12ReEqRjNQ96UI,8393
+tests/test_analyzer.py,sha256=9TVJacQF44AnMYOWfLhqQgz1V2TeKjijVT1rJb50NMw,8360
 tests/test_backends.py,sha256=65jCvPfB2YjUHkhmjkC49Rpd1SyPRQ7_lnLXHI5QC-U,17361
 tests/test_density.py,sha256=JVQunV445qij5WdcpKn-5GKqT3endzjXvBPkIaXPADo,18914
-tests/test_extensions.py,sha256=1Zv9dG_dmmC2mlbX91YIPyGLSToPC0202-ffLAfVcr4,5203
+tests/test_extensions.py,sha256=K71MvMhLdZz66ifJYqRl0IEr7-PCXSMxKokg-4NCc8Q,5170
 tests/test_matching_cli.py,sha256=JF7LQixpFQMpbijoPwtNSe9FoYXZmBR6J3T5BvsLbR8,11202
 tests/test_matching_data.py,sha256=U6ISe4lBLDj-OzgA6QAaoO_aegCJjPtXqHbPiPZ2tkA,6091
 tests/test_matching_exhaustive.py,sha256=bRPCN0pyZk3DmXMrWRcEGAksYoch8P7fRiwE3g0yIf4,4039
@@ -65,13 +65,13 @@ tme/cli.py,sha256=48Q7QGIuCwNSms5wOnGmWjMSlZrXjAQUz_9zsIFT9tY,3652
 tme/density.py,sha256=iupy3gL7gwE2Q9FHcZfHY_AcbOSUAo3i9tSbC-KF-DA,82405
 tme/extensions.cpython-311-darwin.so,sha256=xUCjh1eMK3iKWMC04oFG3LzxBRAqHpl7idIIV4ezx3Q,415840
 tme/mask.py,sha256=4xHqWtOabgdYDDCXHpXflaZmLlmE2_9-bxjs4AaCIYM,10472
-tme/matching_data.py,sha256=l9eIJOkANOn-sHWeCVU-1RskMloqH0KmUFQcOQTt6qQ,30361
-tme/matching_exhaustive.py,sha256=YswFCaRTaPaGQVRf7SgrmsqEHlnC83u8G4y-PIlR-8w,18270
+tme/matching_data.py,sha256=y6UBUykYpZHsl5NEDyKthKztuwcKFIaOIlB0eEh-Sww,31255
+tme/matching_exhaustive.py,sha256=DnWl7zBfnK83MFZzLDQ9WsBvI2mZkTikxpiXIWrdmxQ,18385
 tme/matching_optimization.py,sha256=r2zMyOKyvoP8OX9xValJE0DLWbLpNdyY1MLWcOi-H1U,45715
-tme/matching_scores.py,sha256=ZYw7ff0y0Zs6fzUwTCQ34Gth4Ckx2ZV7qNk1bH4NXdw,34849
+tme/matching_scores.py,sha256=sX4hMki9WjwgenNWJ1wNu04J808wOgGuVqbRYSA2TB8,35860
 tme/matching_utils.py,sha256=Qn65B2F8v1yHLetqpiirvUbj-A9-J1giq3lE9cSb-1E,27426
 tme/memory.py,sha256=lOD_XJtAWo3IFG5fJIdYgLSpn8o0ORO6Vua_tlm5YEk,9683
-tme/orientations.py,sha256=93nmt4GPiuBsMWfwSuQVHUBrgY-5k53PlaCaDXZCZjc,21345
+tme/orientations.py,sha256=xMVC5FWDNYSoiIpRXqTi3K1QiSsH2ilUoDEUqEJLYzk,21642
 tme/parser.py,sha256=d59A7meIUJ8OzfB6eaVqd7CwLj9oYnOfubqlvygbu1U,24210
 tme/preprocessor.py,sha256=7DAGRfCPc9ge-hEuYsNA_KCeajVnpWl-w4NzGQS53GM,40377
 tme/rotations.py,sha256=wVnvZrRCHb1wd5E6yDsLxiP2G0Oin66A4W14vsz4i50,10641
@@ -79,19 +79,19 @@ tme/structure.py,sha256=ojk-IQnIe1JrzxUVHlNlyLxtwzNkNykvxTaXvyBMmbs,72135
 tme/types.py,sha256=NAY7C4qxE6yz-DXVtClMvFfoOV-spWGLNfpLATZ1LcU,442
 tme/analyzer/__init__.py,sha256=aJyVGlQwO6Ij-_NZ5rBlgBQQSIj0dqpUwlmvIvnQqBM,89
 tme/analyzer/_utils.py,sha256=48Xq2Hi_P3VYl1SozYz3WwEj3BwL1QdtFYx7dTxv0bc,6344
-tme/analyzer/aggregation.py,sha256=tl6Vw0VhdhPYwb-Y82eAK83odL2JD0d5C3N39qJRVZM,27875
+tme/analyzer/aggregation.py,sha256=3ieZlNcDp0dhnQ85DdVLsCFrRfqU4jim3JRcMwORF_U,28518
 tme/analyzer/base.py,sha256=2uEYKa0xLwggJUenAA6-l8DIO5sYtklHMg9dR3DzAOI,3887
 tme/analyzer/peaks.py,sha256=y40YVa2zJFLWD3aoP0xtuhK8yOrQV0AVH6ElWyVjXkM,32587
 tme/analyzer/proxy.py,sha256=NH0J93ESl5NAVd_22rLZGDWjCx1ov8I995Fr3mYqUSM,4162
 tme/backends/__init__.py,sha256=VUvEnpfjmZQVn_9QHkLPBoa6y8M3gDyNNL9kq4OQNGE,5249
 tme/backends/_cupy_utils.py,sha256=scxCSK9BcsDaIbHR2mFxAJRIxdc0r5JTC1gCJmVJa64,24558
-tme/backends/_jax_utils.py,sha256=0rybsyG2QbGDdQg-1ruDIt1F15DZcz_TguZYjbepCA4,5838
+tme/backends/_jax_utils.py,sha256=UrC0RwwxwnNdz6QusPOgC8eISTX_C1KbZmt7Karl_Nk,6096
 tme/backends/_numpyfftw_utils.py,sha256=RnYgbSs_sL-NlwYGRqU9kziq0U5qzl9GecVx9uUeSJs,7904
-tme/backends/cupy_backend.py,sha256=G3rsdmxWTes8pM82TTnWvoq1bwGrNsBTbfJMSc5gfTA,8269
-tme/backends/jax_backend.py,sha256=eHN0PvWVpjoKaYUJ0oRcdXwF7muKrgVAUzb96QRDXWI,12069
-tme/backends/matching_backend.py,sha256=oWARMu_Mt6apnvpYGCbQ1_x7tl_PZoxVtVb8HJgtQ-o,32953
+tme/backends/cupy_backend.py,sha256=Ms2sObxr0xc_tdHctcL659G8YWOTqPG59F3665yZH4c,8163
+tme/backends/jax_backend.py,sha256=d14u__H9leb5aXCz383T1xjTrGOzxgfWxXjgRrRWZMI,12676
+tme/backends/matching_backend.py,sha256=B9LXxKpfNYdJBIksyGm7g7V7SC9G3mrluscz3aqISsU,33199
 tme/backends/mlx_backend.py,sha256=3aAKv7VkK_jfxwDMPQZsLp_lRU8MBJ6zNAzrs6VZd3s,6718
-tme/backends/npfftw_backend.py,sha256=P151VbhML5myCMX1o4ArHYRNz_Dz6XoGVyQIesEK2zI,18698
+tme/backends/npfftw_backend.py,sha256=8-Hs3slqnWvsHpqICxKxb-fbmg3IbSok3sThXew-P7M,18800
 tme/backends/pytorch_backend.py,sha256=yn1hhowUMmzMLHdh2jqw2B0CvBIVcSAWOvRMMDkSHmw,14229
 tme/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tme/data/c48n309.npy,sha256=NwH64mOEbm3tStq5c98o81fY1vMOoq4nvXDAh7Z7iZg,296768
@@ -126,8 +126,8 @@ tme/filters/ctf.py,sha256=HNB03Pw3gn2Y8pIM2391iRHQzp2NcHZWFDC61v3K0oM,24520
 tme/filters/reconstruction.py,sha256=4gCGpuWuM-yCDQvuzyex_Q8bmsFqUI_8IoZnv3B_3UA,7564
 tme/filters/wedge.py,sha256=JWiN9IwRXnc7A7syq11kSuin6QjhlVmi38nwGHu4yxE,19541
 tme/filters/whitening.py,sha256=Zwk-0jMVqy_F8TdQ0ht5unMz2JBOePOC3okUhQpU0bo,6348
-pytme-0.3.1.dist-info/METADATA,sha256=uGvw-jEiKmPKGI-cQurxjdCgvwrRhcCLZzbIQmF2NqM,5079
-pytme-0.3.1.dist-info/WHEEL,sha256=SPbiHAOPLnBtml4sk5MIwvWF6YnGfOfy9C__W6Bpeg4,109
-pytme-0.3.1.dist-info/entry_points.txt,sha256=pbUSmB0J4Ghlg0w7jHfaFSvPMuvRWzeSuvdjdRPisAU,288
-pytme-0.3.1.dist-info/top_level.txt,sha256=ovCUR7UXXouH3zYt_fJLoqr_vtjp1wudFgjVAnztQLE,18
-pytme-0.3.1.dist-info/RECORD,,
+pytme-0.3.1.post1.dist-info/METADATA,sha256=2EXMjS9Dx7TEo0gDEgGjvWD7rKnhJ0eTTb4AalGW9kE,4985
+pytme-0.3.1.post1.dist-info/WHEEL,sha256=SPbiHAOPLnBtml4sk5MIwvWF6YnGfOfy9C__W6Bpeg4,109
+pytme-0.3.1.post1.dist-info/entry_points.txt,sha256=pbUSmB0J4Ghlg0w7jHfaFSvPMuvRWzeSuvdjdRPisAU,288
+pytme-0.3.1.post1.dist-info/top_level.txt,sha256=ovCUR7UXXouH3zYt_fJLoqr_vtjp1wudFgjVAnztQLE,18
+pytme-0.3.1.post1.dist-info/RECORD,,

scripts/postprocess.py

@@ -188,7 +188,7 @@ def parse_args():
     )
     additional_group.add_argument(
         "--n-false-positives",
-        type=int,
+        type=float,
         default=None,
         required=False,
         help="Number of accepted false-positives picks to determine minimum score.",
@@ -318,11 +318,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
         data = load_matching_output(foreground)
         scores, _, rotations, rotation_mapping, *_ = data
 
-        # We could normalize to unit sdev, but that might lead to unexpected
-        # results for flat background distributions
-        # scores -= scores.mean()
         indices = tuple(slice(0, x) for x in scores.shape)
-
         indices_update = scores > scores_out[indices]
         scores_out[indices][indices_update] = scores[indices_update]
 
@@ -369,9 +365,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
     scores_norm = np.full(out_shape_norm, fill_value=0, dtype=np.float32)
     for background in backgrounds:
         data_norm = load_matching_output(background)
-
-        scores = data_norm[0]
-        # scores -= scores.mean()
+        scores, _, rotations, rotation_mapping, *_ = data_norm
 
         indices = tuple(slice(0, x) for x in scores.shape)
         indices_update = scores > scores_norm[indices]
@@ -381,6 +375,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
     update = tuple(slice(0, int(x)) for x in np.minimum(out_shape, scores.shape))
     scores_out = np.full(out_shape, fill_value=0, dtype=np.float32)
     scores_out[update] = data[0][update] - scores_norm[update]
+    scores_out[update] += scores_norm[update].mean()
 
     # scores_out[update] = np.divide(scores_out[update], 1 - scores_norm[update])
     scores_out = np.fmax(scores_out, 0, out=scores_out)
@@ -485,8 +480,11 @@ def main():
     if orientations is None:
         translations, rotations, scores, details = [], [], [], []
 
-        # Data processed by normalize_input is guaranteed to have this shape
-        scores, offset, rotation_array, rotation_mapping, meta = data
+        var = None
+        # Data processed by normalize_input is guaranteed to have this shape)
+        scores, _, rotation_array, rotation_mapping, *_ = data
+        if len(data) == 6:
+            scores, _, rotation_array, rotation_mapping, var, *_ = data
 
         cropped_shape = np.subtract(
             scores.shape, np.multiply(args.min_boundary_distance, 2)
@@ -509,13 +507,16 @@ def main():
             )
             args.n_false_positives = max(args.n_false_positives, 1)
             n_correlations = np.size(scores[cropped_slice]) * len(rotation_mapping)
+            std = np.std(scores[cropped_slice])
+            if var is not None:
+                std = np.asarray(np.sqrt(var)).reshape(())
+
             minimum_score = np.multiply(
                 erfcinv(2 * args.n_false_positives / n_correlations),
-                np.sqrt(2) * np.std(scores[cropped_slice]),
+                np.sqrt(2) * std,
             )
-            print(f"Determined minimum score cutoff: {minimum_score}.")
-            minimum_score = max(minimum_score, 0)
-            args.min_score = minimum_score
+            print(f"Determined cutoff --min-score {minimum_score}.")
+            args.min_score = max(minimum_score, 0)
 
         args.batch_dims = None
         if hasattr(cli_args, "batch_dims"):

tests/test_analyzer.py

@@ -165,7 +165,6 @@ class TestMaxScoreOverRotations:
         assert res[0].dtype == be._float_dtype
         assert res[1].size == self.data.ndim
         assert np.allclose(res[2].shape, self.data.shape)
-        assert len(res) == 4
 
     @pytest.mark.parametrize("use_memmap", [False, True])
     @pytest.mark.parametrize("score_threshold", [0, 1e10, -1e10])
@@ -181,7 +180,7 @@ class TestMaxScoreOverRotations:
 
         data2 = self.data * 2
         score_analyzer(state, data2, rotation_matrix=self.rotation_matrix)
-        scores, translation_offset, rotations, mapping = score_analyzer.result(state)
+        scores, offset, rotations, mapping, *_ = score_analyzer.result(state)
 
         assert np.all(scores >= score_threshold)
         max_scores = np.maximum(self.data, data2)
@@ -214,7 +213,7 @@ class TestMaxScoreOverRotations:
         ret = MaxScoreOverRotations.merge(
             results=states, use_memmap=use_memmap, score_threshold=score_threshold
         )
-        scores, translation, rotations, mapping = ret
+        scores, translation, rotations, mapping, *_ = ret
         assert np.all(scores >= score_threshold)
         max_scores = np.maximum(self.data, data2)
         max_scores = np.maximum(max_scores, score_threshold)

tests/test_extensions.py

@@ -53,7 +53,6 @@ class TestExtensions:
     @pytest.mark.parametrize("min_distance", [0, 5, 10])
     def test_find_candidate_indices(self, dimension, dtype, min_distance):
         coordinates = COORDINATES[dimension].astype(dtype)
-        print(coordinates.shape)
 
         min_distance = np.array([min_distance]).astype(dtype)[0]
 

tme/analyzer/aggregation.py

@@ -132,12 +132,14 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             - scores : BackendArray of shape `self._shape` filled with `score_threshold`.
             - rotations : BackendArray of shape `self._shape` filled with -1.
             - rotation_mapping : dict, empty mapping from rotation bytes to indices.
+            - ssum : BackendArray, accumulator for sum of squared scores.
         """
         scores = be.full(
             shape=self._shape, dtype=be._float_dtype, fill_value=self._score_threshold
         )
         rotations = be.full(self._shape, dtype=be._int_dtype, fill_value=-1)
-        return scores, rotations, {}
+        ssum = be.full((1), dtype=be._float_dtype, fill_value=0)
+        return scores, rotations, {}, ssum
 
     def __call__(
         self,
@@ -156,6 +158,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             - scores : BackendArray, current maximum scores.
             - rotations : BackendArray, current rotation indices.
             - rotation_mapping : dict, mapping from rotation bytes to indices.
+            - ssum : BackendArray, accumulator for sum of squared scores.
         scores : BackendArray
             Array of new scores to update analyzer with.
         rotation_matrix : BackendArray
@@ -168,7 +171,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         # be.tobytes behaviour caused overhead for certain GPU/CUDA combinations
         # If the analyzer is not shared and each rotation is unique, we can
         # use index to rotation mapping and invert prior to merging.
-        prev_scores, rotations, rotation_mapping = state
+        prev_scores, rotations, rotation_mapping, ssum = state
 
         rotation_index = len(rotation_mapping)
         rotation_matrix = be.astype(rotation_matrix, be._float_dtype)
@@ -180,13 +183,14 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             rotation = be.tobytes(rotation_matrix)
             rotation_index = rotation_mapping.setdefault(rotation, rotation_index)
 
+        ssum = be.add(ssum, be.ssum(scores), out=ssum)
         scores, rotations = be.max_score_over_rotations(
             scores=scores,
             max_scores=prev_scores,
             rotations=rotations,
             rotation_index=rotation_index,
         )
-        return scores, rotations, rotation_mapping
+        return scores, rotations, rotation_mapping, ssum
 
     @staticmethod
     def _invert_rmap(rotation_mapping: dict) -> dict:
@@ -224,6 +228,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             - scores : BackendArray, current maximum scores.
             - rotations : BackendArray, current rotation indices.
             - rotation_mapping : dict, mapping from rotation indices to matrices.
+            - ssum : BackendArray, accumulator for sum of squared scores.
         targetshape : Tuple[int], optional
             Shape of the target for convolution mode correction.
         templateshape : Tuple[int], optional
@@ -240,9 +245,9 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         Returns
         -------
         tuple
-            Final result tuple (scores, offset, rotations, rotation_mapping).
+            Final result tuple (scores, offset, rotations, rotation_mapping, ssum).
         """
-        scores, rotations, rotation_mapping = state
+        scores, rotations, rotation_mapping, ssum = state
 
         # Apply postprocessing if parameters are provided
         if fourier_shift is not None:
@@ -269,11 +274,13 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         if self._inversion_mapping:
             rotation_mapping = {be.tobytes(v): k for k, v in rotation_mapping.items()}
 
+        n_rotations = max(len(rotation_mapping), 1)
         return (
             scores,
             be.to_numpy_array(self._offset),
             rotations,
             self._invert_rmap(rotation_mapping),
+            be.to_numpy_array(ssum) / (scores.size * n_rotations),
         )
 
     def _harmonize_states(states: List[Tuple]):
@@ -287,18 +294,18 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             if states[i] is None:
                 continue
 
-            scores, offset, rotations, rotation_mapping = states[i]
+            scores, offset, rotations, rotation_mapping, ssum = states[i]
             if out_shape is None:
                 out_shape = np.zeros(scores.ndim, int)
             out_shape = np.maximum(out_shape, np.add(offset, scores.shape))
 
             new_param = {}
             for key, value in rotation_mapping.items():
-                rotation_bytes = be.tobytes(value)
+                rotation_bytes = np.asarray(value).tobytes()
                 new_param[rotation_bytes] = key
                 if rotation_bytes not in new_rotation_mapping:
                     new_rotation_mapping[rotation_bytes] = len(new_rotation_mapping)
-            states[i] = (scores, offset, rotations, new_param)
+            states[i] = (scores, offset, rotations, new_param, ssum)
         out_shape = tuple(int(x) for x in out_shape)
         return new_rotation_mapping, out_shape, states
 
@@ -329,11 +336,10 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         if len(results) == 1:
             ret = results[0]
             if use_memmap:
-                scores, offset, rotations, rotation_mapping = ret
+                scores, offset, rotations, rotation_mapping, ssum = ret
                 scores = array_to_memmap(scores)
                 rotations = array_to_memmap(rotations)
-                ret = (scores, offset, rotations, rotation_mapping)
-
+                ret = (scores, offset, rotations, rotation_mapping, ssum)
             return ret
 
         # Determine output array shape and create consistent rotation map
@@ -368,6 +374,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             )
             rotations_out = np.full(out_shape, fill_value=-1, dtype=rotations_dtype)
 
+        total_ssum = 0
         for i in range(len(results)):
             if results[i] is None:
                 continue
@@ -385,7 +392,9 @@ class MaxScoreOverRotations(AbstractAnalyzer):
                     shape=out_shape,
                     dtype=rotations_dtype,
                 )
-            scores, offset, rotations, rotation_mapping = results[i]
+            scores, offset, rotations, rotation_mapping, ssum = results[i]
+
+            total_ssum = np.add(total_ssum, ssum)
             stops = np.add(offset, scores.shape).astype(int)
             indices = tuple(slice(*pos) for pos in zip(offset, stops))
 
@@ -428,6 +437,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
             np.zeros(scores_out.ndim, dtype=int),
             rotations_out,
             cls._invert_rmap(master_rotation_mapping),
+            total_ssum / len(results),
         )
 
 

tme/backends/_jax_utils.py

@@ -112,10 +112,14 @@ def _identity(arr: BackendArray, arr_filter: BackendArray) -> BackendArray:
     return arr
 
 
+def _mask_scores(arr, mask):
+    return arr.at[:].multiply(mask)
+
+
 @partial(
     pmap,
-    in_axes=(0,) + (None,) * 6,
-    static_broadcasted_argnums=[6, 7, 8, 9],
+    in_axes=(0,) + (None,) * 7,
+    static_broadcasted_argnums=[7, 8, 9, 10],
     axis_name="batch",
 )
 def scan(
@@ -125,6 +129,7 @@ def scan(
     rotations: BackendArray,
     template_filter: BackendArray,
     target_filter: BackendArray,
+    score_mask: BackendArray,
     fast_shape: Tuple[int],
     rotate_mask: bool,
     analyzer_class: object,
@@ -159,6 +164,10 @@ def scan(
     if template_filter.shape != ():
         _template_filter_func = _apply_fourier_filter
 
+    _score_mask_func = _identity
+    if score_mask.shape != ():
+        _score_mask_func = _mask_scores
+
     def _sample_transform(ret, rotation_matrix):
         state, index = ret
         template_rot, template_mask_rot = be.rigid_transform(
@@ -185,6 +194,8 @@ def scan(
             n_observations=n_observations,
             eps=eps,
         )
+        scores = _score_mask_func(scores, score_mask)
+
         state = analyzer(state, scores, rotation_matrix, rotation_index=index)
         return (state, index + 1), None
 

tme/backends/cupy_backend.py

@@ -81,6 +81,17 @@ class CupyBackend(NumpyFFTWBackend):
             """,
             "norm_scores",
         )
+
+        # Sum of square computation similar to the demeaned variance in pytom
+        self.ssum = cp.ReductionKernel(
+            f"{ftype} arr",
+            f"{ftype} ret",
+            "arr * arr",
+            "a + b",
+            "ret = a",
+            "0",
+            f"ssum_{ftype}",
+        )
         self.texture_available = find_spec("voltools") is not None
 
     def to_backend_array(self, arr: NDArray) -> CupyArray:
@@ -139,17 +150,6 @@ class CupyBackend(NumpyFFTWBackend):
         peaks = self._array_backend.array(self._array_backend.nonzero(max_filter)).T
         return peaks
 
-    # The default methods in Cupy were oddly slow
-    def var(self, a, *args, **kwargs):
-        out = a - self._array_backend.mean(a, *args, **kwargs)
-        self._array_backend.square(out, out)
-        out = self._array_backend.mean(out, *args, **kwargs)
-        return out
-
-    def std(self, a, *args, **kwargs):
-        out = self.var(a, *args, **kwargs)
-        return self._array_backend.sqrt(out)
-
     def _get_texture(self, arr: CupyArray, order: int = 3, prefilter: bool = False):
         key = id(arr)
         if key in TEXTURE_CACHE:

tme/backends/jax_backend.py

@@ -197,6 +197,13 @@ class JaxBackend(NumpyFFTWBackend):
     def _tuple_to_dict(self, data: Tuple) -> Dict:
         return {x[0]: self._from_hashable(*x[1]) for x in data}
 
+    def _unbatch(self, data, target_ndim, index):
+        if not isinstance(data, type(self.zeros(1))):
+            return data
+        elif data.ndim <= target_ndim:
+            return data
+        return data[index]
+
     def scan(
         self,
         matching_data: type,
@@ -212,11 +219,13 @@ class JaxBackend(NumpyFFTWBackend):
         :py:class:`tme.analyzer.MaxScoreOverRotations`.
         """
         from ._jax_utils import scan as scan_inner
+        from ..analyzer import MaxScoreOverRotations
 
         pad_target = True if len(splits) > 1 else False
         convolution_mode = "valid" if pad_target else "same"
         target_pad = matching_data.target_padding(pad_target=pad_target)
 
+        score_mask = self.full((1,), fill_value=1, dtype=bool)
         target_shape = tuple(
             (x.stop - x.start + p) for x, p in zip(splits[0][0], target_pad)
         )
@@ -273,6 +282,9 @@ class JaxBackend(NumpyFFTWBackend):
 
                 analyzer_kwargs.append(self._dict_to_tuple(cur_args))
 
+                if pad_target:
+                    score_mask = base._score_mask(fast_shape, shift)
+
                 _target = self.astype(base._target, self._float_dtype)
                 translation_offsets.append(translation_offset)
                 targets.append(self.topleft_pad(_target, fast_shape))
@@ -306,17 +318,23 @@ class JaxBackend(NumpyFFTWBackend):
                 matching_data.rotations,
                 template_filter,
                 target_filter,
+                score_mask,
                 fast_shape,
                 rotate_mask,
                 callback_class,
                 analyzer_kwargs,
             )
 
+            ndim = targets.ndim - 1
             for index in range(targets.shape[0]):
                 kwargs = self._tuple_to_dict(analyzer_kwargs[index])
                 analyzer = callback_class(**kwargs)
 
-                state = (states[0][index], states[1][index], rotation_mapping)
+                state = [self._unbatch(x, ndim, index) for x in states]
+
+                if isinstance(analyzer, MaxScoreOverRotations):
+                    state[2] = rotation_mapping
+
                 ret.append(analyzer.result(state, **kwargs))
         return ret
 

tme/backends/matching_backend.py

@@ -863,6 +863,17 @@ class MatchingBackend(ABC):
             Indices of ``k`` largest elements in ``arr``.
         """
 
+    @abstractmethod
+    def ssum(self, arr, *args, **kwargs) -> BackendArray:
+        """
+        Compute the sum of squares of ``arr``.
+
+        Returns
+        -------
+        BackendArray
+            Sum of squares with shape ().
+        """
+
     def indices(self, *args, **kwargs) -> BackendArray:
         """
         Creates an array representing the index grid of an input.

tme/backends/npfftw_backend.py

@@ -201,6 +201,9 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         sorted_indices = self.unravel_index(indices=sorted_indices, shape=arr.shape)
         return sorted_indices
 
+    def ssum(self, arr, *args, **kwargs):
+        return self.sum(self.square(arr), *args, **kwargs)
+
     def indices(self, *args, **kwargs) -> NDArray:
         return self._array_backend.indices(*args, **kwargs)
 

tme/matching_data.py

@@ -544,6 +544,30 @@ class MatchingData:
             batch_mask=be.to_numpy_array(self._batch_mask),
         )
 
+    def _score_mask(self, fast_shape: Tuple[int], shift: Tuple[int]) -> BackendArray:
+        """
+        Create a boolean mask to exclude scores derived from padding in template matching.
+        """
+        padding = self.target_padding(True)
+        offset = tuple(x // 2 for x in padding)
+        shape = tuple(y - x for x, y in zip(padding, self.target.shape))
+
+        subset = []
+        for i in range(len(offset)):
+            if self._batch_mask[i]:
+                subset.append(slice(None))
+            else:
+                subset.append(slice(offset[i], offset[i] + shape[i]))
+
+        score_mask = np.zeros(fast_shape, dtype=bool)
+        score_mask[tuple(subset)] = 1
+        score_mask = np.roll(
+            score_mask,
+            shift=tuple(-x for x in shift),
+            axis=tuple(i for i in range(len(shift))),
+        )
+        return be.to_backend_array(score_mask)
+
     def computation_schedule(
         self,
         matching_method: str = "FLCSphericalMask",

tme/matching_exhaustive.py

@@ -223,6 +223,10 @@ def scan(
     )
     conv, fwd, inv, shift = matching_data.fourier_padding()
 
+    score_mask = be.full(shape=(1,), fill_value=1, dtype=bool)
+    if pad_target:
+        score_mask = matching_data._score_mask(fwd, shift)
+
     template_filter = _setup_template_filter_apply_target_filter(
         matching_data=matching_data,
         fast_shape=fwd,
@@ -275,6 +279,7 @@ def scan(
             callback=callback_classes[index % n_callback_classes],
             interpolation_order=interpolation_order,
             template_filter=be.to_sharedarr(template_filter, shm_handler),
+            score_mask=be.to_sharedarr(score_mask, shm_handler),
             **setup,
         )
         for index, rotation in enumerate(matching_data._split_rotations_on_jobs(n_jobs))
@@ -420,8 +425,6 @@ def scan_subsets(
     outer_jobs, inner_jobs = job_schedule
     if be._backend_name == "jax":
         func = be.scan
-        if kwargs.get("projection_matching", False):
-            func = be.scan_projections
 
         corr_scoring = MATCHING_EXHAUSTIVE_REGISTER.get("CORR", (None, None))[1]
         results = func(

tme/matching_scores.py

@@ -356,6 +356,7 @@ def corr_scoring(
     callback: CallbackClass,
     interpolation_order: int,
     template_mask: shm_type = None,
+    score_mask: shm_type = None,
 ) -> CallbackClass:
     """
     Calculates a normalized cross-correlation between a target f and a template g.
@@ -394,6 +395,8 @@ def corr_scoring(
         Spline order for template rotations.
     template_mask : Union[Tuple[type, tuple of ints, type], BackendArray], optional
         Template mask data buffer, its shape and datatype, None by default.
+    score_mask : Union[Tuple[type, tuple of ints, type], BackendArray], optional
+        Score mask data buffer, its shape and datatype, None by default.
 
     Returns
     -------
@@ -404,6 +407,7 @@ def corr_scoring(
     inv_denominator = be.from_sharedarr(inv_denominator)
     numerator = be.from_sharedarr(numerator)
     template_filter = be.from_sharedarr(template_filter)
+    score_mask = be.from_sharedarr(score_mask)
 
     n_obs = None
     if template_mask is not None:
@@ -413,6 +417,7 @@ def corr_scoring(
     norm_template = conditional_execute(normalize_template, n_obs is not None)
     norm_sub = conditional_execute(be.subtract, numerator.shape != (1,))
     norm_mul = conditional_execute(be.multiply, inv_denominator.shape != (1))
+    norm_mask = conditional_execute(be.multiply, score_mask.shape != (1,))
 
     arr = be.zeros(fast_shape, be._float_dtype)
     ft_temp = be.zeros(fast_ft_shape, be._complex_dtype)
@@ -447,6 +452,8 @@ def corr_scoring(
 
         arr = norm_sub(arr, numerator, out=arr)
         arr = norm_mul(arr, inv_denominator, out=arr)
+        arr = norm_mask(arr, score_mask, out=arr)
+
         callback(arr, rotation_matrix=rotation)
 
     return callback
@@ -463,6 +470,7 @@ def flc_scoring(
     rotations: BackendArray,
     callback: CallbackClass,
     interpolation_order: int,
+    score_mask: shm_type = None,
 ) -> CallbackClass:
     """
     Computes a normalized cross-correlation between ``target`` (f),
@@ -522,6 +530,7 @@ def flc_scoring(
     ft_target = be.from_sharedarr(ft_target)
     ft_target2 = be.from_sharedarr(ft_target2)
     template_filter = be.from_sharedarr(template_filter)
+    score_mask = be.from_sharedarr(score_mask)
 
     arr = be.zeros(fast_shape, float_dtype)
     temp = be.zeros(fast_shape, float_dtype)
@@ -532,6 +541,7 @@ def flc_scoring(
     template_mask_rot = be.zeros(template.shape, be._float_dtype)
 
     tmpl_filter_func = _create_filter_func(template.shape, template_filter.shape)
+    norm_mask = conditional_execute(be.multiply, score_mask.shape != (1,))
 
     eps = be.eps(float_dtype)
     center = be.divide(be.to_backend_array(template.shape) - 1, 2)
@@ -567,6 +577,8 @@ def flc_scoring(
         arr = _correlate_fts(ft_target, ft_temp, ft_temp, arr, fast_shape)
 
         arr = be.norm_scores(arr, temp2, temp, n_obs, eps, arr)
+        arr = norm_mask(arr, score_mask, out=arr)
+
         callback(arr, rotation_matrix=rotation)
 
     return callback
@@ -585,6 +597,7 @@ def mcc_scoring(
     callback: CallbackClass,
     interpolation_order: int,
     overlap_ratio: float = 0.3,
+    score_mask: shm_type = None,
 ) -> CallbackClass:
     """
     Computes a normalized cross-correlation score between ``target`` (f),
@@ -755,12 +768,14 @@ def flc_scoring2(
     rotations: BackendArray,
     callback: CallbackClass,
     interpolation_order: int,
+    score_mask: shm_type = None,
 ) -> CallbackClass:
     template = be.from_sharedarr(template)
     template_mask = be.from_sharedarr(template_mask)
     ft_target = be.from_sharedarr(ft_target)
     ft_target2 = be.from_sharedarr(ft_target2)
     template_filter = be.from_sharedarr(template_filter)
+    score_mask = be.from_sharedarr(score_mask)
 
     tar_batch, tmpl_batch = _get_batch_dim(ft_target, template)
 
@@ -785,6 +800,7 @@ def flc_scoring2(
         filter_shape=template_filter.shape,
         arr_padded=True,
     )
+    norm_mask = conditional_execute(be.multiply, score_mask.shape != (1,))
 
     eps = be.eps(be._float_dtype)
     for index in range(rotations.shape[0]):
@@ -816,6 +832,8 @@ def flc_scoring2(
         arr = _correlate_fts(ft_target, ft_temp, ft_denom, arr, shape, axes)
 
         arr = be.norm_scores(arr, temp2, temp, n_obs, eps, arr)
+        arr = norm_mask(arr, score_mask, out=arr)
+
         callback(arr, rotation_matrix=rotation)
 
     return callback
@@ -834,12 +852,14 @@ def corr_scoring2(
     interpolation_order: int,
     target_filter: shm_type = None,
     template_mask: shm_type = None,
+    score_mask: shm_type = None,
 ) -> CallbackClass:
     template = be.from_sharedarr(template)
     ft_target = be.from_sharedarr(ft_target)
     inv_denominator = be.from_sharedarr(inv_denominator)
     numerator = be.from_sharedarr(numerator)
     template_filter = be.from_sharedarr(template_filter)
+    score_mask = be.from_sharedarr(score_mask)
 
     tar_batch, tmpl_batch = _get_batch_dim(ft_target, template)
 
@@ -869,6 +889,7 @@ def corr_scoring2(
     norm_template = conditional_execute(normalize_template, n_obs is not None)
     norm_sub = conditional_execute(be.subtract, numerator.shape != (1,))
     norm_mul = conditional_execute(be.multiply, inv_denominator.shape != (1,))
+    norm_mask = conditional_execute(be.multiply, score_mask.shape != (1,))
 
     template_filter_func = _create_filter_func(
         arr_shape=template.shape,
@@ -896,6 +917,8 @@ def corr_scoring2(
 
         arr = norm_sub(arr, numerator, out=arr)
         arr = norm_mul(arr, inv_denominator, out=arr)
+        arr = norm_mask(arr, score_mask, out=arr)
+
         callback(arr, rotation_matrix=rotation)
 
     return callback

tme/orientations.py

@@ -327,11 +327,18 @@ class Orientations:
             "_rlnAnglePsi",
             "_rlnClassNumber",
         ]
+
+        target_identifer = "_rlnMicrographName"
+        if version == "# version 50001":
+            header[3] = "_rlnCenteredCoordinateXAngst"
+            header[4] = "_rlnCenteredCoordinateYAngst"
+            header[5] = "_rlnCenteredCoordinateZAngst"
+            target_identifer = "_rlnTomoName"
+
         if source_path is not None:
-            header.append("_rlnMicrographName")
+            header.append(target_identifer)
 
         header.append("_pytmeScore")
-
         header = "\n".join(header)
         with open(filename, mode="w", encoding="utf-8") as ofile:
             if version is not None: