Update contrast_mask to be able to handle timesteps

canismarko · Aug 12, 2019 · 3d21a63 · 3d21a63
1 parent 59e4d9a
commit 3d21a63
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/tests/test_math.py b/tests/test_math.py
@@ -480,11 +480,12 @@ def test_extract_signals_pca(self):
     def test_contrast_mask(self):
 
         # The first time idx - User selected in the xp.XanesFrameset.frames() method
-        frames = self.coins()[0]
+        frames = self.coins()
 
+        print(np.shape(frames))
         # Set up the initial numpy frames
-        mean_frames_image = np.mean(frames, axis=0)
-        single_frame_image = frames[10]
+        mean_frames_image = np.mean(frames, axis=(0, 1))
+        single_frame_image = frames[0][10]
 
         # Check difference input values
         sensitivity_vals = [1, 0.5, 1.8]
@@ -520,7 +521,7 @@ def test_contrast_mask(self):
                 single_frames_check.append(contrast_mask(frames=frames,
                                                        sensitivity=sensitivity,
                                                        min_size=min_size,
-                                                       frame_idx=10))
+                                                       frame_idx=(0, 10)))
 
         # Check all the values
         np.testing.assert_equal(all_masks[0], mean_frames_check[0])

diff --git a/xanespy/xanes_math.py b/xanespy/xanes_math.py
@@ -1114,7 +1114,7 @@ def contrast_mask(frames: np.ndarray, sensitivity: float = 1, min_size=0, frame_
     min_size : float, optional
       Objects below this size (in pixels) will be
       removed. Passing zero (default) will result in no effect.
-    frame_idx : str, int, optional
+    frame_idx : str, tuple(time_step_index, energy_index), optional
       Allows the user to select which image to
 
     Returns:
@@ -1129,10 +1129,11 @@ def contrast_mask(frames: np.ndarray, sensitivity: float = 1, min_size=0, frame_
 
     # Obtain the correct image set
     if frame_idx == 'mean':
-        image = np.mean(frames, axis=0)
+        image = np.mean(frames, axis=(0, 1))
 
-    elif isinstance(frame_idx, int):
-        image = frames[frame_idx]
+    elif isinstance(frame_idx, tuple):
+        time_idx, energy_idx = frame_idx
+        image = frames[time_idx][energy_idx]
 
     # Determining threshold
     img_bottom = image.min()