Merge 277a732 into 5d44a50

shimmingtoolbox/unwrap/prelude.py

@@ -10,6 +10,7 @@
 import pathlib
 import tempfile
 import logging
+import numpy as np
 
 from shimmingtoolbox.utils import run_subprocess
 
@@ -80,4 +81,9 @@ def prelude(wrapped_phase, mag, affine, mask=None, threshold=None, is_unwrapping
 
     fname_phase_unwrapped = glob.glob(os.path.join(path_tmp, 'rawPhase_unwrapped*'))[0]
 
-    return nib.load(fname_phase_unwrapped).get_fdata()
+    phase_unwrapped = nib.load(fname_phase_unwrapped).get_fdata()
+
+    for _ in range(wrapped_phase.ndim - phase_unwrapped.ndim):
+        phase_unwrapped = np.expand_dims(phase_unwrapped, -1)
+
+    return phase_unwrapped

test/test_prelude.py

@@ -168,3 +168,42 @@ def test_threshold(self):
         """
         unwrapped_phase_e1 = prelude(self.phase_e1, self.mag_e1, self.affine_phase_e1, threshold=200)
         assert(unwrapped_phase_e1.shape == self.phase_e1.shape)
+
+    def test_3rd_dim_singleton(self):
+        """
+        Call prelude on data with a singleton on the z dimension
+        """
+
+        # Prepare singleton
+        phase_singleton = np.expand_dims(self.phase_e1[..., 0], -1)
+        mag_singleton = np.expand_dims(self.mag_e1[..., 0], -1)
+
+        unwrapped_phase_singleton = prelude(phase_singleton, mag_singleton, self.affine_phase_e1)
+
+        assert unwrapped_phase_singleton.ndim == 3
+
+    def test_2nd_dim_singleton(self):
+        """
+        Call prelude on data with a singleton on the 2nd dimension
+        """
+
+        # Prepare singleton
+        phase_singleton = np.expand_dims(self.phase_e1[:, 0, 0], -1)
+        mag_singleton = np.expand_dims(self.mag_e1[:, 0, 0], -1)
+
+        unwrapped_phase_singleton = prelude(phase_singleton, mag_singleton, self.affine_phase_e1)
+
+        assert unwrapped_phase_singleton.ndim == 2
+
+    def test_2nd_and_3rd_dim_singleton(self):
+        """
+        Call prelude on data with a singleton on the 2nd and 3rd dimension
+        """
+
+        # Prepare singleton
+        phase_singleton = np.expand_dims(np.expand_dims(self.phase_e1[:, 0, 0], -1), -1)
+        mag_singleton = np.expand_dims(np.expand_dims(self.mag_e1[:, 0, 0], -1), -1)
+
+        unwrapped_phase_singleton = prelude(phase_singleton, mag_singleton, self.affine_phase_e1)
+
+        assert unwrapped_phase_singleton.ndim == 3