pytme 0.2.0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.1__cp311-cp311-macosx_14_0_arm64.whl

scripts/match_template_filters.py

@@ -178,6 +178,7 @@ def parse_rotation_logic(args, ndim):
 
 
 # TODO: Think about whether wedge mask should also be added to target
+# For now leave it at the cost of incorrect upper bound on the scores
 def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Compose]:
     from tme.preprocessing import LinearWhiteningFilter, BandPassFilter
     from tme.preprocessing.tilt_series import (
@@ -216,7 +217,7 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
             if args.tilt_weighting not in ("angle", None):
                 raise ValueError(
                     "Tilt weighting schemes other than 'angle' or 'None' require "
-                    "a specification of electron doses."
+                    "a specification of electron doses via --tilt_angles."
                 )
 
             wedge = Wedge(
@@ -239,31 +240,58 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
         wedge.sampling_rate = template.sampling_rate
         template_filter.append(wedge)
         if not isinstance(wedge, WedgeReconstructed):
-            template_filter.append(ReconstructFromTilt(
-                reconstruction_filter = args.reconstruction_filter
-            ))
+            template_filter.append(
+                ReconstructFromTilt(
+                    reconstruction_filter=args.reconstruction_filter,
+                    interpolation_order=args.reconstruction_interpolation_order,
+                )
+            )
 
-    if args.ctf_file is not None:
+    if args.ctf_file is not None or args.defocus is not None:
         from tme.preprocessing.tilt_series import CTF
 
-        ctf = CTF.from_file(args.ctf_file)
-        n_tilts_ctfs, n_tils_angles = len(ctf.defocus_x), len(wedge.angles)
-        if n_tilts_ctfs != n_tils_angles:
-            raise ValueError(
-                f"CTF file contains {n_tilts_ctfs} micrographs, but match_template "
-                f"recieved {n_tils_angles} tilt angles. Expected one angle "
-                "per micrograph."
+        needs_reconstruction = True
+        if args.ctf_file is not None:
+            ctf = CTF.from_file(args.ctf_file)
+            n_tilts_ctfs, n_tils_angles = len(ctf.defocus_x), len(wedge.angles)
+            if n_tilts_ctfs != n_tils_angles:
+                raise ValueError(
+                    f"CTF file contains {n_tilts_ctfs} micrographs, but match_template "
+                    f"recieved {n_tils_angles} tilt angles. Expected one angle "
+                    "per micrograph."
+                )
+            ctf.angles = wedge.angles
+            ctf.opening_axis, ctf.tilt_axis = args.wedge_axes
+        else:
+            needs_reconstruction = False
+            ctf = CTF(
+                defocus_x=[args.defocus],
+                phase_shift=[args.phase_shift],
+                defocus_y=None,
+                angles=[0],
+                shape=None,
+                return_real_fourier=True,
             )
-        ctf.angles = wedge.angles
-        ctf.opening_axis, ctf.tilt_axis = args.wedge_axes
-
-        if isinstance(template_filter[-1], ReconstructFromTilt):
+        ctf.sampling_rate = template.sampling_rate
+        ctf.flip_phase = not args.no_flip_phase
+        ctf.amplitude_contrast = args.amplitude_contrast
+        ctf.spherical_aberration = args.spherical_aberration
+        ctf.acceleration_voltage = args.acceleration_voltage * 1e3
+        ctf.correct_defocus_gradient = args.correct_defocus_gradient
+
+        if not needs_reconstruction:
+            template_filter.append(ctf)
+        elif isinstance(template_filter[-1], ReconstructFromTilt):
             template_filter.insert(-1, ctf)
         else:
             template_filter.insert(0, ctf)
-            template_filter.insert(1, ReconstructFromTilt(
-                reconstruction_filter = args.reconstruction_filter
-            ))
+            template_filter.insert(
+                1,
+                ReconstructFromTilt(
+                    reconstruction_filter=args.reconstruction_filter,
+                    interpolation_order=args.reconstruction_interpolation_order,
+                ),
+            )
 
     if args.lowpass or args.highpass is not None:
         lowpass, highpass = args.lowpass, args.highpass
@@ -292,6 +320,14 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
         template_filter.append(whitening_filter)
         target_filter.append(whitening_filter)
 
+    needs_reconstruction = any(
+        [isinstance(t, ReconstructFromTilt) for t in template_filter]
+    )
+    if needs_reconstruction and args.reconstruction_filter is None:
+        warnings.warn(
+            "Consider using a --reconstruction_filter such as 'ramp' to avoid artifacts."
+        )
+
     template_filter = Compose(template_filter) if len(template_filter) else None
     target_filter = Compose(target_filter) if len(target_filter) else None
 
@@ -509,7 +545,7 @@ def parse_args():
         dest="no_pass_smooth",
         action="store_false",
         default=True,
-        help="Whether a hard edge filter should be used for --lowpass and --highpass."
+        help="Whether a hard edge filter should be used for --lowpass and --highpass.",
     )
     filter_group.add_argument(
         "--pass_format",
@@ -518,7 +554,7 @@ def parse_args():
         required=False,
         choices=["sampling_rate", "voxel", "frequency"],
         help="How values passed to --lowpass and --highpass should be interpreted. "
-        "By default, they are assumed to be in units of sampling rate, e.g. Ångstrom."
+        "By default, they are assumed to be in units of sampling rate, e.g. Ångstrom.",
     )
     filter_group.add_argument(
         "--whiten_spectrum",
@@ -560,23 +596,90 @@ def parse_args():
         "grigorieff (exposure filter as defined in Grant and Grigorieff 2015)."
         "relion and grigorieff require electron doses in --tilt_angles weights column.",
     )
-    # filter_group.add_argument(
-    #     "--ctf_file",
-    #     dest="ctf_file",
-    #     type=str,
-    #     required=False,
-    #     default=None,
-    #     help="Path to a file with CTF parameters from CTFFIND4.",
-    # )
     filter_group.add_argument(
         "--reconstruction_filter",
         dest="reconstruction_filter",
         type=str,
         required=False,
-        choices = ["ram-lak", "ramp", "shepp-logan", "cosine", "hamming"],
+        choices=["ram-lak", "ramp", "ramp-cont", "shepp-logan", "cosine", "hamming"],
         default=None,
         help="Filter applied when reconstructing (N+1)-D from N-D filters.",
     )
+    filter_group.add_argument(
+        "--reconstruction_interpolation_order",
+        dest="reconstruction_interpolation_order",
+        type=int,
+        default=1,
+        required=False,
+        help="Analogous to --interpolation_order but for reconstruction.",
+    )
+
+    ctf_group = parser.add_argument_group("Contrast Transfer Function")
+    ctf_group.add_argument(
+        "--ctf_file",
+        dest="ctf_file",
+        type=str,
+        required=False,
+        default=None,
+        help="Path to a file with CTF parameters from CTFFIND4. Each line will be "
+        "interpreted as tilt obtained at the angle specified in --tilt_angles. ",
+    )
+    ctf_group.add_argument(
+        "--defocus",
+        dest="defocus",
+        type=float,
+        required=False,
+        default=None,
+        help="Defocus in units of sampling rate (typically Ångstrom). "
+        "Superseded by --ctf_file.",
+    )
+    ctf_group.add_argument(
+        "--phase_shift",
+        dest="phase_shift",
+        type=float,
+        required=False,
+        default=0,
+        help="Phase shift in degrees. Superseded by --ctf_file.",
+    )
+    ctf_group.add_argument(
+        "--acceleration_voltage",
+        dest="acceleration_voltage",
+        type=float,
+        required=False,
+        default=300,
+        help="Acceleration voltage in kV, defaults to 300.",
+    )
+    ctf_group.add_argument(
+        "--spherical_aberration",
+        dest="spherical_aberration",
+        type=float,
+        required=False,
+        default=2.7e7,
+        help="Spherical aberration in units of sampling rate (typically Ångstrom).",
+    )
+    ctf_group.add_argument(
+        "--amplitude_contrast",
+        dest="amplitude_contrast",
+        type=float,
+        required=False,
+        default=0.07,
+        help="Amplitude contrast, defaults to 0.07.",
+    )
+    ctf_group.add_argument(
+        "--no_flip_phase",
+        dest="no_flip_phase",
+        action="store_false",
+        required=False,
+        help="Whether the phase of the computed CTF should not be flipped.",
+    )
+    ctf_group.add_argument(
+        "--correct_defocus_gradient",
+        dest="correct_defocus_gradient",
+        action="store_true",
+        required=False,
+        help="[Experimental] Whether to compute a more accurate 3D CTF incorporating "
+        "defocus gradients.",
+    )
 
     performance_group = parser.add_argument_group("Performance")
     performance_group.add_argument(
@@ -654,12 +757,11 @@ def parse_args():
     )
 
     args = parser.parse_args()
+    args.version = __version__
 
     if args.interpolation_order < 0:
         args.interpolation_order = None
 
-    args.ctf_file = None
-
     if args.temp_directory is None:
         default = abspath(".")
         if os.environ.get("TMPDIR", None) is not None:
@@ -719,17 +821,20 @@ def main():
     try:
         template = Density.from_file(args.template)
     except Exception:
+        drop = target.metadata.get("batch_dimension", ())
+        keep = [i not in drop for i in range(target.data.ndim)]
         template = Density.from_structure(
             filename_or_structure=args.template,
-            sampling_rate=target.sampling_rate,
+            sampling_rate=target.sampling_rate[keep],
         )
 
-    if not np.allclose(target.sampling_rate, template.sampling_rate):
-        print(
-            f"Resampling template to {target.sampling_rate}. "
-            "Consider providing a template with the same sampling rate as the target."
-        )
-        template = template.resample(target.sampling_rate, order=3)
+    if target.sampling_rate.size == template.sampling_rate.size:
+        if not np.allclose(target.sampling_rate, template.sampling_rate):
+            print(
+                f"Resampling template to {target.sampling_rate}. "
+                "Consider providing a template with the same sampling rate as the target."
+            )
+            template = template.resample(target.sampling_rate, order=3)
 
     template_mask = load_and_validate_mask(
         mask_target=template, mask_path=args.template_mask
@@ -862,31 +967,52 @@ def main():
     if args.memory is None:
         args.memory = int(args.memory_scaling * available_memory)
 
-    target_padding = np.zeros_like(template.shape)
-    if args.pad_target_edges:
-        target_padding = template.shape
+    callback_class = MaxScoreOverRotations
+    if args.peak_calling:
+        callback_class = PeakCallerMaximumFilter
+
+    matching_data = MatchingData(
+        target=target,
+        template=template.data,
+        target_mask=target_mask,
+        template_mask=template_mask,
+        invert_target=args.invert_target_contrast,
+        rotations=parse_rotation_logic(args=args, ndim=template.data.ndim),
+    )
 
-    template_box = template.shape
+    template_filter, target_filter = setup_filter(args, template, target)
+    matching_data.template_filter = template_filter
+    matching_data.target_filter = target_filter
+
+    target_dims = target.metadata.get("batch_dimension", None)
+    matching_data._set_batch_dimension(target_dims=target_dims, template_dims=None)
+    args.score = "FLC2" if target_dims is not None else args.score
+    args.target_batch, args.template_batch = target_dims, None
+
+    template_box = matching_data._output_template_shape
     if not args.pad_fourier:
         template_box = np.ones(len(template_box), dtype=int)
 
-    callback_class = MaxScoreOverRotations
-    if args.peak_calling:
-        callback_class = PeakCallerMaximumFilter
+    target_padding = np.zeros(
+        (backend.size(matching_data._output_template_shape)), dtype=int
+    )
+    if args.pad_target_edges:
+        target_padding = matching_data._output_template_shape
 
     splits, schedule = compute_parallelization_schedule(
         shape1=target.shape,
-        shape2=template_box,
-        shape1_padding=target_padding,
+        shape2=tuple(int(x) for x in template_box),
+        shape1_padding=tuple(int(x) for x in target_padding),
         max_cores=args.cores,
         max_ram=args.memory,
         split_only_outer=args.use_gpu,
         matching_method=args.score,
         analyzer_method=callback_class.__name__,
         backend=backend._backend_name,
-        float_nbytes=backend.datatype_bytes(backend._default_dtype),
+        float_nbytes=backend.datatype_bytes(backend._float_dtype),
         complex_nbytes=backend.datatype_bytes(backend._complex_dtype),
-        integer_nbytes=backend.datatype_bytes(backend._default_dtype_int),
+        integer_nbytes=backend.datatype_bytes(backend._int_dtype),
+        split_axes=target_dims,
     )
 
     if splits is None:
@@ -897,20 +1023,6 @@ def main():
         exit(-1)
 
     matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[args.score]
-    matching_data = MatchingData(target=target, template=template.data)
-    matching_data.rotations = parse_rotation_logic(args=args, ndim=target.data.ndim)
-
-    template_filter, target_filter = setup_filter(args, template, target)
-    matching_data.template_filter = template_filter
-    matching_data.target_filter = target_filter
-
-    matching_data.template_filter = template_filter
-    matching_data._invert_target = args.invert_target_contrast
-    if target_mask is not None:
-        matching_data.target_mask = target_mask
-    if template_mask is not None:
-        matching_data.template_mask = template_mask.data
-
     n_splits = np.prod(list(splits.values()))
     target_split = ", ".join(
         [":".join([str(x) for x in axis]) for axis in splits.items()]
@@ -944,13 +1056,23 @@ def main():
         "Lowpass": args.lowpass,
         "Highpass": args.highpass,
         "Smooth Pass": args.no_pass_smooth,
-        "Pass Format" : args.pass_format,
+        "Pass Format": args.pass_format,
         "Spectral Whitening": args.whiten_spectrum,
         "Wedge Axes": args.wedge_axes,
         "Tilt Angles": args.tilt_angles,
         "Tilt Weighting": args.tilt_weighting,
-        "CTF": args.ctf_file,
+        "Reconstruction Filter": args.reconstruction_filter,
     }
+    if args.ctf_file is not None or args.defocus is not None:
+        filter_args["CTF File"] = args.ctf_file
+        filter_args["Defocus"] = args.defocus
+        filter_args["Phase Shift"] = args.phase_shift
+        filter_args["No Flip Phase"] = args.no_flip_phase
+        filter_args["Acceleration Voltage"] = args.acceleration_voltage
+        filter_args["Spherical Aberration"] = args.spherical_aberration
+        filter_args["Amplitude Contrast"] = args.amplitude_contrast
+        filter_args["Correct Defocus"] = args.correct_defocus_gradient
+
     filter_args = {k: v for k, v in filter_args.items() if v is not None}
     if len(filter_args):
         print_block(
@@ -999,16 +1121,16 @@ def main():
             candidates[0] *= target_mask.data
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", category=UserWarning)
+            nbytes = backend.datatype_bytes(backend._float_dtype)
+            dtype = np.float32 if nbytes == 4 else np.float16
+            rot_dim = matching_data.rotations.shape[1]
             candidates[3] = {
                 x: euler_from_rotationmatrix(
-                    np.frombuffer(i, dtype=matching_data.rotations.dtype).reshape(
-                        candidates[0].ndim, candidates[0].ndim
-                    )
+                    np.frombuffer(i, dtype=dtype).reshape(rot_dim, rot_dim)
                 )
                 for i, x in candidates[3].items()
             }
-
-    candidates.append((target.origin, template.origin, target.sampling_rate, args))
+    candidates.append((target.origin, template.origin, template.sampling_rate, args))
     write_pickle(data=candidates, filename=args.output)
 
     runtime = time() - start