Kernels¶

def interpolate_block(
    chunked_coords: np.ndarray,
    neighbours_coords: np.ndarray,
    idx_query: np.ndarray,
    chunk_idx: slice,
    src_values: np.ndarray,
    config: InterpolationConfig,
) -> tuple[np.ndarray, np.ndarray]:
    """Interpolate a block of points

    This function handles two different philosophies:
    1. Source-to-Destination (DEST_SRC_MAP[kernel] = True):
       - chunked_coords are source points, neighbours_coords are destination points
       - For each source point, distribute its value to neighboring destination points
       - Used by: DISTANCE_WEIGHTED, FEM

    2. Destination-to-Source (DEST_SRC_MAP[kernel] = False):
       - chunked_coords are destination points, neighbours_coords are source points
       - For each destination point, compute value from neighboring source points
       - Used by: AVERAGE, CLOSEST
    """
    is_src_to_dest = DEST_SRC_MAP[config.kernel]

    if is_src_to_dest:
        # Initialize output for destination points
        interpolated = np.zeros([neighbours_coords.shape[0], config.num_components])
    else:
        # Initialize output only for the current destination chunk
        block_size = max(chunk_idx.stop - chunk_idx.start, 0)
        interpolated = np.zeros([block_size, config.num_components])

    unmapped = np.zeros([1, config.num_components])

    i = chunk_idx.start
    for point in chunked_coords[chunk_idx]:
        neighbours_idx = np.asarray(idx_query[i], dtype=int)
        mapped = False

        # Ensure query is not empty
        if neighbours_idx.shape[0] > 0:
            distances = euclidean_distances(
                neighbours_coords[neighbours_idx], point.reshape(1, -1)
            )
            # clip the nodes based on max distance
            keep = distances.flatten() < config.max_distance
            distances = distances[keep]
            neighbours_idx = neighbours_idx[keep]

            # Ensure that after clipping something remains
            if neighbours_idx.shape[0] == 0:
                if is_src_to_dest:
                    # Source point has no destination neighbors - accumulate as unmapped
                    unmapped = unmapped + src_values[i, :]
                else:
                    # Destination point has no source neighbors - this is an error condition
                    if config.accept_no_neighbor:
                        dest_idx = i - chunk_idx.start
                        interpolated[dest_idx, :] = 0
                    else:
                        raise InterpolationError(
                            f"Destination point at {point} has no source neighbors within max_distance"
                        )
                i = i + 1
                continue

            # if coincident node found, assign directly
            if distances[0] < config.coincidence_tolerance:
                logging.debug(f"Coincident node found {point}")
                if is_src_to_dest:
                    # Source point coincides with destination point
                    # neighbours_idx[0] is the destination index
                    interpolated[neighbours_idx[0], :] = src_values[i, :]
                else:
                    # Destination point coincides with source point
                    # neighbours_idx[0] is the source index, i is relative to chunk
                    dest_idx = i - chunk_idx.start
                    interpolated[dest_idx, :] = src_values[neighbours_idx[0], :]
                i = i + 1
                continue

            # Perform the interpolation based on kernel and philosophy
            if is_src_to_dest:
                index = i
            else:
                index = i - chunk_idx.start

            # Source-to-Destination: distribute source value to destination neighbors
            if config.kernel == INTERPOLATION_KERNEL.DISTANCE_WEIGHTED:
                mapped = _dist_weight_kernel(
                    distances,
                    neighbours_idx,
                    i,
                    src_values,
                    interpolated,
                )
            elif config.kernel == INTERPOLATION_KERNEL.FEM:
                mapped = _FEM_interpolation_kernel(
                    distances,
                    neighbours_idx,
                    i,
                    neighbours_coords,
                    chunked_coords,
                    src_values,
                    interpolated,
                )

            # Destination-to-Source: compute destination value from source neighbors
            if config.kernel == INTERPOLATION_KERNEL.AVERAGE:
                mapped = _average_kernel(
                    neighbours_idx,
                    index,
                    src_values,
                    interpolated,
                )
            elif config.kernel == INTERPOLATION_KERNEL.AVERAGE_WEIGHTED:
                mapped = _average_kernel(
                    neighbours_idx,
                    index,
                    src_values,
                    interpolated,
                    distances=distances,
                )
            elif config.kernel == INTERPOLATION_KERNEL.CLOSEST:
                mapped = _closest_source_kernel(
                    distances,
                    neighbours_idx,
                    index,
                    src_values,
                    interpolated,
                )

            if not mapped:
                # Should be possible only in source-to-destination
                unmapped = unmapped + src_values[i, :]

        else:
            # No neighbors found at all
            if is_src_to_dest:
                unmapped = unmapped + src_values[i, :]
            else:
                if config.accept_no_neighbor:
                    dest_idx = i - chunk_idx.start
                    interpolated[dest_idx, :] = 0
                else:
                    raise InterpolationError(
                        f"Destination point at {point} has no neighbors"
                    )

        i += 1

    return interpolated, unmapped