readd numba utils

Author: Jammy2211

autoarray/mask/mask_1d_util.py

@@ -1,6 +1,43 @@
 import numpy as np
 
+from autoarray import numba_util
 
+
+@numba_util.jit()
+def total_pixels_1d_from(mask_1d: np.ndarray) -> int:
+    """
+    Returns the total number of unmasked pixels in a mask.
+
+    Parameters
+    ----------
+    mask_1d
+        A 2D array of bools, where `False` values are unmasked and included when counting pixels.
+
+    Returns
+    -------
+    int
+        The total number of pixels that are unmasked.
+
+    Examples
+    --------
+
+    mask = np.array([[True, False, True],
+                 [False, False, False]
+                 [True, False, True]])
+
+    total_regular_pixels = total_regular_pixels_from(mask=mask)
+    """
+
+    total_regular_pixels = 0
+
+    for x in range(mask_1d.shape[0]):
+        if not mask_1d[x]:
+            total_regular_pixels += 1
+
+    return total_regular_pixels
+
+
+@numba_util.jit()
 def native_index_for_slim_index_1d_from(
     mask_1d: np.ndarray,
 ) -> np.ndarray:
@@ -25,12 +62,22 @@ def native_index_for_slim_index_1d_from(
 
     Examples
     --------
-    >>> mask_1d = np.array([True, False, True, False, False, True])
-    >>> native_index_for_slim_index_1d_from(mask_1d)
-    array([1, 3, 4])
+    mask_2d = np.array([[True, True, True],
+                     [True, False, True]
+                     [True, True, True]])
+
+    native_index_for_slim_index_1d =  native_index_for_slim_index_1d_from(mask_2d=mask_2d)
 
     """
-    # Create an array of native indexes corresponding to unmasked pixels
-    native_index_for_slim_index_1d = np.flatnonzero(~mask_1d)
 
-    return native_index_for_slim_index_1d
+    total_pixels = total_pixels_1d_from(mask_1d=mask_1d)
+    native_index_for_slim_index_1d = np.zeros(shape=total_pixels)
+
+    slim_index = 0
+
+    for x in range(mask_1d.shape[0]):
+        if not mask_1d[x]:
+            native_index_for_slim_index_1d[slim_index] = x
+            slim_index += 1
+
+    return native_index_for_slim_index_1d