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

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

@@ -87,11 +87,6 @@ def parse_args():
         help="Output prefix. Defaults to basename of first input. Extension is "
         "added with respect to chosen output format.",
     )
-    output_group.add_argument(
-        "--angles-clockwise",
-        action="store_true",
-        help="Report Euler angles in clockwise format expected by RELION.",
-    )
     output_group.add_argument(
         "--output-format",
         choices=[
@@ -173,13 +168,6 @@ def parse_args():
         default=None,
         help="Box size of extracted subtomograms, defaults to the centered template.",
     )
-    additional_group.add_argument(
-        "--mask-subtomograms",
-        action="store_true",
-        default=False,
-        help="Whether to mask subtomograms using the template mask. The mask will be "
-        "rotated according to determined angles.",
-    )
     additional_group.add_argument(
         "--invert-target-contrast",
         action="store_true",
@@ -219,20 +207,15 @@ def load_template(
 ):
     try:
         template = Density.from_file(filepath)
-        center = np.divide(np.subtract(template.shape, 1), 2)
         template_is_density = True
     except Exception:
-        template = Structure.from_file(filepath)
-        center = template.center_of_mass()
-        template = Density.from_structure(template, sampling_rate=sampling_rate)
+        template = Density.from_structure(filepath, sampling_rate=sampling_rate)
         template_is_density = False
 
-    translation = np.zeros_like(center)
-    if centering and template_is_density:
-        template, translation = template.centered(0)
-        center = np.divide(np.subtract(template.shape, 1), 2)
-
-    return template, center, translation, template_is_density
+    if centering:
+        template = template.centered(0)
+    center = np.divide(np.subtract(template.shape, 1), 2)
+    return template, center, template_is_density
 
 
 def load_matching_output(path: str) -> List:
@@ -375,10 +358,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)
@@ -448,34 +431,24 @@ def main():
     if hasattr(cli_args, "no_centering"):
         cli_args.centering = not cli_args.no_centering
 
-    _, template_extension = splitext(cli_args.template)
-    ret = load_template(
+    template, *_ = load_template(
         filepath=cli_args.template,
         sampling_rate=sampling_rate,
         centering=cli_args.centering,
     )
-    template, center_of_mass, translation, template_is_density = ret
 
     template_mask = template.empty
     template_mask.data[:] = 1
     if cli_args.template_mask is not None:
         template_mask = Density.from_file(cli_args.template_mask)
-        template_mask.pad(template.shape, center=False)
-        origin_translation = np.divide(
-            np.subtract(template.origin, template_mask.origin), template.sampling_rate
-        )
-        translation = np.add(translation, origin_translation)
-
-        template_mask = template_mask.rigid_transform(
-            rotation_matrix=np.eye(template_mask.data.ndim),
-            translation=-translation,
-            order=1,
-        )
+    if cli_args.centering:
+        template_mask.pad(template.shape, center=True)
 
     if args.mask_edges and args.min_boundary_distance == 0:
         max_shape = np.max(template.shape)
         args.min_boundary_distance = np.ceil(np.divide(max_shape, 2))
 
+    # Do the actual peak calling
     orientations = args.orientations
     if orientations is None:
         translations, rotations, scores, details = [], [], [], []
@@ -518,18 +491,20 @@ def main():
             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"):
-            args.batch_dims = cli_args.batch_dims
+        projection_dims = None
+        batch_dims = getattr(cli_args, "batch_dims", None)
+        if getattr(cli_args, "match_projection", False):
+            projection_dims = batch_dims
 
         peak_caller_kwargs = {
             "shape": scores.shape,
             "num_peaks": args.num_peaks,
             "min_distance": args.min_distance,
             "min_boundary_distance": args.min_boundary_distance,
-            "batch_dims": args.batch_dims,
             "min_score": args.min_score,
             "max_score": args.max_score,
+            "batch_dims": batch_dims,
+            "projection_dims": projection_dims,
         }
 
         peak_caller = PEAK_CALLERS[args.peak_caller](**peak_caller_kwargs)
@@ -551,13 +526,9 @@ def main():
             exit(-1)
 
         for translation, _, score, detail in zip(*candidates):
-            rotation_index = rotation_array[tuple(translation)]
-            rotation = rotation_mapping.get(
-                rotation_index, np.zeros(template.data.ndim, int)
-            )
-            if rotation.ndim == 2:
-                rotation = euler_from_rotationmatrix(rotation)
-            rotations.append(rotation)
+            index = rotation_array[tuple(translation)]
+            rotation = rotation_mapping.get(index, np.eye(template.data.ndim))
+            rotations.append(euler_from_rotationmatrix(rotation, seq="ZYZ"))
 
         if len(rotations):
             rotations = np.vstack(rotations).astype(float)
@@ -583,11 +554,10 @@ def main():
 
     if args.peak_oversampling > 1:
         if data[0].ndim != data[2].ndim:
-            print(
+            exit(
                 "Input pickle does not contain template matching scores."
                 " Cannot oversample peaks."
             )
-            exit(-1)
         peak_caller = peak_caller = PEAK_CALLERS[args.peak_caller](shape=scores.shape)
         orientations.translations = peak_caller.oversample_peaks(
             scores=data[0],
@@ -597,8 +567,6 @@ def main():
 
     if args.local_optimization:
         target = Density.from_file(cli_args.target, use_memmap=True)
-        orientations.translations = orientations.translations.astype(np.float32)
-        orientations.rotations = orientations.rotations.astype(np.float32)
         for index, (translation, angles, *_) in enumerate(orientations):
             score_object = create_score_object(
                 score="FLC",
@@ -619,12 +587,11 @@ def main():
                 x0=[*init_translation, *angles],
             )
             orientations.translations[index] = np.add(translation, center)
-            orientations.rotations[index] = angles
+            orientations.rotations[index] = euler_from_rotationmatrix(
+                rotation_matrix, seq="ZYZ"
+            )
             orientations.scores[index] = score * -1
 
-    if args.angles_clockwise:
-        orientations.rotations *= -1
-
     if args.output_format in ("orientations", "relion4", "relion5"):
         file_format, extension = "text", "tsv"
 
@@ -691,12 +658,6 @@ def main():
                 sampling_rate=sampling_rate,
                 origin=np.multiply(cand_start, sampling_rate),
             )
-            if args.mask_subtomograms:
-                rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
-                mask_transfomed = template_mask.rigid_transform(
-                    rotation_matrix=rotation_matrix, order=1
-                )
-                out_density.data = out_density.data * mask_transfomed.data
             out_density.to_file(
                 join(working_directory, f"{args.output_prefix}_{index}.mrc")
             )
@@ -713,10 +674,11 @@ def main():
         out = np.zeros_like(template.data)
         for index in range(len(cand_slices)):
             subset = Density(target.data[obs_slices[index]])
-            rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
 
+            # We invert to pull the local into the global reference system
+            matrix = euler_to_rotationmatrix(orientations.rotations[index]).T
             subset = subset.rigid_transform(
-                rotation_matrix=np.linalg.inv(rotation_matrix),
+                rotation_matrix=matrix,
                 order=1,
                 use_geometric_center=True,
             )
@@ -728,37 +690,37 @@ def main():
         ret.to_file(f"{args.output_prefix}.mrc")
         exit(0)
 
-    template, center, *_ = load_template(
+    template, center, template_is_density, *_ = load_template(
         filepath=cli_args.template,
         sampling_rate=sampling_rate,
         centering=cli_args.centering,
     )
 
+    _, ext = splitext(cli_args.template)
     for index, (translation, angles, *_) in enumerate(orientations):
-        rotation_matrix = euler_to_rotationmatrix(angles)
+        rotation = euler_to_rotationmatrix(angles, seq="ZYZ")
         if template_is_density:
             transformed_template = template.rigid_transform(
-                rotation_matrix=rotation_matrix, use_geometric_center=True
+                rotation_matrix=rotation, use_geometric_center=True
             )
             # Just adapting the coordinate system not the in-box position
             shift = np.multiply(np.subtract(translation, center), sampling_rate)
             transformed_template.origin = np.add(target_origin, shift)
-
         else:
             template = Structure.from_file(cli_args.template)
-            new_center_of_mass = np.add(
-                np.multiply(translation, sampling_rate), target_origin
-            )
-            translation = np.subtract(new_center_of_mass, center)
+            shift = np.add(np.multiply(translation, sampling_rate), target_origin)
+            translation = np.subtract(shift, template.center_of_mass())
+
+            # Since we move the template's center of mass to the geometric center
+            # during matching and analysis we use the center of mass
+            # directly for rotating structures into the correct orientation
             transformed_template = template.rigid_transform(
                 translation=translation,
-                rotation_matrix=rotation_matrix,
+                rotation_matrix=rotation,
+                use_geometric_center=False,
             )
-        # template_extension should contain '.'
-        transformed_template.to_file(
-            f"{args.output_prefix}_{index}{template_extension}"
-        )
-        index += 1
+
+        transformed_template.to_file(f"{args.output_prefix}_{index}{ext}")
 
 
 if __name__ == "__main__":

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

@@ -126,7 +126,7 @@ def main():
         if args.align_axis is not None:
             rmat = data.align_to_axis(axis=args.align_axis, flip=args.flip_axis)
             data = data.rigid_transform(
-                rotation_matrix=rmat, translation=0, use_geometric_center=True
+                rotation_matrix=rmat.T, use_geometric_center=True
             )
         data = Density.from_structure(data, sampling_rate=sampling_rate)
 
@@ -138,11 +138,11 @@ def main():
         if args.align_axis is not None:
             rmat = data.align_to_axis(axis=args.align_axis, flip=args.flip_axis)
             data = data.rigid_transform(
-                rotation_matrix=rmat, translation=0, use_geometric_center=True
+                rotation_matrix=rmat.T, use_geometric_center=True
             )
 
     if not args.no_centering:
-        data, _ = data.centered(0)
+        data = data.centered(0)
 
     if args.box_size is None:
         scale = np.divide(data.sampling_rate, args.sampling_rate)
@@ -177,9 +177,8 @@ def main():
             lowpass=lowpass,
             highpass=None,
             use_gaussian=True,
-            return_real_fourier=True,
             sampling_rate=data.sampling_rate,
-        )(shape=data.shape)["data"]
+        )(shape=data.shape, return_real_fourier=True)["data"]
         bpf_mask = be.to_backend_array(bpf_mask)
 
         data_ft = be.rfftn(be.to_backend_array(data.data), s=data.shape)

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

@@ -26,14 +26,16 @@ from tme.backends import backend as be
 from tme.rotations import align_vectors
 from tme.matching_utils import create_mask, load_pickle
 from tme import Preprocessor, Density, Orientations
-from tme.filters import BandPassReconstructed, CTFReconstructed
+from tme.filters import BandPassReconstructed, CTFReconstructed, WedgeReconstructed
 
 preprocessor = Preprocessor()
 SLIDER_MIN, SLIDER_MAX = 0, 25
 
 
-def gaussian_filter(template: NDArray, sigma: float, **kwargs: dict) -> NDArray:
-    return preprocessor.gaussian_filter(template=template, sigma=sigma, **kwargs)
+def _apply_fourier_filter(arr, arr_filter):
+    arr_ft = np.fft.rfftn(arr, s=arr.shape)
+    arr_ft = np.multiply(arr_ft, arr_filter, out=arr_ft)
+    return np.real(np.fft.irfftn(arr_ft, s=arr.shape))
 
 
 def bandpass_filter(
@@ -48,13 +50,8 @@ def bandpass_filter(
         highpass=highpass_angstrom,
         sampling_rate=np.max(sampling_rate),
         use_gaussian=not hard_edges,
-        return_real_fourier=True,
-    )
-    template_ft = np.fft.rfftn(template, s=template.shape)
-
-    mask = bpf(shape=template.shape)["data"]
-    np.multiply(template_ft, mask, out=template_ft)
-    return np.fft.irfftn(template_ft, s=template.shape).real
+    )(shape=template.shape, return_real_fourier=True)["data"]
+    return _apply_fourier_filter(template, bpf)
 
 
 def ctf_filter(
@@ -70,9 +67,7 @@ def ctf_filter(
 ) -> NDArray:
     fast_shape = [next_fast_len(x) for x in np.multiply(template.shape, 2)]
     template_pad = be.topleft_pad(template, fast_shape)
-    template_ft = np.fft.rfftn(template_pad, s=template_pad.shape)
     ctf = CTFReconstructed(
-        shape=fast_shape,
         defocus_x=[defocus_angstrom],
         acceleration_voltage=acceleration_voltage * 1e3,
         spherical_aberration=spherical_aberration * 1e7,
@@ -80,22 +75,18 @@ def ctf_filter(
         phase_shift=phase_shift,
         defocus_angle=defocus_angle,
         sampling_rate=np.max(sampling_rate),
-        return_real_fourier=True,
         flip_phase=flip_phase,
-    )
-    np.multiply(template_ft, ctf()["data"], out=template_ft)
-    template_pad = np.fft.irfftn(template_ft, s=template_pad.shape).real
-    template = be.topleft_pad(template_pad, template.shape)
-    return template
+    )(shape=template.shape, return_real_fourier=True)["data"]
+    template = _apply_fourier_filter(template, ctf)
+    return be.topleft_pad(template_pad, template.shape)
 
 
-def difference_of_gaussian_filter(
-    template: NDArray, sigmas: Tuple[float, float], **kwargs: dict
-) -> NDArray:
-    low_sigma, high_sigma = sigmas
-    return preprocessor.difference_of_gaussian_filter(
-        template=template, low_sigma=low_sigma, high_sigma=high_sigma, **kwargs
-    )
+def gaussian_filter(template: NDArray, sigma: float, **kwargs: dict) -> NDArray:
+    return preprocessor.gaussian_filter(template=template, sigma=sigma, **kwargs)
+
+
+def difference_of_gaussian_filter(template, sigmas: Tuple[float, float], **kwargs):
+    return gaussian_filter(template, sigmas[0]) - gaussian_filter(template, sigmas[1])
 
 
 def edge_gaussian_filter(
@@ -132,11 +123,7 @@ def local_gaussian_filter(
     )
 
 
-def mean(
-    template: NDArray,
-    width: int,
-    **kwargs: dict,
-) -> NDArray:
+def mean(template: NDArray, width: int, **kwargs: dict) -> NDArray:
     return preprocessor.mean_filter(template=template, width=width)
 
 
@@ -147,45 +134,23 @@ def wedge(
     tilt_step: float = 0,
     opening_axis: int = 2,
     tilt_axis: int = 0,
-    omit_negative_frequencies: bool = False,
-    infinite_plane: bool = False,
-    weight_angle: bool = False,
     **kwargs,
 ) -> NDArray:
-    template_ft = np.fft.fftn(template)
-
-    if tilt_step <= 0:
-        wedge_mask = preprocessor.continuous_wedge_mask(
-            start_tilt=tilt_start,
-            stop_tilt=tilt_stop,
-            tilt_axis=tilt_axis,
-            opening_axis=opening_axis,
-            shape=template.shape,
-            omit_negative_frequencies=omit_negative_frequencies,
-            infinite_plane=infinite_plane,
-        )
-    else:
-        weights = None
-        tilt_angles = np.arange(-tilt_start, tilt_stop + tilt_step, tilt_step)
-        if weight_angle:
-            weights = np.cos(np.radians(tilt_angles))
-
-        wedge_mask = preprocessor.step_wedge_mask(
-            tilt_angles=tilt_angles,
-            tilt_axis=tilt_axis,
-            opening_axis=opening_axis,
-            shape=template.shape,
-            weights=weights,
-            omit_negative_frequencies=omit_negative_frequencies,
-        )
-
-    np.multiply(template_ft, wedge_mask, out=template_ft)
-    template = np.real(np.fft.ifftn(template_ft))
-    return template
+    mask = wedge_mask(
+        template=template,
+        fftshift=False,
+        tilt_start=tilt_start,
+        tilt_stop=tilt_stop,
+        tilt_step=tilt_step,
+        return_real_fourier=True,
+        weight_angle=False,
+    )
+    return _apply_fourier_filter(template, mask)
 
 
 def compute_power_spectrum(template: NDArray) -> NDArray:
-    return np.fft.fftshift(np.log(np.abs(np.fft.fftn(template))))
+    return np.fft.fftshift(np.log(1 + np.abs(np.fft.fftn(template))))
+    # return np.fft.fftshift(np.log(np.abs(np.fft.fftn(template))))
 
 
 def invert_contrast(template: NDArray) -> NDArray:
@@ -474,6 +439,7 @@ def membrane_mask(
     **kwargs,
 ) -> NDArray:
     return create_mask(
+        center=(center_x, center_y, center_z),
         mask_type="membrane",
         shape=template.shape,
         radius=radius,
@@ -514,39 +480,30 @@ def wedge_mask(
     tilt_step: float = 0,
     opening_axis: int = 2,
     tilt_axis: int = 0,
-    omit_negative_frequencies: bool = False,
-    infinite_plane: bool = False,
-    weight_angle: bool = False,
     **kwargs,
 ) -> NDArray:
-    if tilt_step <= 0:
-        wedge_mask = preprocessor.continuous_wedge_mask(
-            start_tilt=tilt_start,
-            stop_tilt=tilt_stop,
-            tilt_axis=tilt_axis,
-            opening_axis=opening_axis,
-            shape=template.shape,
-            omit_negative_frequencies=omit_negative_frequencies,
-            infinite_plane=infinite_plane,
-        )
-        wedge_mask = np.fft.fftshift(wedge_mask)
-        return wedge_mask
-
-    weights = None
-    tilt_angles = np.arange(-tilt_start, tilt_stop + tilt_step, tilt_step)
-    if weight_angle:
-        weights = np.cos(np.radians(tilt_angles))
-
-    wedge_mask = preprocessor.step_wedge_mask(
-        tilt_angles=tilt_angles,
+    angles = (tilt_start, tilt_stop)
+    continuous_wedge = tilt_step == 0
+    if not continuous_wedge:
+        angles = np.arange(-tilt_start, tilt_stop + tilt_step, tilt_step)
+
+    return_real_fourier = kwargs.get("return_real_fourier", False)
+    func = WedgeReconstructed(
+        angles=angles,
         tilt_axis=tilt_axis,
         opening_axis=opening_axis,
-        shape=template.shape,
-        weights=weights,
-        omit_negative_frequencies=omit_negative_frequencies,
+        frequency_cutoff=0.5,
+        create_continuous_wedge=continuous_wedge,
+        weight_wedge=kwargs.get("weight_angle", False),
     )
-
-    wedge_mask = np.fft.fftshift(wedge_mask)
+    wedge_mask = func(shape=template.shape, return_real_fourier=return_real_fourier)[
+        "data"
+    ]
+    if kwargs.get("fftshift", True):
+        axes = [i for i in range(wedge_mask.ndim)]
+        if return_real_fourier:
+            _ = axes.pop(-1)
+        wedge_mask = np.fft.fftshift(wedge_mask, axes=axes)
     return wedge_mask
 
 
@@ -569,8 +526,7 @@ def threshold_mask(
         mask[mask < np.exp(-np.square(sigma))] = 0
 
     if invert:
-        np.invert(mask, out=mask)
-
+        mask = 1 - mask
     return mask
 
 
@@ -757,15 +713,6 @@ class MaskWidget(widgets.Container):
             if self.method_dropdown.value == "Shape":
                 new_layer.metadata = {}
 
-        # origin_layer = metadata["origin_layer"]
-        # if origin_layer in self.viewer.layers:
-        #     origin_layer = self.viewer.layers[origin_layer]
-        #     if np.allclose(origin_layer.data.shape, processed_data.shape):
-        #         in_mask = np.sum(np.fmax(origin_layer.data * processed_data, 0))
-        #         in_mask /= np.sum(np.fmax(origin_layer.data, 0))
-        #         in_mask *= 100
-        #         self.density_field.value = f"Positive Density in Mask: {in_mask:.2f}%"
-
 
 class AlignmentWidget(widgets.Container):
     def __init__(self, viewer):