From 33e1acb17a8b9684f27ca2afe232aa9c4b432d60 Mon Sep 17 00:00:00 2001
From: Dimitrios Bellos <34098719+DimitriosBellos@users.noreply.github.com>
Date: Fri, 26 Nov 2021 10:52:39 +0000
Subject: [PATCH] The guanaco.correct_file functionality can now be used as
 parakeet.analyse.correct (#19)

---
 setup.py                             |  1 +
 src/parakeet/analyse.py              | 45 +++++++++++++++-
 src/parakeet/command_line/analyse.py | 79 ++++++++++++++++++++++++++++
 3 files changed, 124 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index c9472ba..9c8f5dc 100644
--- a/setup.py
+++ b/setup.py
@@ -101,6 +101,7 @@ def main():
                 "parakeet.analyse.average_all_particles=parakeet.command_line.analyse:average_all_particles",
                 "parakeet.analyse.extract_particles=parakeet.command_line.analyse:extract_particles",
                 "parakeet.analyse.refine=parakeet.command_line.analyse:refine",
+                "parakeet.analyse.correct=parakeet.command_line.analyse:correct",
             ]
         },
         extras_require={
diff --git a/src/parakeet/analyse.py b/src/parakeet/analyse.py
index 79bb3a3..ea95ec2 100644
--- a/src/parakeet/analyse.py
+++ b/src/parakeet/analyse.py
@@ -26,7 +26,7 @@
 
 def reconstruct(image_filename, rec_filename, microscope, simulation, device="gpu"):
     """
-    Reconstruct the volume using 3D CTF correction
+    Reconstruct the volume and use 3D CTF correction beforehand if the input image is uncorrected
 
     """
 
@@ -72,6 +72,49 @@ def reconstruct(image_filename, rec_filename, microscope, simulation, device="gp
     )
 
 
+def correct(image_filename, corrected_filename, microscope, simulation, device="gpu"):
+    """
+    Correct the images using 3D CTF correction
+
+    """
+
+    # Ensure mrc file
+    assert os.path.splitext(image_filename)[1] == ".mrc"
+
+    # Get the parameters for the CTF correction
+    nx = microscope.detector.nx
+    pixel_size = microscope.detector.pixel_size
+    energy = microscope.beam.energy
+    defocus = -microscope.lens.c_10
+    num_defocus = int((nx * pixel_size) / 100)
+
+    # Set the spherical aberration
+    if simulation["inelastic_model"] == "cc_corrected":
+        print("Setting spherical aberration to zero")
+        spherical_aberration = 0
+    else:
+        spherical_aberration = microscope.lens.c_30
+
+    astigmatism = microscope.lens.c_12
+    astigmatism_angle = microscope.lens.phi_12
+    phase_shift = 0
+
+    # Do the reconstruction
+    guanaco.correct_file(
+        input_filename=image_filename,
+        output_filename=corrected_filename,
+        centre=None,
+        energy=energy,
+        defocus=defocus,
+        num_defocus=num_defocus,
+        spherical_aberration=spherical_aberration,
+        astigmatism=astigmatism,
+        astigmatism_angle=astigmatism_angle,
+        phase_shift=phase_shift,
+        device=device,
+    )
+
+
 def average_particles(
     scan,
     sample_filename,
diff --git a/src/parakeet/command_line/analyse.py b/src/parakeet/command_line/analyse.py
index 2a4fa1e..73f478b 100644
--- a/src/parakeet/command_line/analyse.py
+++ b/src/parakeet/command_line/analyse.py
@@ -101,6 +101,85 @@ def reconstruct(args=None):
     logger.info("Time taken: %.2f seconds" % (time.time() - start_time))
 
 
+def correct(args=None):
+    """
+    Correct the images using 3D CTF correction
+
+    """
+
+    # Get the start time
+    start_time = time.time()
+
+    # Create the argument parser
+    parser = argparse.ArgumentParser(description="3D CTF correction of the images")
+
+    # Add some command line arguments
+    parser.add_argument(
+        "-c",
+        "--config",
+        type=str,
+        default=None,
+        dest="config",
+        help="The yaml file to configure the simulation",
+    )
+    parser.add_argument(
+        "-d",
+        "--device",
+        choices=["cpu", "gpu"],
+        default=None,
+        dest="device",
+        help="Choose the device to use",
+    )
+    parser.add_argument(
+        "-i",
+        "--image",
+        type=str,
+        default="image.mrc",
+        dest="image",
+        help="The filename for the image",
+    )
+    parser.add_argument(
+        "-cr",
+        "--corrected",
+        type=str,
+        default="corrected_image.mrc",
+        dest="corrected",
+        help="The filename for the corrected image",
+    )
+
+    # Parse the arguments
+    args = parser.parse_args(args=args)
+
+    # Configure some basic logging
+    parakeet.command_line.configure_logging()
+
+    # Set the command line args in a dict
+    command_line = {}
+    if args.device is not None:
+        command_line["device"] = args.device
+
+    # Load the full configuration
+    config = parakeet.config.load(args.config, command_line)
+
+    # Print some options
+    parakeet.config.show(config)
+
+    # Create the microscope
+    microscope = parakeet.microscope.new(**config["microscope"])
+
+    # Do the reconstruction
+    parakeet.analyse.correct(
+        args.image,
+        args.corrected,
+        microscope=microscope,
+        simulation=config["simulation"],
+        device=config["device"],
+    )
+
+    # Write some timing stats
+    logger.info("Time taken: %.2f seconds" % (time.time() - start_time))
+
+
 def average_particles(args=None):
     """
     Perform sub tomogram averaging