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

tme/density.py

@@ -2196,7 +2196,7 @@ class Density:
 
         Parameters
         ----------
-        target : Density
+        target : :py:class:`Density`
             The target map for template matching.
         template : Structure
             The template that should be aligned to the target.
@@ -2259,3 +2259,60 @@ class Density:
         coordinates = np.array(np.where(data > 0))
         weights = self.data[tuple(coordinates)]
         return align_to_axis(coordinates.T, weights=weights, axis=axis, flip=flip)
+
+    @staticmethod
+    def fourier_shell_correlation(density1: "Density", density2: "Density") -> NDArray:
+        """
+        Computes the Fourier Shell Correlation (FSC) between two instances of `Density`.
+
+        The Fourier transforms of the input maps are divided into shells
+        based on their spatial frequency. The correlation between corresponding shells
+        in the two maps is computed to give the FSC.
+
+        Parameters
+        ----------
+        density1 : :py:class:`Density`
+            Reference for comparison.
+        density2 : :py:class:`Density`
+            Target for comparison.
+
+        Returns
+        -------
+        NDArray
+            An array of shape (N, 2), where N is the number of shells.
+            The first column represents the spatial frequency for each shell
+            and the second column represents the corresponding FSC.
+
+        References
+        ----------
+        .. [1] https://github.com/tdgrant1/denss/blob/master/saxstats/saxstats.py
+        """
+        side = density1.data.shape[0]
+        df = 1.0 / side
+
+        qx_ = np.fft.fftfreq(side) * side * df
+        qx, qy, qz = np.meshgrid(qx_, qx_, qx_, indexing="ij")
+        qr = np.sqrt(qx**2 + qy**2 + qz**2)
+
+        qmax = np.max(qr)
+        qstep = np.min(qr[qr > 0])
+        nbins = int(qmax / qstep)
+        qbins = np.linspace(0, nbins * qstep, nbins + 1)
+        qbin_labels = np.searchsorted(qbins, qr, "right") - 1
+
+        F1 = np.fft.fftn(density1.data)
+        F2 = np.fft.fftn(density2.data)
+
+        qbin_labels = qbin_labels.reshape(-1)
+        numerator = np.bincount(
+            qbin_labels, weights=np.real(F1 * np.conj(F2)).reshape(-1)
+        )
+        term1 = np.bincount(qbin_labels, weights=np.abs(F1).reshape(-1) ** 2)
+        term2 = np.bincount(qbin_labels, weights=np.abs(F2).reshape(-1) ** 2)
+        np.multiply(term1, term2, out=term1)
+        denominator = np.sqrt(term1)
+        FSC = np.divide(numerator, denominator)
+
+        qidx = np.where(qbins < qx.max())
+
+        return np.vstack((qbins[qidx], FSC[qidx])).T

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:
@@ -487,16 +494,22 @@ class Orientations:
 
     @classmethod
     def _from_star(
-        cls, filename: str, delimiter: str = "\t"
+        cls, filename: str, delimiter: str = None
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         parser = StarParser(filename, delimiter=delimiter)
 
-        ret = parser.get("data_particles", None)
-        if ret is None:
-            ret = parser.get("data_", None)
+        keyword_order = ("data_particles", "particles", "data")
+        for keyword in keyword_order:
+            ret = parser.get(keyword, None)
+            if ret is None:
+                ret = parser.get(f"{keyword}_", None)
+            if ret is not None:
+                break
 
         if ret is None:
-            raise ValueError(f"No data_particles section found in {filename}.")
+            raise ValueError(
+                f"Could not find either {keyword_order} section found in {filename}."
+            )
 
         translation = np.vstack(
             (ret["_rlnCoordinateX"], ret["_rlnCoordinateY"], ret["_rlnCoordinateZ"])