Phase plate ctf and argparse grouping (#157)

* add argument groups; add phase shift for CTF * fix the phase shift upon loading instead * switch to -phase_shift in CTF; add ctf param to tests * make check proper, thanks tests:) * add link to refer to other phase shift ctf implementation * add to tests the phase shift * Update src/pytom_tm/weights.py Co-authored-by: Sander Roet <[email protected]> * fix phase shfit sign * update docstring to reflect the inverse ctf curve * Update src/pytom_tm/entry_points.py Co-authored-by: Sander Roet <[email protected]> * add test job file from before 0.6.1 * make path relative * file extension was wrong * make botched job instead dynamically --------- Co-authored-by: Sander Roet <[email protected]>
SBC-Utrecht · Apr 5, 2024 · 1291893 · 1291893
1 parent b218916
commit 1291893
Show file tree

Hide file tree

Showing 6 changed files with 119 additions and 66 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pytom-match-pick"
-version = "0.6.0"
+version = "0.6.1"
 description = "PyTOM's GPU template matching module as an independent package"
 readme = "README.md"
 license = {file = "LICENSE"}

diff --git a/src/pytom_tm/entry_points.py b/src/pytom_tm/entry_points.py
@@ -212,6 +212,13 @@ def pytom_create_template(argv=None):
         default=2.7,
         help="Spherical aberration in mm.",
     )
+    parser.add_argument(
+        "--phase-shift",
+        type=float,
+        required=False,
+        default=.0,
+        help="Phase shift (in degrees) for the CTF to model phase plates.",
+    )
     parser.add_argument(
         "--cut-after-first-zero",
         action="store_true",
@@ -324,6 +331,7 @@ def pytom_create_template(argv=None):
             "spherical_aberration": args.Cs * 1e-3,
             "cut_after_first_zero": args.cut_after_first_zero,
             "flip_phase": args.flip_phase,
+            "phase_shift_deg": args.phase_shift,
         }
 
     template = generate_template_from_map(
@@ -568,38 +576,24 @@ def match_template(argv=None):
     parser = argparse.ArgumentParser(
         description="Run template matching. -- Marten Chaillet (@McHaillet)"
     )
-    parser.add_argument(
+    io_group = parser.add_argument_group("Template, search volume, and output")
+    io_group.add_argument(
         "-t",
         "--template",
         type=pathlib.Path,
         required=True,
         action=CheckFileExists,
         help="Template; MRC file.",
     )
-    parser.add_argument(
-        "-m",
-        "--mask",
-        type=pathlib.Path,
-        required=True,
-        action=CheckFileExists,
-        help="Mask with same box size as template; MRC file.",
-    )
-    parser.add_argument(
-        "--non-spherical-mask",
-        action="store_true",
-        required=False,
-        help="Flag to set when the mask is not spherical. It adds the required "
-        "computations for non-spherical masks and roughly doubles computation time.",
-    )
-    parser.add_argument(
+    io_group.add_argument(
         "-v",
         "--tomogram",
         type=pathlib.Path,
         required=True,
         action=CheckFileExists,
         help="Tomographic volume; MRC file.",
     )
-    parser.add_argument(
+    io_group.add_argument(
         "-d",
         "--destination",
         type=pathlib.Path,
@@ -608,52 +602,46 @@ def match_template(argv=None):
         action=CheckDirExists,
         help="Folder to store the files produced by template matching.",
     )
-    parser.add_argument(
-        "-a",
-        "--tilt-angles",
-        nargs="+",
-        type=str,
+    mask_group = parser.add_argument_group("Mask")
+    mask_group.add_argument(
+        "-m",
+        "--mask",
+        type=pathlib.Path,
         required=True,
-        action=ParseTiltAngles,
-        help="Tilt angles of the tilt-series, either the minimum and maximum values of "
-        "the tilts (e.g. --tilt-angles -59.1 60.1) or a .rawtlt/.tlt file with all the "
-        "angles (e.g. --tilt-angles tomo101.rawtlt). In case all the tilt angles are "
-        "provided a more elaborate Fourier space constraint can be used",
+        action=CheckFileExists,
+        help="Mask with same box size as template; MRC file.",
     )
-    parser.add_argument(
-        "--per-tilt-weighting",
+    mask_group.add_argument(
+        "--non-spherical-mask",
         action="store_true",
-        default=False,
         required=False,
-        help="Flag to activate per-tilt-weighting, only makes sense if a file with all "
-        "tilt angles have been provided. In case not set, while a tilt angle file is "
-        "provided, the minimum and maximum tilt angle are used to create a binary "
-        "wedge. The base functionality creates a fanned wedge where each tilt is "
-        "weighted by cos(tilt_angle). If dose accumulation and CTF parameters are "
-        "provided these will all be incorporated in the tilt-weighting.",
+        help="Flag to set when the mask is not spherical. It adds the required "
+             "computations for non-spherical masks and roughly doubles computation time.",
     )
-    parser.add_argument(
+    rotation_group = parser.add_argument_group('Angular search')
+    rotation_group.add_argument(
         "--angular-search",
         type=str,
         required=True,
         help="Options are: [7.00, 35.76, 19.95, 90.00, 18.00, "
-        "12.85, 38.53, 11.00, 17.86, 25.25, 50.00, 3.00].\n"
-        "Alternatively, a .txt file can be provided with three Euler angles "
-        "(in radians) per line that define the angular search. "
-        "Angle format is ZXZ anti-clockwise (see: "
-        "https://www.ccpem.ac.uk/user_help/rotation_conventions.php).",
+             "12.85, 38.53, 11.00, 17.86, 25.25, 50.00, 3.00].\n"
+             "Alternatively, a .txt file can be provided with three Euler angles "
+             "(in radians) per line that define the angular search. "
+             "Angle format is ZXZ anti-clockwise (see: "
+             "https://www.ccpem.ac.uk/user_help/rotation_conventions.php).",
     )
-    parser.add_argument(
+    rotation_group.add_argument(
         "--z-axis-rotational-symmetry",
         type=int,
         required=False,
         action=LargerThanZero,
         default=1,
         help="Integer value indicating the rotational symmetry of the template around "
-        "the z-axis. The length of the rotation search will be shortened through "
-        "division by this value. Only works for template symmetry around the z-axis.",
+             "the z-axis. The length of the rotation search will be shortened through "
+             "division by this value. Only works for template symmetry around the z-axis.",
     )
-    parser.add_argument(
+    volume_group = parser.add_argument_group('Volume control')
+    volume_group.add_argument(
         "-s",
         "--volume-split",
         nargs=3,
@@ -664,7 +652,7 @@ def match_template(argv=None):
         "can be relevant if the volume does not fit into GPU memory. "
         "Format is x y z, e.g. --volume-split 1 2 1",
     )
-    parser.add_argument(
+    volume_group.add_argument(
         "--search-x",
         nargs=2,
         type=int,
@@ -673,7 +661,7 @@ def match_template(argv=None):
         help="Start and end indices of the search along the x-axis, "
         "e.g. --search-x 10 490 ",
     )
-    parser.add_argument(
+    volume_group.add_argument(
         "--search-y",
         nargs=2,
         type=int,
@@ -682,7 +670,7 @@ def match_template(argv=None):
         help="Start and end indices of the search along the y-axis, "
         "e.g. --search-x 10 490 ",
     )
-    parser.add_argument(
+    volume_group.add_argument(
         "--search-z",
         nargs=2,
         type=int,
@@ -691,7 +679,32 @@ def match_template(argv=None):
         help="Start and end indices of the search along the z-axis, "
         "e.g. --search-x 30 230 ",
     )
-    parser.add_argument(
+    filter_group = parser.add_argument_group('Filter control')
+    filter_group.add_argument(
+        "-a",
+        "--tilt-angles",
+        nargs="+",
+        type=str,
+        required=True,
+        action=ParseTiltAngles,
+        help="Tilt angles of the tilt-series, either the minimum and maximum values of "
+             "the tilts (e.g. --tilt-angles -59.1 60.1) or a .rawtlt/.tlt file with all the "
+             "angles (e.g. --tilt-angles tomo101.rawtlt). In case all the tilt angles are "
+             "provided a more elaborate Fourier space constraint can be used",
+    )
+    filter_group.add_argument(
+        "--per-tilt-weighting",
+        action="store_true",
+        default=False,
+        required=False,
+        help="Flag to activate per-tilt-weighting, only makes sense if a file with all "
+             "tilt angles have been provided. In case not set, while a tilt angle file is "
+             "provided, the minimum and maximum tilt angle are used to create a binary "
+             "wedge. The base functionality creates a fanned wedge where each tilt is "
+             "weighted by cos(tilt_angle). If dose accumulation and CTF parameters are "
+             "provided these will all be incorporated in the tilt-weighting.",
+    )
+    filter_group.add_argument(
         "--voxel-size-angstrom",
         type=float,
         required=False,
@@ -700,7 +713,7 @@ def match_template(argv=None):
         "try to read from the MRC files. Argument is important for band-pass "
         "filtering!",
     )
-    parser.add_argument(
+    filter_group.add_argument(
         "--low-pass",
         type=float,
         required=False,
@@ -709,7 +722,7 @@ def match_template(argv=None):
         "if the template was already filtered to a certain resolution. "
         "Value is the resolution in A.",
     )
-    parser.add_argument(
+    filter_group.add_argument(
         "--high-pass",
         type=float,
         required=False,
@@ -719,7 +732,7 @@ def match_template(argv=None):
         "e.g. 500 could be appropriate as the CTF is often incorrectly modelled "
         "up to 50nm.",
     )
-    parser.add_argument(
+    filter_group.add_argument(
         "--dose-accumulation",
         type=str,
         required=False,
@@ -728,7 +741,7 @@ def match_template(argv=None):
         "tilt angle, assuming the same ordering of tilts as the tilt angle file. "
         "Format should be a .txt file with on each line a dose value in e-/A2.",
     )
-    parser.add_argument(
+    filter_group.add_argument(
         "--defocus-file",
         type=str,
         required=False,
@@ -742,28 +755,35 @@ def match_template(argv=None):
         "same ordering as tilt angle list. The .txt file should contain a single "
         "defocus value (in nm) per line.",
     )
-    parser.add_argument(
+    filter_group.add_argument(
         "--amplitude-contrast",
         type=float,
         required=False,
         action=BetweenZeroAndOne,
         help="Amplitude contrast fraction for CTF.",
     )
-    parser.add_argument(
+    filter_group.add_argument(
         "--spherical-abberation",
         type=float,
         required=False,
         action=LargerThanZero,
         help="Spherical abberation for CTF in mm.",
     )
-    parser.add_argument(
+    filter_group.add_argument(
         "--voltage",
         type=float,
         required=False,
         action=LargerThanZero,
         help="Voltage for CTF in keV.",
     )
-    parser.add_argument(
+    filter_group.add_argument(
+        "--phase-shift",
+        type=float,
+        required=False,
+        action=LargerThanZero,
+        help="Phase shift (in degrees) for the CTF to model phase plates.",
+    )
+    filter_group.add_argument(
         "--spectral-whitening",
         action="store_true",
         default=False,
@@ -772,15 +792,17 @@ def match_template(argv=None):
         "apply it to the tomogram patch and template. Effectively puts more weight on "
         "high resolution features and sharpens the correlation peaks.",
     )
-    parser.add_argument(
+    device_group = parser.add_argument_group('Device control')
+    device_group.add_argument(
         "-g",
         "--gpu-ids",
         nargs="+",
         type=int,
         required=True,
         help="GPU indices to run the program on.",
     )
-    parser.add_argument(
+    debug_group = parser.add_argument_group('Logging/debugging')
+    debug_group.add_argument(
         "--log",
         type=str,
         required=False,
@@ -810,6 +832,7 @@ def match_template(argv=None):
                 "amplitude": args.amplitude_contrast,
                 "voltage": args.voltage,
                 "cs": args.spherical_abberation,
+                "phase_shift_deg": args.phase_shift,
             }
             for defocus in args.defocus_file
         ]

diff --git a/src/pytom_tm/tmjob.py b/src/pytom_tm/tmjob.py
@@ -61,6 +61,13 @@ def load_json_to_tmjob(file_name: pathlib.Path, load_for_extraction: bool = True
         pytom_tm_version_number=data.get('pytom_tm_version_number', '0.3.0'),
         job_loaded_for_extraction=load_for_extraction,
     )
+    # if the file originates from an old version set the phase shift for compatibility
+    if (
+        version.parse(job.pytom_tm_version_number) < version.parse('0.6.1') and
+        job.ctf_data is not None
+    ):
+        for tilt in job.ctf_data:
+            tilt['phase_shift_deg'] = .0
     job.rotation_file = pathlib.Path(data['rotation_file'])
     job.whole_start = data['whole_start']
     job.sub_start = data['sub_start']

diff --git a/src/pytom_tm/weights.py b/src/pytom_tm/weights.py
@@ -469,6 +469,7 @@ def _create_tilt_weighted_wedge(
         - 'amplitude'; fraction of amplitude contrast between 0 and 1
         - 'voltage'; in keV
         - 'cs'; spherical abberation in mm
+        - 'phase_shift_deg'; phase shift for phase plates in deg
 
     Returns
     -------
@@ -514,7 +515,8 @@ def _create_tilt_weighted_wedge(
                     ctf_params_per_tilt[i]['amplitude'],
                     ctf_params_per_tilt[i]['voltage'] * 1e3,
                     ctf_params_per_tilt[i]['cs'] * 1e-3,
-                    flip_phase=True  # creating a per tilt ctf is hard if the phase is not flipped
+                    flip_phase=True,  # creating per tilt ctf requires phase flip atm
+                    phase_shift_deg=ctf_params_per_tilt[i]['phase_shift_deg'],
                 ), axes=0,
             )
             tilt[:, :, image_size // 2] = np.concatenate(
@@ -571,7 +573,8 @@ def create_ctf(
         voltage: float,
         spherical_aberration: float,
         cut_after_first_zero: bool = False,
-        flip_phase: bool = False
+        flip_phase: bool = False,
+        phase_shift_deg: float = .0,
 ) -> npt.NDArray[float]:
     """Create a CTF in a 3D volume in reduced format.
 
@@ -593,6 +596,11 @@ def create_ctf(
         whether to cut ctf after first zero crossing
     flip_phase: bool, default False
         make ctf fully positive/negative to imitate ctf correction by phase flipping
+    phase_shift_deg: float, default .0
+        additional phase shift to model phase plates, similar to
+        `https://github.com/dtegunov/tom_deconv` except the ctf defintion in tom
+        produces the inverse curve of what we have here
+
 
     Returns
     -------
@@ -609,7 +617,7 @@ def create_ctf(
     tan_term = np.arctan(amplitude_contrast / np.sqrt(1 - amplitude_contrast ** 2))
 
     # determine the ctf
-    ctf = - np.sin(chi + tan_term)
+    ctf = - np.sin(chi + tan_term + np.deg2rad(phase_shift_deg))
 
     if cut_after_first_zero:  # find frequency to cut first zero
         def chi_1d(q):