Add check for aln file, skip if exists (#39)

* Add check for aln file, skip if exists * No longer skips by default * AB suggested changes * Update README.md * Final changes pre-merge * only write tilt-series file if data shapes differ Co-authored-by: alisterburt <[email protected]>
teamtomo · Sep 1, 2022 · 0c3db23 · 0c3db23
1 parent 06d4f71
commit 0c3db23
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -27,6 +27,7 @@ align_tilt_series(
     correct_tilt_angle_offset=True,
     basename='TS_01',  # basename for files passed to AreTomo
     output_directory='TS_01_AreTomo',
+    skip_if_completed=False # set to True to skip if results from a previous alignment are present
 )
 ```
 

diff --git a/lil_aretomo/aretomo.py b/lil_aretomo/aretomo.py
@@ -21,7 +21,8 @@ def align_tilt_series(
         output_pixel_size: Optional[float] = 10,
         nominal_rotation_angle: Optional[float] = None,
         correct_tilt_angle_offset: bool = False,
-        gpu_ids: Optional[Sequence[int]] = None
+        gpu_ids: Optional[Sequence[int]] = None,
+        skip_if_completed: bool = False
 ) -> AreTomoOutput:
     """Align a single-axis tilt-series using AreTomo.
 
@@ -39,7 +40,8 @@ def align_tilt_series(
         CCW is positive.
     n_patches_xy: number of patches in each dimension to use for local alignments.
     correct_tilt_angle_offset: flag controlling whether or not to correct for sample tilt in the output.
-    gpu_ids: integer ids for GPUs on the system (zero indexed)
+    gpu_ids: integer ids for GPUs on the system (zero indexed).
+    skip_if_completed: skip alignment if previous results found.
     """
     if check_aretomo_on_path() is False:
         raise RuntimeError("AreTomo executable was not found. \
@@ -67,7 +69,11 @@ def align_tilt_series(
         n_patches_xy=n_patches_xy,
         gpu_ids=gpu_ids,
     )
-    with open(output_directory / 'log.txt', mode='w') as log:
-        subprocess.run(command, stdout=log, stderr=log)
-    return AreTomoOutput(tilt_series_file=tilt_series_file, reconstruction_file=reconstruction_file)
+    aretomo_output = AreTomoOutput(tilt_series_file=tilt_series_file, reconstruction_file=reconstruction_file)
+    if aretomo_output.contains_alignment_results is False or skip_if_completed is False:
+        with open(output_directory / 'log.txt', mode='w') as log:
+            subprocess.run(command, stdout=log, stderr=log)
+        if aretomo_output.contains_alignment_results is False:
+            raise RuntimeError(f'{basename} failed to align correctly.')
+    return aretomo_output
 
diff --git a/lil_aretomo/output.py b/lil_aretomo/output.py
@@ -17,3 +17,7 @@ def basename(self) -> str:
     @property
     def aln_file(self) -> Path:
         return self.directory / f'{self.basename}.aln'
+
+    @property
+    def contains_alignment_results(self) -> bool:
+        return self.aln_file.exists() 
diff --git a/lil_aretomo/utils.py b/lil_aretomo/utils.py
@@ -19,11 +19,16 @@ def prepare_output_directory(
     directory.mkdir(exist_ok=True, parents=True)
 
     tilt_series_file = directory / f'{basename}.mrc'
-    mrcfile.write(
-        tilt_series_file,
-        tilt_series.astype(np.float32),
-        voxel_size=(pixel_size, pixel_size, 1)
-    )
+    if tilt_series_file.exists():
+        with mrcfile.open(tilt_series_file, header_only=True) as mrc:
+            data_on_disk_shape = (mrc.header.nz, mrc.header.ny, mrc.header.nx)
+    if not np.array_equal(tilt_series.shape, data_on_disk_shape):
+        mrcfile.write(
+            tilt_series_file,
+            tilt_series.astype(np.float32),
+            voxel_size=(pixel_size, pixel_size, 1),
+            overwrite=True
+        )
 
     rawtlt_file = directory / f'{basename}.rawtlt'
     np.savetxt(rawtlt_file, tilt_angles, fmt='%.2f', delimiter='')