[HOTFIX!] Previous chunking fix was breaking the interpolation

Author: Leguark

gempy_engine/modules/evaluator/generic_evaluator.py

@@ -8,63 +8,92 @@
 from gempy_engine.modules.kernel_constructor.kernel_constructor_interface import yield_evaluation_grad_kernel, yield_evaluation_kernel
 
 
-def generic_evaluator(solver_input: SolverInput, weights: np.ndarray, options: InterpolationOptions) -> ExportedFields:
+def generic_evaluator(
+    solver_input: SolverInput,
+    weights: np.ndarray,
+    options: InterpolationOptions
+) -> ExportedFields:
     grid_size = solver_input.xyz_to_interpolate.shape[0]
-    matrix_size = grid_size * weights.shape[0]
-    scalar_field: np.ndarray = BackendTensor.t.zeros(grid_size, dtype=weights.dtype)
+    max_op_size = options.evaluation_chunk_size
+    num_weights = weights.shape[0]
+    
+
+    chunk_size_grid = max(1, int(max_op_size / num_weights))  # Ensure at least 1 point per chunk
+    n_chunks = int(np.ceil(grid_size / chunk_size_grid)) 
+
+    # Pre‑allocate outputs
+    scalar_field = BackendTensor.t.zeros(grid_size, dtype=weights.dtype)
     gx_field: Optional[np.ndarray] = None
     gy_field: Optional[np.ndarray] = None
     gz_field: Optional[np.ndarray] = None
-    gradient = options.compute_scalar_gradient
-
-    # * Chunking the evaluation
-    max_size = options.evaluation_chunk_size
-    n_chunks = int(np.ceil(matrix_size / max_size))
-    chunk_size = int(np.ceil(grid_size / n_chunks))
-    for i in range(n_chunks): # TODO: It seems the chunking is not properly implemented
-        slice_array = slice(i * chunk_size, (i + 1) * chunk_size)
-        scalar_field_chunk, gx_field_chunk, gy_field_chunk, gz_field_chunk = _eval_on(
+    if options.compute_scalar_gradient:
+        gx_field = BackendTensor.t.zeros(grid_size, dtype=weights.dtype)
+        gy_field = BackendTensor.t.zeros(grid_size, dtype=weights.dtype)
+        if options.number_dimensions == 3:
+            gz_field = BackendTensor.t.zeros(grid_size, dtype=weights.dtype)
+
+    # Chunked evaluation over grid indices
+    for i in range(n_chunks):
+
+        start = i * chunk_size_grid
+        end = min(grid_size, start + chunk_size_grid)  # Ensure 'end' doesn't exceed grid_size
+        slice_array = slice(start, end)
+
+        # Avoid processing empty slices if start == end
+        if start >= end:
+            continue
+
+        sf_chunk, gx_chunk, gy_chunk, gz_chunk = _eval_on(
             solver_input=solver_input,
             weights=weights,
             options=options,
             slice_array=slice_array
         )
 
-        scalar_field[slice_array] = scalar_field_chunk
-        if gradient is True:
-            if i == 0:
-                gx_field = BackendTensor.t.zeros(grid_size, dtype=weights.dtype)
-                gy_field = BackendTensor.t.zeros(grid_size, dtype=weights.dtype)
-                gz_field = BackendTensor.t.zeros(grid_size, dtype=weights.dtype)
-
-            gx_field[slice_array] = gx_field_chunk
-            gy_field[slice_array] = gy_field_chunk
-            gz_field[slice_array] = gz_field_chunk
+        scalar_field[slice_array] = sf_chunk
+        if options.compute_scalar_gradient:
+            gx_field[slice_array] = gx_chunk  # type: ignore
+            gy_field[slice_array] = gy_chunk  # type: ignore
+            if gz_field is not None:
+                gz_field[slice_array] = gz_chunk  # type: ignore
 
     if n_chunks > 5:
         print(f"Chunking done: {n_chunks} chunks")
 
     return ExportedFields(scalar_field, gx_field, gy_field, gz_field)
 
 
-def _eval_on(solver_input, weights, options, slice_array: slice = None):
-    eval_kernel = yield_evaluation_kernel(solver_input, options.kernel_options, slice_array=slice_array)
-    scalar_field: np.ndarray = (eval_kernel.T @ weights).reshape(-1)
-    scalar_field[-50:]
+def _eval_on(
+    solver_input: SolverInput,
+    weights: np.ndarray,
+    options: InterpolationOptions,
+    slice_array: slice
+):
+    eval_kernel = yield_evaluation_kernel(
+        solver_input, options.kernel_options, slice_array=slice_array
+    )
+    scalar_field = (eval_kernel.T @ weights).reshape(-1)
+
     gx_field: Optional[np.ndarray] = None
     gy_field: Optional[np.ndarray] = None
     gz_field: Optional[np.ndarray] = None
-    if options.compute_scalar_gradient is True:
-        eval_gx_kernel = yield_evaluation_grad_kernel(solver_input, options.kernel_options, axis=0, slice_array=slice_array)
-        eval_gy_kernel = yield_evaluation_grad_kernel(solver_input, options.kernel_options, axis=1, slice_array=slice_array)
-        gx_field = (eval_gx_kernel.T @ weights).reshape(-1)
-        gy_field = (eval_gy_kernel.T @ weights).reshape(-1)
+
+    if options.compute_scalar_gradient:
+        eval_gx = yield_evaluation_grad_kernel(
+            solver_input, options.kernel_options, axis=0, slice_array=slice_array
+        )
+        eval_gy = yield_evaluation_grad_kernel(
+            solver_input, options.kernel_options, axis=1, slice_array=slice_array
+        )
+        gx_field = (eval_gx.T @ weights).reshape(-1)
+        gy_field = (eval_gy.T @ weights).reshape(-1)
 
         if options.number_dimensions == 3:
-            eval_gz_kernel = yield_evaluation_grad_kernel(solver_input, options.kernel_options, axis=2, slice_array=slice_array)
-            gz_field = (eval_gz_kernel.T @ weights).reshape(-1)
-        elif options.number_dimensions == 2:
-            gz_field = None
-        else:
-            raise ValueError("Number of dimensions have to be 2 or 3")
-    return scalar_field, gx_field, gy_field, gz_field
+            eval_gz = yield_evaluation_grad_kernel(
+                solver_input, options.kernel_options, axis=2, slice_array=slice_array
+            )
+            gz_field = (eval_gz.T @ weights).reshape(-1)
+        elif options.number_dimensions != 2:
+            raise ValueError("`number_dimensions` must be 2 or 3")
+
+    return scalar_field, gx_field, gy_field, gz_field