inned supeer sampling array uses jax.ops.segment_sum for fast compile

Author: Jammy2211

autoarray/inversion/pixelization/mappers/mapper_util.py

@@ -236,7 +236,6 @@ def adaptive_rectangular_mappings_weights_via_interpolation_from(
     # --- Step 3. Transform oversampled grid into index space ---
     grid_over_sampled_scaled = (source_plane_data_grid_over_sampled - mu) / scale
     grid_over_sampled_transformed = transform(grid_over_sampled_scaled)
-    grid_over_index = source_grid_size * grid_over_sampled_transformed
     grid_over_index = (source_grid_size - 3) * grid_over_sampled_transformed + 1
 
     # --- Step 4. Floor/ceil indices ---

autoarray/inversion/pixelization/mappers/rectangular.py

@@ -97,15 +97,6 @@ def pix_sub_weights(self) -> PixSubWeights:
         dimension of the array `pix_indexes_for_sub_slim_index` 1 and all entries in `pix_weights_for_sub_slim_index`
         are equal to 1.0.
         """
-
-        # mappings, weights = (
-        #     mapper_util.rectangular_mappings_weights_via_interpolation_from(
-        #         shape_native=self.shape_native,
-        #         source_plane_mesh_grid=self.source_plane_mesh_grid.array,
-        #         source_plane_data_grid=self.source_plane_data_grid.over_sampled,
-        #     )
-        # )
-
         mappings, weights = (
             mapper_util.adaptive_rectangular_mappings_weights_via_interpolation_from(
                 source_grid_size=self.shape_native[0],

autoarray/operators/over_sampling/over_sampler.py

@@ -1,5 +1,6 @@
 import numpy as np
 import jax.numpy as jnp
+import jax
 from typing import Union
 
 from autoconf import conf
@@ -164,6 +165,14 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]):
         # Ensure correct concatenation by making 0 a JAX array
         self.start_indices = np.concatenate((np.array([0]), self.split_indices[:-1]))
 
+        # Compute segment ids for each element in the flattened array
+        self.segment_ids = np.empty(np.sum(sub_size**2), dtype=np.int32)
+
+        for seg_id, (start, end) in enumerate(zip(self.start_indices, self.split_indices)):
+            self.segment_ids[start:end] = seg_id
+
+        self.segment_ids = jnp.array(self.segment_ids)
+
     @property
     def sub_is_uniform(self) -> bool:
         """
@@ -234,18 +243,18 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D":
             pass
 
         if self.sub_is_uniform:
+
             binned_array_2d = array.reshape(
                 self.mask.shape_slim, self.sub_size[0] ** 2
             ).mean(axis=1)
+
         else:
 
             # Compute the group means
-            binned_array_2d = jnp.array(
-                [
-                    array[start:end].mean()
-                    for start, end in zip(self.start_indices, self.split_indices)
-                ]
-            )
+
+            sums = jax.ops.segment_sum(array, self.segment_ids, self.mask.pixels_in_mask)
+            counts = jax.ops.segment_sum(jnp.ones_like(array), self.segment_ids, self.mask.pixels_in_mask)
+            binned_array_2d = sums / counts
 
         return Array2D(
             values=binned_array_2d,

autoarray/plot/mat_plot/two_d.py

@@ -564,7 +564,7 @@ def _plot_rectangular_mapper(
         else:
             ax = self.setup_subplot(aspect=aspect_inv)
 
-            shape_native = mapper.source_plane_mesh_grid.shape_native
+        shape_native = mapper.source_plane_mesh_grid.shape_native
 
         if pixel_values is not None: