Merge pull request #9 from EuanPyle/PR4

PR4

example/test.py

@@ -14,5 +14,4 @@
     local_align=True,
     #n_patches_xy=(5,6),
     thickness_for_alignment=600,
-    correct_tilt_angle_offset=True
 )

lil_aretomo/aretomo.py

@@ -2,7 +2,7 @@
 from typing import List, Optional
 
 from .utils import (
-    prepare_imod_directory,
+    prepare_output_directory,
     align_tilt_series_aretomo,
     find_binning_factor,
     check_aretomo_availability,
@@ -16,10 +16,9 @@ def run_aretomo_alignment(
         output_directory: Path,
         aretomo_executable: Optional[Path] = None,
         local_align: Optional[bool] = False,
-        target_pixel_size: Optional[float] = 10,
+        target_pixel_size: Optional[float] = 10,	
         nominal_rotation_angle: Optional[float] = None,
         n_patches_xy: Optional[tuple[int, int]] = (5, 4),
-        correct_tilt_angle_offset: Optional[bool] = False,
         thickness_for_alignment: Optional[float] = 800
 ):
     """Run aretomo alignment on a single tilt-series
@@ -30,27 +29,25 @@ def run_aretomo_alignment(
     tilt_series_file: file containing tilt-series images
     tilt_angles: nominal stage tilt-angles from the microscope.
     pixel_size: pixel size of the tilt-series in angstroms-per-pixel
-    output_directory: tilt-series directory for IMOD.
+    output_directory: tilt-series directory.
     aretomo_executable: path to the AreTomo executable file
     (optional) local_align: carry out local tilt series alignments? Yes or no, default is no
     (optional) target_pixel_size: the ideal pixel size at which TSA is carried out. Default is 10A
     (optional) nominal_rotation_angle: initial estimate for the rotation angle of the tilt
     axis. AreTomo does not need this information but it might help.
     (optional) n_patches_xy: if local_align is True, AreTomo will carry out patch tracking. 
     Specify the number of patches in X and Y here as tuple. Default is 5 in X,4 in Y.
-    (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.
     (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.
     """
     if check_aretomo_availability() is False and aretomo_executable is None:
         raise RuntimeError('AreTomo executable not found.')
 
-    prepare_imod_directory(
+    tilt_series_file=tilt_series_file = prepare_output_directory(
         tilt_series_file=tilt_series_file,
         tilt_angles=tilt_angles,
-        imod_directory=output_directory
+        output_directory=output_directory
     )
 
     binning = find_binning_factor(
@@ -60,12 +57,11 @@ def run_aretomo_alignment(
 
     align_tilt_series_aretomo(
         tilt_series_file=tilt_series_file,
-        imod_directory=output_directory,
-        binning=binning,
+        output_directory=output_directory,
+	    binning=binning,
         aretomo_executable=aretomo_executable,
         nominal_rotation_angle=nominal_rotation_angle,
         local_align=local_align,
         n_patches_xy=n_patches_xy,
-        correct_tilt_angle_offset=correct_tilt_angle_offset,
-        thickness_for_alignment=thickness_for_alignment
+        thickness_for_alignment=thickness_for_alignment	
     )

lil_aretomo/utils.py

@@ -6,52 +6,45 @@
 
 import numpy as np
 
-
-def prepare_imod_directory(
-        tilt_series_file: Path,
-        tilt_angles: List[float],
-        imod_directory: Path
+def prepare_output_directory(
+        tilt_series_file: Path, 
+        tilt_angles: List[float], 
+        output_directory: Path
 ):
     ts_dir_name = tilt_series_file.stem
-    imod_directory.mkdir(exist_ok=True, parents=True)
-
-    # Rename file .mrc if .st
-    if tilt_series_file.suffix == '.st':
-        tilt_series_file = tilt_series_file.with_suffix('.mrc')
-
-    tilt_series_file_for_imod = imod_directory / tilt_series_file.name
-    force_symlink(tilt_series_file.absolute(), tilt_series_file_for_imod)
-
-    rawtlt_file = imod_directory / f'{ts_dir_name}.rawtlt'
+    output_directory.mkdir(exist_ok=True, parents=True)
+
+    # Link tilt-series file into output directory
+    tilt_series_filename = tilt_series_file.with_suffix('.mrc').name
+    linked_tilt_series_file = output_directory / tilt_series_filename
+    force_symlink(tilt_series_file.absolute(), linked_tilt_series_file)
+
+    rawtlt_file = output_directory / f'{ts_dir_name}.rawtlt'
     np.savetxt(rawtlt_file, tilt_angles, fmt='%.2f', delimiter='')
-
+    return linked_tilt_series_file
+
+#####
 
 def align_tilt_series_aretomo(
-        tilt_series_file: Path,
-        imod_directory: Path,
+        tilt_series_file: Path, 
+        output_directory: Path, 
         binning: float,
         aretomo_executable: Path,
         nominal_rotation_angle: bool or float,
         local_align: bool,
         n_patches_xy: tuple[int, int],
-        correct_tilt_angle_offset: bool,
         thickness_for_alignment: float
-):
-    # Rename file .mrc if .st
-    if tilt_series_file.suffix == '.st':
-        tilt_series_file = tilt_series_file.with_suffix('.mrc')
-
-    output_file_name = Path(
-        f'{imod_directory}/{tilt_series_file.stem}_aln{tilt_series_file.suffix}')
-
-    # Run AreTomo
+):    
+    output_file_name = Path(f'{output_directory}/{tilt_series_file.stem}_aln{tilt_series_file.suffix}')
+
+	#Run AreTomo
     aretomo_command = [
         f'{str(aretomo_executable)}',
         '-InMrc', f'{tilt_series_file}',
         '-OutMrc', f'{output_file_name}',
         '-OutBin', f'{binning}',
-        '-AngFile', f'{imod_directory}/{tilt_series_file.stem}.rawtlt',
-        '-AlignZ', f'{thickness_for_alignment}',
+        '-AngFile', f'{output_directory}/{tilt_series_file.stem}.rawtlt',
+        '-AlignZ', f'{thickness_for_alignment}',	
         '-VolZ', '0',
         '-OutXF', '1'
     ]
@@ -63,26 +56,22 @@ def align_tilt_series_aretomo(
     if local_align:
         aretomo_command.append('-Patch')
         aretomo_command.append(f'{n_patches_xy[0]}')
-        aretomo_command.append(f'{n_patches_xy[1]}')
-
-    if correct_tilt_angle_offset:
-        aretomo_command.append('-TiltCor')
-        aretomo_command.append('1')
+        aretomo_command.append(f'{n_patches_xy[1]}')        
 
     subprocess.run(aretomo_command)
 
-    # Rename .tlt
-    tlt_file_name = Path(f'{imod_directory}/{tilt_series_file.stem}_aln.tlt')
+    #Rename .tlt
+    tlt_file_name = Path(f'{output_directory}/{tilt_series_file.stem}_aln.tlt')
     new_tlt_stem = tlt_file_name.stem[:-4]
-    new_output_name_tlt = Path(f'{imod_directory}/{new_tlt_stem}').with_suffix('.tlt')
+    new_output_name_tlt = Path(f'{output_directory}/{new_tlt_stem}').with_suffix('.tlt')
     tlt_file_name.rename(new_output_name_tlt)
 
 
 def find_binning_factor(
         pixel_size: float,
         target_pixel_size: float
 ) -> int:
-    """Find closest power of two binning factor to reach target pixel size."""
+    #Find closest binning to reach target pixel size
     factors = 2 ** np.arange(7)
     binned_pixel_sizes = factors * pixel_size
     delta_pixel = np.abs(binned_pixel_sizes - target_pixel_size)