Merge pull request #87 from rosalindfranklininstitute/issue85

Update to AreTomo 1.3+

conda-build/meta.yaml

@@ -1,6 +1,6 @@
 package:
   name: ot2rec
-  version: 0.3.1
+  version: 0.3.2
 
 source:
   path: ..

construct.yaml

@@ -1,10 +1,10 @@
 name: Ot2Rec
-version: 0.3.1
+version: 0.3.2
 company: Rosalind Franklin Institute
 channels:
   - https://conda.anaconda.org/conda-forge
   - file:///usr/share/miniconda3/conda-bld/
 specs:
   - python 3.9.*
-  - Ot2Rec 0.3.1
+  - Ot2Rec 0.3.2
   - conda

setup.py

@@ -17,7 +17,7 @@
 
 
 setup(
-    version='0.3.1',
+    version='0.3.2',
     name='Ot2Rec',
     description='Ot2Rec',
     url='https://github.com/rosalindfranklininstitute/Ot2Rec',

src/Ot2Rec/aretomo.py

@@ -15,6 +15,7 @@
 
 import argparse
 import os
+import shutil
 import subprocess
 import warnings
 from glob import glob
@@ -59,6 +60,8 @@ def __init__(self,
 
         self.md_out = {}
 
+        self.sta = {}
+
         self._get_internal_metadata()
 
     def _get_internal_metadata(self):
@@ -90,7 +93,15 @@ def _get_internal_metadata(self):
                 self.md_out["aretomo_output_dir"] = {}
                 self.md_out["aretomo_align_stats"] = {}
             self.md_out["aretomo_output_dir"][curr_ts] = subfolder
-            self.md_out["aretomo_align_stats"][curr_ts] = subfolder + f"/{self.rootname}_{curr_ts:04d}{self.suffix}.st.aln"
+            self.md_out["aretomo_align_stats"][curr_ts] = (
+                f"{subfolder}/"
+                f"{self.rootname}_{curr_ts:04d}{self.suffix}.st.aln"
+            )
+
+        if self.params["AreTomo_setup"]["out_imod"] != "N/A":
+            self.sta_folder = f"{self.basis_folder}/STA"
+            self.md_out["aretomo_STA_dir"] = self.sta_folder
+            os.makedirs(self.sta_folder, exist_ok=True)
 
     def _get_aretomo_align_command(self, i):
         """
@@ -111,6 +122,8 @@ def _get_aretomo_align_command(self, i):
             '0',
             '-OutBin',
             str(self.params['AreTomo_setup']['output_binning']),
+            '-DarkTol',
+            str(self.params['AreTomo_setup']['dark_tol']),
         ]
 
         return cmd
@@ -134,18 +147,16 @@ def _get_aretomo_recon_command(self, i):
             str(self.params['AreTomo_recon']['volz']),
             '-OutBin',
             str(self.params['AreTomo_setup']['output_binning']),
-            '-Align',
-            '0'
         ]
 
+        if self.params['AreTomo_setup']['aretomo_mode'] == 1:
+            cmd.append('-Align')
+            cmd.append('0')
+
         if self.params['AreTomo_recon']['recon_algo'] == "WBP":
             # WBP
             cmd.append('-Wbp')
             cmd.append('1')
-        elif self.params['AreTomo_recon']['recon_algo'] == "SART":
-            # SART
-            cmd.append('-Wbp')
-            cmd.append('0')
 
         return cmd
 
@@ -163,6 +174,16 @@ def _run_aretomo(self, i):
             recon_cmd = self._get_aretomo_recon_command(i)
             cmd = recon_cmd
 
+        out_imod = self.params['AreTomo_setup']['out_imod']
+        if out_imod != "N/A":
+            outimod_lookup = {
+                "RELION4": "1",
+                "Warp": "2",
+                "Local alignment": "3",
+            }
+            cmd.append('-OutImod')
+            cmd.append(outimod_lookup[out_imod])
+
         # Add extra kwargs
         kwargs = self.params["AreTomo_kwargs"].keys()
         for kwarg in kwargs:
@@ -175,12 +196,23 @@ def _run_aretomo(self, i):
         self.md_out["aretomo_cmd"][curr_ts] = " ".join(cmd)
 
         # Run aretomo
-        aretomo_run = subprocess.run(cmd,
-                                     stdout=subprocess.PIPE,
-                                     stderr=subprocess.STDOUT,
-                                     encoding='ascii',
-                                     check=True,
-                                     )
+        aretomo_run = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            encoding='ascii',
+            check=True,
+        )
+
+        # If STA files are generated save folder names to move to common folder
+        if self.params["AreTomo_setup"]["out_imod"] != "N/A":
+            output_mrc = self.params["AreTomo_setup"]["output_mrc"][i]
+            self.sta[curr_ts] = (
+                f'{self.basis_folder}/'
+                f'{self.rootname}_{curr_ts:04d}{self.suffix}/'
+                f'{os.path.splitext(os.path.basename(output_mrc))[0]}_Imod/'
+            )
+
         self.logObj(f"\nStdOut:{aretomo_run.stdout}\n")
         self.logObj(f"\nStdErr:{aretomo_run.stderr}\n")
 
@@ -199,6 +231,21 @@ def export_metadata(self):
         """
         Method to export metadata as yaml
         """
+        # If STA files created, move to common folder
+        if self.params["AreTomo_setup"]["out_imod"] != "N/A": 
+            for ts in list(self.sta.keys()):
+                sta_ts_folder = (
+                    f"{self.sta_folder}/{self.sta[ts].split(os.sep)[-2]}"
+                )
+                if os.path.exists(sta_ts_folder):
+                    self.logObj(
+                        f"STA folder {sta_ts_folder} not empty, overwriting."
+                    )
+                    shutil.rmtree(sta_ts_folder)
+                shutil.move(
+                    src=self.sta[ts],
+                    dst=self.sta_folder,
+                )
 
         yaml_file = self.proj_name + "_aretomo_mdout.yaml"
 
@@ -282,6 +329,7 @@ def _get_process_list(file_list, rootname, suffix, ext):
             )
     return ts_list
 
+
 def update_yaml(args):
     """
     Method to update yaml file for AreTomo
@@ -318,7 +366,7 @@ def update_yaml(args):
     if args["aretomo_mode"] != 1:  # for workflows with alignment
         # Set InMrc
         st_file_list = _find_files_with_ext(
-            ".st",
+            args["input_ext"],
             rootname,
             suffix,
             str(args["input_mrc_folder"])
@@ -357,7 +405,12 @@ def update_yaml(args):
         aretomo_params.params["AreTomo_setup"]["tilt_angles"] = tlt_file_list
 
         # Set process list
-        ts_list = _get_process_list(st_file_list, rootname, suffix, ".st")
+        ts_list = _get_process_list(
+            st_file_list,
+            rootname,
+            suffix,
+            args["input_ext"]
+        )
         aretomo_params.params["System"]["process_list"] = ts_list
 
         # Set output mrc

src/Ot2Rec/magicgui.py

@@ -772,7 +772,7 @@ def get_args_imod_route(
         "max": 2,
         "step": 1,
     },
-    rot_angle={"label": "Rotation angle of electron beam*, can be obtained from mdoc",
+    rot_angle={"label": "Rotation angle of electron beam, can be obtained from mdoc",
                "min": -180.0000,
                "max": 180.0000
     },
@@ -785,6 +785,10 @@ def get_args_imod_route(
         "label": "Folder containing input mrc's",
         "mode": "d",
     },
+    input_ext={
+        "label": "Extension of input mrc files",
+        "tooltip": ".st or .mrc usually"
+    },
     output_path={
         "label": "Path to output folder",
         "mode": "d",
@@ -809,6 +813,14 @@ def get_args_imod_route(
         "label": "Reconstruction algorithm",
         "choices": ["WBP", "SART"]
     },
+    out_imod={
+        "label": "Produce STA output for",
+        "choices": ["N/A", "RELION4", "Warp", "Local alignment"]
+    },
+    dark_tol={
+        "label": "Tolerance to remove dark images",
+        "tooltip": "Default 0.7, low number = fewer images removed",
+    }
 )
 def get_args_aretomo(
         project_name="",
@@ -818,12 +830,15 @@ def get_args_aretomo(
         rootname="",
         suffix="",
         input_mrc_folder=Path("./aretomo"),
+        input_ext=".st",
         output_path=Path("./aretomo"),
         tilt_angles="",
         volz=-1,
         sample_thickness=-1,
         output_binning=4,
-        recon_algo="WBP",
+        recon_algo="SART",
+        out_imod="N/A",
+        dark_tol=0.7,
 ):
     return locals()
 

src/Ot2Rec/params.py

@@ -326,14 +326,16 @@ def new_aretomo_yaml(args):
             "input_mrc": None,
             "output_mrc": None,
             "tilt_angles": None,
-            "output_binning": None
+            "output_binning": None,
+            "out_imod": args["out_imod"],
+            "dark_tol": args["dark_tol"]
         },
 
         "AreTomo_recon": {
             "volz": None,
             "sample_thickness": None,
             "pixel_size": args["pixel_size"],
-            "recon_algo": None,
+            "recon_algo": args["recon_algo"],
         },
 
         "AreTomo_kwargs": {

tests/test_aretomo_recon.py

@@ -114,3 +114,39 @@ def test_aretomo_called(self, aretomo_mock):
         self.assertTrue(aretomo_mock.called)
 
         tmpdir.cleanup()
+
+    @patch("subprocess.run")
+    def test_aretomo_STA_output_folders_created(self, aretomo_mock):
+        # Create expected input
+        tmpdir = self._create_expected_folder_structure()
+        os.chdir(tmpdir.name)
+        args = self._create_expected_input_args()
+        args["out_imod"] = "Warp"   
+
+        # Create yaml
+        aretomo.create_yaml(args)
+
+        # Read params
+        params = prmMod.read_yaml(
+            project_name="TS",
+            filename="./TS_aretomo_recon.yaml",
+        )
+
+        # Run
+        logger = logMod.Logger("./o2r_aretomo_recon.log")
+        aretomo_obj = aretomo.AreTomo(
+            project_name="TS",
+            params_in=params,
+            logger_in=logger
+        )
+        sta_folder = "./aretomo_recon/STA"
+        aretomo_obj._run_aretomo(0)
+        os.mkdir(f"{tmpdir.name}/aretomo_recon/TS_0001/TS_0001_ali_rec_Imod")
+
+        aretomo_obj.export_metadata()
+
+        self.assertTrue(
+            os.path.isdir(f"{sta_folder}/TS_0001_ali_rec_Imod")
+        )
+
+        tmpdir.cleanup()