over sampler xp simplified

Author: Jammy2211

autoarray/inversion/pixelization/mappers/mapper_util.py

@@ -516,7 +516,7 @@ def mapping_matrix_from(
     flat_pixidx = xp.where(flat_pixidx < 0, OUT, flat_pixidx)
 
     # 5) Multiply by sub_fraction of the slim row
-    flat_frac = sub_fraction[flat_parent]  # (M_sub*B,)
+    flat_frac = xp.take(sub_fraction, flat_parent, axis=0)  # (M_sub*B,)
     flat_contrib = flat_w * flat_frac  # (M_sub*B,)
 
     # 6) Scatter into (M × (S+1)), summing duplicates

autoarray/operators/over_sampling/decorator.py

@@ -31,6 +31,7 @@ def over_sample(func):
     def wrapper(
         obj: object,
         grid: Union[np.ndarray, Grid2D, Grid2DIrregular, Grid1D],
+        xp=np,
         *args,
         **kwargs,
     ) -> Union[np.ndarray, Array1D, Array2D, ArrayIrregular, List]:
@@ -49,12 +50,12 @@ def wrapper(
         """
 
         if isinstance(grid, Grid2DIrregular) or isinstance(grid, Grid1D):
-            return func(obj=obj, grid=grid, *args, **kwargs)
+            return func(obj=obj, grid=grid, xp=xp, *args, **kwargs)
 
         if obj is not None:
-            values = func(obj, grid.over_sampled, *args, **kwargs)
+            values = func(obj, grid.over_sampled, xp, *args, **kwargs)
         else:
-            values = func(grid.over_sampled, *args, **kwargs)
+            values = func(grid.over_sampled, xp, *args, **kwargs)
 
         return grid.over_sampler.binned_array_2d_from(array=values)
 

autoarray/operators/over_sampling/over_sampler.py

@@ -14,7 +14,7 @@
 
 @register_pytree_node_class
 class OverSampler:
-    def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D], xp=np):
+    def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]):
         """
          Over samples grid calculations using a uniform sub-grid.
 
@@ -149,7 +149,7 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D], xp=np):
         self.sub_total = int(np.sum(self.sub_size**2))
         self.sub_length = self.sub_size**self.mask.dimensions
         self.sub_fraction = Array2D(
-            values=xp.array(1.0 / self.sub_length.array), mask=self.mask
+            values=1.0 / self.sub_length.array, mask=self.mask
         )
 
         # Used for JAX based adaptive over sampling.
@@ -171,10 +171,6 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D], xp=np):
         ):
             self.segment_ids[start:end] = seg_id
 
-        self.segment_ids = xp.array(self.segment_ids)
-
-        self.xp = xp
-
     @property
     def sub_is_uniform(self) -> bool:
         """
@@ -260,12 +256,10 @@ def binned_array_2d_from(self, array: Array2D) -> "Array2D":
                 array, self.segment_ids, self.mask.pixels_in_mask
             )
             counts = jax.ops.segment_sum(
-                self.xp.ones_like(array), self.segment_ids, self.mask.pixels_in_mask
+                np.ones_like(array), self.segment_ids, self.mask.pixels_in_mask
             )
             binned_array_2d = sums / counts
 
-            binned_array_2d = self.xp.array(binned_array_2d)
-
         return Array2D(
             values=binned_array_2d,
             mask=self.mask,