over sampling unit tests

Author: Jammy2211

autoarray/operators/over_sampling/over_sampler.py

@@ -147,6 +147,16 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]):
             over_sample_size=sub_size, mask=mask
         )
 
+
+    @property
+    def sub_is_uniform(self) -> bool:
+        """
+        Returns True if the sub_size is uniform across all pixels in the mask.
+        """
+        return np.all(
+            np.isclose(self.sub_size.array, self.sub_size.array[0])
+        )
+
     def tree_flatten(self):
         return (self.mask, self.sub_size), ()
 
@@ -185,7 +195,7 @@ def sub_pixel_areas(self) -> np.ndarray:
         """
         The area of every sub-pixel in the mask.
         """
-        sub_pixel_areas = jnp.zeros(self.sub_total)
+        sub_pixel_areas = np.zeros(self.sub_total)
 
         k = 0
 
@@ -221,15 +231,24 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D":
         except AttributeError:
             pass
 
-        # binned_array_2d = over_sample_util.binned_array_2d_from(
-        #     array_2d=jnp.array(array),
-        #     mask_2d=jnp.array(self.mask),
-        #     sub_size=jnp.array(self.sub_size).astype("int"),
-        # )
+        if self.sub_is_uniform:
+            binned_array_2d = array.reshape(
+                self.mask.shape_slim, self.sub_size[0] ** 2
+            ).mean(axis=1)
+        else:
+
+            # Define group sizes
+            group_sizes = jnp.array(self.sub_size.array.astype("int") ** 2)
+
+            # Compute the cumulative sum of group sizes to get split points
+            split_indices = jnp.cumsum(group_sizes)
+
+            # Ensure correct concatenation by making 0 a JAX array
+            start_indices = jnp.concatenate((jnp.array([0]), split_indices[:-1]))
 
-        binned_array_2d = array.reshape(
-            self.mask.shape_slim, self.sub_size[0] ** 2
-        ).mean(axis=1)
+            # Compute the group means
+            binned_array_2d = jnp.array(
+                [array[start:end].mean() for start, end in zip(start_indices, split_indices)])
 
         return Array2D(
             values=binned_array_2d,

test_autoarray/operators/over_sample/test_over_sampler.py

@@ -70,6 +70,16 @@ def test__binned_array_2d_from():
         pixel_scales=1.0,
     )
 
+    over_sampling = aa.OverSampler(
+        mask=mask, sub_size=aa.Array2D(values=[2, 2], mask=mask)
+    )
+
+    arr = np.array([1.0, 5.0, 7.0, 10.0, 10.0, 10.0, 10.0, 10.0])
+
+    binned_array_2d = over_sampling.binned_array_2d_from(array=arr)
+
+    assert binned_array_2d.slim == pytest.approx(np.array([5.75, 10.0]), 1.0e-4)
+
     over_sampling = aa.OverSampler(
         mask=mask, sub_size=aa.Array2D(values=[1, 2], mask=mask)
     )