fix(test): adjust custom function tests

src/extraction/core/functions/custom/localisation.py

@@ -129,7 +129,8 @@ def nuc_est_conv(
 
 
 def nuc_conv_3d(cell_mask, trap_image, pixel_size=0.23, spacing=0.6):
-    cell_mask = np.dstack([cell_mask] * trap_image.shape[-1])
+    print(cell_mask.shape, trap_image.shape)
+    cell_mask = np.stack([cell_mask] * trap_image.shape[0])
     ratio = spacing / pixel_size
     cell_fluo = trap_image[cell_mask]
     num_cell_fluo = len(np.nonzero(cell_fluo)[0])

tests/aliby/network/extraction/test_functions.py

@@ -11,14 +11,14 @@ ext.load_funs()
 
 def test_custom_output():
     self = ext
-    mask = np.zeros((6, 6, 2), dtype=bool)
-    mask[2:4, 2:4, 0] = True
-    mask[3:5, 3:5, 1] = True
-    img = np.random.randint(1, 11, size=6**2 * 5).reshape(6, 6, 5)
+    mask = np.zeros((2, 6, 6), dtype=bool)
+    mask[0, 2:4, 2:4] = True
+    mask[1, 3:5, 3:5] = True
+    img = np.random.randint(1, 11, size=6**2 * 5).reshape(5, 6, 6)
 
     for i, f in self._custom_funs.items():
         if "3d" in i:
             res = f(mask, img)
         else:
-            res = f(mask, np.maximum.reduce(img, axis=2))
-        assert len(res) == mask.shape[2], "Output doesn't match input"
+            res = f(mask, np.maximum.reduce(img, axis=0))
+        assert len(res) == mask.shape[0], "Output doesn't match input"