but of typing simplication

Author: Jammy2211

autoarray/operators/over_sampling/over_sample_util.py

@@ -44,9 +44,9 @@ def over_sample_size_convert_to_array_2d_from(
     if isinstance(over_sample_size, int):
         over_sample_size = np.full(
             fill_value=over_sample_size, shape=mask.pixels_in_mask
-        ).astype("int")
+        )
 
-    return Array2D(values=over_sample_size, mask=mask)
+    return Array2D(values=np.array(over_sample_size).astype("int"), mask=mask)
 
 
 @numba_util.jit()
@@ -444,23 +444,6 @@ def grid_2d_slim_over_sampled_via_mask_from(
 
                 for y1 in range(sub):
                     for x1 in range(sub):
-                        # if use_jax:
-                        #     # while this makes it run, it is very, very slow
-                        #     grid_slim = grid_slim.at[sub_index, 0].set(
-                        #         -(
-                        #             y_scaled
-                        #             - y_sub_half
-                        #             + y1 * y_sub_step
-                        #             + (y_sub_step / 2.0)
-                        #         )
-                        #     )
-                        #     grid_slim = grid_slim.at[sub_index, 1].set(
-                        #         x_scaled
-                        #         - x_sub_half
-                        #         + x1 * x_sub_step
-                        #         + (x_sub_step / 2.0)
-                        #     )
-                        # else:
                         grid_slim[sub_index, 0] = -(
                             y_scaled - y_sub_half + y1 * y_sub_step + (y_sub_step / 2.0)
                         )

autoarray/operators/over_sampling/over_sampler.py

@@ -154,7 +154,7 @@ 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])
+            np.isclose(self.sub_size, self.sub_size[0])
         )
 
     def tree_flatten(self):
@@ -233,12 +233,12 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D":
 
         if self.sub_is_uniform:
             binned_array_2d = array.reshape(
-                self.mask.shape_slim, self.sub_size.array[0] ** 2
+                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)
+            group_sizes = jnp.array(self.sub_size.array ** 2)
 
             # Compute the cumulative sum of group sizes to get split points
             split_indices = jnp.cumsum(group_sizes)
@@ -310,7 +310,7 @@ def sub_mask_native_for_sub_mask_slim(self) -> np.ndarray:
             print(derive_indexes_2d.sub_mask_native_for_sub_mask_slim)
         """
         return over_sample_util.native_sub_index_for_slim_sub_index_2d_from(
-            mask_2d=self.mask.array, sub_size=np.array(self.sub_size)
+            mask_2d=self.mask.array, sub_size=self.sub_size.array
         ).astype("int")
 
     @cached_property
@@ -363,7 +363,7 @@ def slim_for_sub_slim(self) -> np.ndarray:
             print(derive_indexes_2d.slim_for_sub_slim)
         """
         return over_sample_util.slim_index_for_sub_slim_index_via_mask_2d_from(
-            mask_2d=np.array(self.mask), sub_size=np.array(self.sub_size)
+            mask_2d=np.array(self.mask), sub_size=self.sub_size.array
         ).astype("int")
 
     @property
@@ -386,7 +386,7 @@ def uniform_over_sampled(self):
         grid = over_sample_util.grid_2d_slim_over_sampled_via_mask_from(
             mask_2d=np.array(self.mask),
             pixel_scales=self.mask.pixel_scales,
-            sub_size=np.array(self.sub_size).astype("int"),
+            sub_size=self.sub_size.array.astype("int"),
             origin=self.mask.origin,
         )
 

autoarray/structures/grids/uniform_2d.py

@@ -185,7 +185,7 @@ def __init__(
                 over_sample_util.grid_2d_slim_over_sampled_via_mask_from(
                     mask_2d=np.array(self.mask),
                     pixel_scales=self.mask.pixel_scales,
-                    sub_size=np.array(self.over_sampler.sub_size._array).astype("int"),
+                    sub_size=self.over_sampler.sub_size.array.astype("int"),
                     origin=self.mask.origin,
                 )
             )