Merge branch 'master' of github.com:rosalindfranklininstitute/amplus-…

…digital-twin

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={

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,

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