From c96800d8b36ce670f406b56a1776e0fe2410e400 Mon Sep 17 00:00:00 2001
From: EuanPyle <euanpyle@gmail.com>
Date: Wed, 1 Jun 2022 15:29:59 +0100
Subject: [PATCH] Add tiltcor

---
 lil_aretomo/aretomo.py | 8 ++++++--
 lil_aretomo/utils.py   | 7 ++++++-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/lil_aretomo/aretomo.py b/lil_aretomo/aretomo.py
index 1b3f869..e1dbb66 100644
--- a/lil_aretomo/aretomo.py
+++ b/lil_aretomo/aretomo.py
@@ -21,7 +21,8 @@ def run_aretomo_alignment(
         target_pixel_size: Optional[float] = 10,
         nominal_rotation_angle: Optional[float] = None,
         n_patches_xy: Optional[tuple[int, int]] = (5, 4),
-        thickness_for_alignment: Optional[float] = 800
+        thickness_for_alignment: Optional[float] = 800,
+        correct_tilt_angle_offset: Optional[bool] = False	
 ):
     """Run aretomo alignment on a single tilt-series
     
@@ -42,6 +43,8 @@ def run_aretomo_alignment(
     (optional) thickness_for_alignment: thickness in Z in unbinned pixels for which AreTomo will use in the alignment.
     This is useful is there is a lot of empty space at the top and bottom of your tomogram.
     See AreTomo manual for full explanation: this sets -AlignZ. Default is 800.
+    (optional) correct_tilt_angle_offset: Apply tilt angle offset correction, yes or no.
+    Default is no. See AreTomo manual for full explanation: yes; adds the -TiltCor 1 argument
     """
     if check_aretomo_availability() is False:
         e = 'AreTomo is not available for use. Load AreTomo so it can be called from the terminal via: AreTomo.'
@@ -67,5 +70,6 @@ def run_aretomo_alignment(
         nominal_rotation_angle=nominal_rotation_angle,
         local_alignments=local_align,
         n_patches_xy=n_patches_xy,
-        thickness_for_alignment=thickness_for_alignment
+        thickness_for_alignment=thickness_for_alignment,
+        correct_tilt_angle_offset=correct_tilt_angle_offset
     )
diff --git a/lil_aretomo/utils.py b/lil_aretomo/utils.py
index 9cb1a47..d7a83c5 100644
--- a/lil_aretomo/utils.py
+++ b/lil_aretomo/utils.py
@@ -33,7 +33,8 @@ def align_tilt_series_aretomo(
         nominal_rotation_angle: bool or float,
         local_alignments: bool,
         n_patches_xy: tuple[int, int],
-        thickness_for_alignment: float
+        thickness_for_alignment: float,
+        correct_tilt_angle_offset: bool
 ):
     output_file_name = Path(
         f'{output_directory}/{tilt_series_file.stem}_aln{tilt_series_file.suffix}')
@@ -60,6 +61,10 @@ def align_tilt_series_aretomo(
         command.append(f'{n_patches_xy[0]}')
         command.append(f'{n_patches_xy[1]}')
 
+    if correct_tilt_angle_offset:
+        command.append('-TiltCor')
+        command.append('1')
+
     subprocess.run(command)
 
     # Rename .tlt