pytme 0.3.1.post1__tar.gz → 0.3.1.post2__tar.gz

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

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

{pytme-0.3.1.post1 → pytme-0.3.1.post2}/pyproject.toml

@@ -7,7 +7,7 @@ name="pytme"
 authors = [
   { name = "Valentin Maurer", email = "valentin.maurer@embl-hamburg.de" },
 ]
-version="0.3.1.post1"
+version="0.3.1.post2"
 description="Python Template Matching Engine"
 readme="README.md"
 requires-python = ">=3.11"

{pytme-0.3.1.post1 → pytme-0.3.1.post2}/scripts/match_template.py

@@ -359,8 +359,8 @@ def parse_args():
         "--invert-target-contrast",
         action="store_true",
         default=False,
-        help="Invert the target's contrast for cases where templates to-be-matched have "
-        "negative values, e.g. tomograms.",
+        help="Invert the target contrast. Useful for matching on tomograms if the "
+        "template has not been inverted.",
     )
     io_group.add_argument(
         "--scramble-phases",

{pytme-0.3.1.post1 → pytme-0.3.1.post2}/scripts/preprocessor_gui.py

@@ -474,6 +474,7 @@ def membrane_mask(
     **kwargs,
 ) -> NDArray:
     return create_mask(
+        center=(center_x, center_y, center_z),
         mask_type="membrane",
         shape=template.shape,
         radius=radius,

{pytme-0.3.1.post1 → pytme-0.3.1.post2}/scripts/refine_matches.py

@@ -10,11 +10,9 @@ import subprocess
 from sys import exit
 from os import unlink
 from time import time
-from os.path import join
 from typing import Tuple, List, Dict
 
 import numpy as np
-from scipy import optimize
 from sklearn.metrics import roc_auc_score
 
 from tme import Orientations, Density
@@ -66,7 +64,6 @@ def parse_args():
     matching_group.add_argument(
         "-i",
         "--template",
-        dest="template",
         type=str,
         required=True,
         help="Path to a template in PDB/MMCIF or other supported formats (see target).",
@@ -102,7 +99,7 @@ def parse_args():
     )
     matching_group.add_argument(
         "-s",
-        dest="score",
+        "--score",
         type=str,
         default="batchFLCSphericalMask",
         choices=list(MATCHING_EXHAUSTIVE_REGISTER.keys()),
@@ -197,6 +194,7 @@ def create_matching_argdict(args) -> Dict:
         "-n": args.cores,
         "--ctf-file": args.ctf_file,
         "--invert-target-contrast": args.invert_target_contrast,
+        "--backend" : args.backend,
     }
     return arg_dict
 
@@ -252,7 +250,7 @@ class DeepMatcher:
         if args.lowpass_range:
             self.filter_parameters["--lowpass"] = 0
         if args.highpass_range:
-            self.filter_parameters["--highpass"] = 200
+            self.filter_parameters["--highpass"] = 0
 
         self.postprocess_args = create_postprocessing_argdict(args)
         self.log_file = f"{args.output_prefix}_optimization_log.txt"
@@ -309,14 +307,14 @@ class DeepMatcher:
 
         match_template = argdict_to_command(
             self.match_template_args,
-            executable="match_template.py",
+            executable="match_template",
         )
         run_command(match_template)
 
         # Assume we get a new peak for each input in the same order
         postprocess = argdict_to_command(
             self.postprocess_args,
-            executable="postprocess.py",
+            executable="postprocess",
         )
         run_command(postprocess)
 

{pytme-0.3.1.post1 → pytme-0.3.1.post2}/tests/test_orientations.py

@@ -95,18 +95,6 @@ class TestDensity:
             self.orientations.rotations, orientations_new.rotations, atol=1e-3
         )
 
-    @pytest.mark.parametrize("input_format", ("text", "star", "tbl"))
-    @pytest.mark.parametrize("output_format", ("text", "star", "tbl"))
-    def test_file_format_io(self, input_format: str, output_format: str):
-        _, output_file = mkstemp(suffix=f".{input_format}")
-        _, output_file2 = mkstemp(suffix=f".{output_format}")
-
-        self.orientations.to_file(output_file)
-        orientations_new = Orientations.from_file(output_file)
-        orientations_new.to_file(output_file2)
-
-        assert True
-
     @pytest.mark.parametrize("drop_oob", (True, False))
     @pytest.mark.parametrize("shape", (10, 40, 80))
     @pytest.mark.parametrize("odd", (True, False))

{pytme-0.3.1.post1 → pytme-0.3.1.post2}/tme/backends/_jax_utils.py

@@ -17,7 +17,7 @@ from ..backends import backend as be
 from ..matching_utils import normalize_template as _normalize_template
 
 
-__all__ = ["scan"]
+__all__ = ["scan", "setup_scan"]
 
 
 def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
@@ -116,12 +116,49 @@ def _mask_scores(arr, mask):
     return arr.at[:].multiply(mask)
 
 
-@partial(
-    pmap,
-    in_axes=(0,) + (None,) * 7,
-    static_broadcasted_argnums=[7, 8, 9, 10],
-    axis_name="batch",
-)
+def _select_config(analyzer_kwargs, device_idx):
+    return analyzer_kwargs[device_idx]
+
+
+def setup_scan(analyzer_kwargs, callback_class, fast_shape, rotate_mask):
+    """Create separate scan function with initialized analyzer for each device"""
+    device_scans = [
+        partial(
+            scan,
+            fast_shape=fast_shape,
+            rotate_mask=rotate_mask,
+            analyzer=callback_class(**device_config),
+        ) for device_config in analyzer_kwargs
+    ]
+
+    @partial(
+        pmap,
+        in_axes=(0,) + (None,) * 6,
+        axis_name="batch",
+    )
+    def scan_combined(
+        target,
+        template,
+        template_mask,
+        rotations,
+        template_filter,
+        target_filter,
+        score_mask,
+    ):
+        return lax.switch(
+            lax.axis_index("batch"),
+            device_scans,
+            target,
+            template,
+            template_mask,
+            rotations,
+            template_filter,
+            target_filter,
+            score_mask,
+        )
+    return scan_combined
+
+
 def scan(
     target: BackendArray,
     template: BackendArray,
@@ -132,17 +169,10 @@ def scan(
     score_mask: BackendArray,
     fast_shape: Tuple[int],
     rotate_mask: bool,
-    analyzer_class: object,
-    analyzer_kwargs: Tuple[Tuple],
+    analyzer: object,
 ) -> Tuple[BackendArray, BackendArray]:
     eps = jnp.finfo(template.dtype).resolution
 
-    kwargs = lax.switch(
-        lax.axis_index("batch"),
-        [lambda: analyzer_kwargs[i] for i in range(len(analyzer_kwargs))],
-    )
-    analyzer = analyzer_class(**be._tuple_to_dict(kwargs))
-
     if hasattr(target_filter, "shape"):
         target = _apply_fourier_filter(target, target_filter)
 

{pytme-0.3.1.post1 → pytme-0.3.1.post2}/tme/backends/jax_backend.py

@@ -218,7 +218,7 @@ class JaxBackend(NumpyFFTWBackend):
         Emulates output of :py:meth:`tme.matching_exhaustive.scan` using
         :py:class:`tme.analyzer.MaxScoreOverRotations`.
         """
-        from ._jax_utils import scan as scan_inner
+        from ._jax_utils import setup_scan
         from ..analyzer import MaxScoreOverRotations
 
         pad_target = True if len(splits) > 1 else False
@@ -279,8 +279,7 @@ class JaxBackend(NumpyFFTWBackend):
                 cur_args = analyzer_args.copy()
                 cur_args["offset"] = translation_offset
                 cur_args.update(callback_class_args)
-
-                analyzer_kwargs.append(self._dict_to_tuple(cur_args))
+                analyzer_kwargs.append(cur_args)
 
                 if pad_target:
                     score_mask = base._score_mask(fast_shape, shift)
@@ -310,7 +309,13 @@ class JaxBackend(NumpyFFTWBackend):
             create_filter, create_template_filter, create_target_filter = (False,) * 3
             base, targets = None, self._array_backend.stack(targets)
 
-            analyzer_kwargs = tuple(analyzer_kwargs)
+            scan_inner = setup_scan(
+                analyzer_kwargs=analyzer_kwargs,
+                callback_class=callback_class,
+                fast_shape=fast_shape,
+                rotate_mask=rotate_mask
+            )
+
             states = scan_inner(
                 self.astype(targets, self._float_dtype),
                 self.astype(matching_data.template, self._float_dtype),
@@ -319,17 +324,12 @@ class JaxBackend(NumpyFFTWBackend):
                 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])
+                kwargs = analyzer_kwargs[index]
                 analyzer = callback_class(**kwargs)
-
                 state = [self._unbatch(x, ndim, index) for x in states]
 
                 if isinstance(analyzer, MaxScoreOverRotations):

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

{pytme-0.3.1.post1 → pytme-0.3.1.post2}/tme/orientations.py

@@ -494,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"])

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

@@ -375,10 +375,10 @@ 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 = np.fmax(scores_out, 0, out=scores_out)
     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)
     data[0] = scores_out
 
     fg, bg = simple_stats(data[0]), simple_stats(scores_norm)