knn without batch fully working with factory fixes

Author: Jammy2211

autoarray/inversion/pixelization/mappers/factory.py

@@ -70,21 +70,22 @@ def mapper_from(
             preloads=preloads,
             xp=xp,
         )
-    elif isinstance(mapper_grids.source_plane_mesh_grid, Mesh2DDelaunay):
-        return MapperDelaunay(
+    elif isinstance(mapper_grids.source_plane_mesh_grid, Mesh2DDelaunayKNN):
+        return MapperKNNInterpolator(
             mapper_grids=mapper_grids,
             border_relocator=border_relocator,
             regularization=regularization,
             settings=settings,
             preloads=preloads,
             xp=xp,
         )
-    elif isinstance(mapper_grids.source_plane_mesh_grid, Mesh2DDelaunayKNN):
-        return MapperKNNInterpolator(
+
+    elif isinstance(mapper_grids.source_plane_mesh_grid, Mesh2DDelaunay):
+        return MapperDelaunay(
             mapper_grids=mapper_grids,
             border_relocator=border_relocator,
             regularization=regularization,
             settings=settings,
             preloads=preloads,
             xp=xp,
-        )
+        )

autoarray/inversion/pixelization/mappers/knn.py

@@ -50,13 +50,12 @@ def _pix_sub_weights_from_query_points(self, query_points) -> PixSubWeights:
         # ------------------------------------------------------------------
         # Convert outputs to xp backend *only if needed*
         # ------------------------------------------------------------------
-        if xp is jnp:
-            weights = weights_jax
-            mappings = indices_jax
-        else:
-            # xp is numpy
+        if xp is np:
             weights = np.asarray(weights_jax)
             mappings = np.asarray(indices_jax)
+        else:
+            weights = weights_jax
+            mappings = indices_jax
 
         # ------------------------------------------------------------------
         # Sizes: always k for kNN
@@ -90,11 +89,40 @@ def pix_sub_weights(self) -> PixSubWeights:
     @property
     def pix_sub_weights_split_points(self) -> PixSubWeights:
         """
-        kNN mappings + kernel weights computed at split points (for split regularization schemes).
+        kNN mappings + kernel weights computed at split points (for split regularization schemes),
+        with split-point step sizes derived from kNN local spacing (no Delaunay / simplices).
         """
-        # Your Delaunay mesh exposes split points via self.delaunay.split_points.
-        # For KNN mesh, you should expose the same property. If not, route appropriately:
-        #   split_points = self.mesh.split_points
-        split_points = self.delaunay.split_points  # keep consistent with existing API
+        from autoarray.structures.mesh.delaunay_2d import split_points_from
 
+        # TODO: wire these to your pixelization / regularization config rather than hard-code.
+        k_neighbors = 10
+        kernel = "wendland_c4"
+        radius_scale = 1.5
+        areas_factor = 0.5
+
+        xp = self._xp  # np or jnp
+
+        # Mesh points (N, 2)
+        points = xp.asarray(self.source_plane_mesh_grid.array, dtype=xp.float64)
+
+        # kNN distances of each point to its neighbors (include self, then drop it)
+        _, _, dist_self = get_interpolation_weights(
+            points=points,
+            query_points=points,
+            k_neighbors=int(k_neighbors) + 1,
+            kernel=kernel,
+            radius_scale=float(radius_scale),
+        )
+
+        # Local spacing scale: distance to k-th nearest OTHER point
+        r_k = dist_self[:, 1:][:, -1]  # (N,)
+
+        # Split cross step size (length): sqrt(area) ~ r_k
+        split_step = xp.asarray(areas_factor, dtype=xp.float64) * r_k  # (N,)
+
+        # Split points (xp-native)
+        split_points = split_points_from(points=points, area_weights=split_step, xp=xp)
+
+        # Compute kNN mappings/weights at split points
         return self._pix_sub_weights_from_query_points(query_points=split_points)
+

autoarray/inversion/pixelization/mesh/knn.py

@@ -3,10 +3,13 @@
 Uses Wendland compactly supported kernels with normalized weights (partition of unity).
 More robust and faster than MLS, better accuracy than simple IDW.
 """
+import numpy as np
 import jax
 import jax.numpy as jnp
 from functools import partial
 
+from autoarray.structures.mesh.knn_delaunay_2d import Mesh2DDelaunayKNN
+
 
 def get_interpolation_weights(points, query_points, k_neighbors=10, kernel='wendland_c4',
                               radius_scale=1.5):
@@ -268,3 +271,33 @@ def __init__(self,  k_neighbors=10, kernel='wendland_c4',
 
         super().__init__()
 
+    def mesh_grid_from(
+        self,
+        source_plane_data_grid=None,
+        source_plane_mesh_grid=None,
+        preloads=None,
+        xp=np,
+    ):
+        """
+        Return the Delaunay ``source_plane_mesh_grid`` as a ``Mesh2DDelaunay`` object, which provides additional
+        functionality for performing operations that exploit the geometry of a Delaunay mesh.
+
+        Parameters
+        ----------
+        source_plane_data_grid
+            A 2D grid of (y,x) coordinates associated with the unmasked 2D data after it has been transformed to the
+            ``source`` reference frame.
+        source_plane_mesh_grid
+            The centres of every Delaunay pixel in the ``source`` frame, which are initially derived by computing a sparse
+            set of (y,x) coordinates computed from the unmasked data in the image-plane and applying a transformation
+            to this.
+        settings
+            Settings controlling the pixelization for example if a border is used to relocate its exterior coordinates.
+        """
+
+        return Mesh2DDelaunayKNN(
+            values=source_plane_mesh_grid,
+            source_plane_data_grid_over_sampled=source_plane_data_grid,
+            preloads=preloads,
+            _xp=xp,
+        )

autoarray/preloads.py

@@ -27,6 +27,7 @@ def __init__(
         use_voronoi_areas: bool = True,
         areas_factor: float = 0.5,
         skip_areas: bool = False,
+        splitted_only : bool = False
     ):
         """
         Stores preloaded arrays and matrices used during pixelized linear inversions, improving both performance