From 79cf7ffeed5f5f7cd8a4c185e4f824d2596db841 Mon Sep 17 00:00:00 2001
From: Andrew Herzing <andrew.herzing@nist.gov>
Date: Tue, 3 Oct 2023 10:00:32 -0400
Subject: [PATCH] Cleaned up load_dm and load_dm_series in io.py

---
 tomotools/io.py            | 32 +++++++-------------------------
 tomotools/tests/test_io.py |  4 ++--
 2 files changed, 9 insertions(+), 27 deletions(-)

diff --git a/tomotools/io.py b/tomotools/io.py
index e4c0eaf..775add2 100644
--- a/tomotools/io.py
+++ b/tomotools/io.py
@@ -127,7 +127,7 @@ def _get_manual_tilts():
     return stack
 
 
-def loadhspy(filename, tilts=None):
+def load_hspy(filename, tilts=None):
     """
     Read an MRC file to a TomoStack object using the Hyperspy reader.
 
@@ -229,7 +229,7 @@ def loadhspy(filename, tilts=None):
     return stack
 
 
-def loaddm(filename):
+def load_dm(filename):
     """
     Read DM image series.
 
@@ -274,19 +274,8 @@ def loaddm(filename):
 
     tilts = np.arange(mintilt, maxtilt + tiltstep, tiltstep)
 
-    s_new = TomoStack(s)
-    s_new.axes_manager[0].name = 'Tilt'
-    s_new.axes_manager[0].units = 'degrees'
-    s_new.axes_manager[0].offset = tilts[0]
-    s_new.axes_manager[0].scale = tilts[1] - tilts[0]
-    logger.info('Tilts found in metadata')
-
-    s_new.metadata.Tomography.tilts = tilts
-    s_new.metadata.Tomography.shifts = np.zeros([s_new.data.shape[0], 2])
-    s_new.metadata.Tomography.tiltaxis = 0.0
-    s_new.metadata.Tomography.xshift = 0.0
-
-    return s_new
+    s = convert_to_tomo_stack(s, tilts)
+    return s
 
 
 def load_dm_series(input_data):
@@ -328,14 +317,7 @@ def load_dm_series(input_data):
     files_sorted = list(np.array(files)[sorted_order])
     del s
     stack = hspy.load(files_sorted, stack=True, new_axis_name='tilt')
-    stack.metadata.Acquisition_instrument.TEM.Stage.tilt_alpha\
-        = tilts_sorted
-    stack.axes_manager[0].scale = tilts_sorted[1] - tilts_sorted[0]
-    stack.axes_manager[0].units = 'degrees'
-    stack.axes_manager[0].offset = tilts_sorted[0]
-
-    stack = convert_to_tomo_stack(stack)
-    stack.metadata.Tomography.tilts = tilts_sorted
+    stack = convert_to_tomo_stack(stack, tilts_sorted)
     return stack
 
 
@@ -390,9 +372,9 @@ def load(filename, tilts=None):
     ext = os.path.splitext(filename)[1]
     if ext in ['.HDF5', '.hdf5', '.hd5', '.HD5', '.MRC', '.mrc', '.ALI',
                '.ali', '.REC', '.rec', '.hspy', '.HSPY']:
-        stack = loadhspy(filename, tilts)
+        stack = load_hspy(filename, tilts)
     elif ext in ['.dm3', '.DM3', '.dm4', '.DM4']:
-        stack = loaddm(filename)
+        stack = load_dm(filename)
     else:
         raise ValueError("Unknown file type")
     return stack
diff --git a/tomotools/tests/test_io.py b/tomotools/tests/test_io.py
index a822e11..e7ad0c5 100644
--- a/tomotools/tests/test_io.py
+++ b/tomotools/tests/test_io.py
@@ -14,7 +14,7 @@ class TestHspy:
     def test_load_hspy_mrc(self):
         filename = os.path.join(tomotools_path, "tests",
                                 "test_data", "HAADF.mrc")
-        stack = tomotools.io.loadhspy(filename)
+        stack = tomotools.io.load_hspy(filename)
         assert type(stack) is TomoStack
         assert stack.axes_manager.signal_shape == (256, 256)
         assert stack.axes_manager.navigation_shape == (77,)
@@ -23,7 +23,7 @@ def test_load_hspy_mrc(self):
     def test_load_hspy_hdf5(self):
         filename = os.path.join(tomotools_path, "tests",
                                 "test_data", "HAADF_Aligned.hdf5")
-        stack = tomotools.io.loadhspy(filename)
+        stack = tomotools.io.load_hspy(filename)
         with h5py.File(filename, 'r') as h5:
             h5_shape = h5['Experiments']['__unnamed__']['data'].shape
         assert type(stack) is TomoStack