pytme 0.3.1__tar.gz → 0.3.1.post1__tar.gz

{pytme-0.3.1 → pytme-0.3.1.post1}/PKG-INFO

@@ -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 → pytme-0.3.1.post1}/pyproject.toml

@@ -7,7 +7,7 @@ name="pytme"
 authors = [
   { name = "Valentin Maurer", email = "valentin.maurer@embl-hamburg.de" },
 ]
-version="0.3.1"
+version="0.3.1.post1"
 description="Python Template Matching Engine"
 readme="README.md"
 requires-python = ">=3.11"
@@ -33,8 +33,7 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
-cupy = ["cupy-cuda12x>13.0.0"]
-cupy_voltools = ["cupy-cuda12x>13.0.0", "voltools"]
+cupy = ["cupy-cuda12x>13.0.0", "voltools"]
 pytorch = ["torch", "torchvision"]
 jax = ["jax[cuda12]", "jaxlib"]
 jax_cpu = ["jax", "jaxlib"]

{pytme-0.3.1 → pytme-0.3.1.post1}/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 → pytme-0.3.1.post1}/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)

{pytme-0.3.1 → pytme-0.3.1.post1}/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]
 

{pytme-0.3.1 → pytme-0.3.1.post1}/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),
         )
 
 

{pytme-0.3.1 → pytme-0.3.1.post1}/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
 

{pytme-0.3.1 → pytme-0.3.1.post1}/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:

{pytme-0.3.1 → pytme-0.3.1.post1}/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
 

{pytme-0.3.1 → pytme-0.3.1.post1}/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.

{pytme-0.3.1 → pytme-0.3.1.post1}/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)
 

{pytme-0.3.1 → pytme-0.3.1.post1}/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",

{pytme-0.3.1 → pytme-0.3.1.post1}/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(

{pytme-0.3.1 → pytme-0.3.1.post1}/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

{pytme-0.3.1 → pytme-0.3.1.post1}/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: