tensor/internal/core/grid.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package core

import (
	"math"
	"math/bits"
)

// Multilinear interpolation over a regular grid of any rank: the
// rank-2 generalisation of Interpolate2D's bilinear read, one weight
// per axis, clamped at the boundary.

// InterpolateGrid reads the grid at every query point by multilinear
// interpolation. grid has one axis per dimension (each at least two
// samples); origins[i] and steps[i] place axis i, with steps strictly
// positive; queries is a rank-2 (m × rank) matrix, one row per query,
// axes in the grid's own order. Queries outside the grid clamp to the
// boundary, matching Interpolate's convention; a position that comes
// out NaN is an error, since there is nothing sensible to clamp it to.
// A rank above 12 is an error: 2^12 corners per query is where this
// implementation stays honest about its cost.
//
// The grid and the queries are widened once and the walk reads the
// payload slices, the same values the accessors returned; the axes
// combine bottom up, deepest first, which is the order the recursive
// walk evaluated, so the bits are the recursive walk's.
func InterpolateGrid(grid *Array, origins, steps []float64, queries *Array) (*Array, error) {
	const name = "InterpolateGrid"
	dims := grid.NDim()
	if dims < 1 {
		return nil, errf("%s: the grid must have at least one axis", name)
	}
	if dims > 12 {
		return nil, errf("%s: the grid has %d axes, above the 12 this multilinear read supports", name, dims)
	}
	if len(origins) != dims || len(steps) != dims {
		return nil, errf("%s: origins and steps need %d entries each, got %d and %d", name, dims, len(origins), len(steps))
	}
	shape := grid.Shape()
	for i := range dims {
		if !isFiniteStep(steps[i]) || steps[i] <= 0 {
			return nil, errf("%s: axis %d needs a positive finite step, got %g", name, i, steps[i])
		}
		if shape[i] < 2 {
			return nil, errf("%s: axis %d needs at least two samples, got %d", name, i, shape[i])
		}
	}
	if queries.NDim() != 2 || queries.Shape()[1] != dims {
		return nil, errf("%s: queries must be rank 2 with %d columns, got shape %s", name, dims, shapeText(queries.Shape()))
	}
	if grid.dt == Complex || queries.dt == Complex {
		return nil, errf("%s: complex grids are not supported", name)
	}
	m := queries.Shape()[0]
	out := &Array{shape: []int{m}, dt: Float}
	out.alloc(m)
	gv, qv := floatPayload(grid), floatPayload(queries)
	// The axis strides and the corner flats are functions of the shape
	// alone: precomputed once, the corner table stepped by the lowest
	// cleared bit so no query re-walks the axes to find its corners.
	strides := make([]int, dims)
	strides[dims-1] = 1
	for i := dims - 2; i >= 0; i-- {
		strides[i] = strides[i+1] * shape[i+1]
	}
	corners := 1 << dims
	flat := make([]int, corners)
	for c := 1; c < corners; c++ {
		flat[c] = flat[c&(c-1)] + strides[bits.TrailingZeros32(uint32(c))]
	}
	buf := make([]float64, corners)
	idx := make([]int, dims)
	frac := make([]float64, dims)
	for q := range m {
		base := 0
		for i := range dims {
			t := (qv[q*dims+i] - origins[i]) / steps[i]
			// NaN compares false against both clamps below and converts
			// to the platform's indefinite integer, which then indexes
			// far outside the grid: the clamp contract only holds for
			// the infinities, so an undefined position is a loud error.
			if math.IsNaN(t) {
				return nil, errf("%s: query %d axis %d is NaN, which cannot be clamped", name, q, i)
			}
			if t < 0 {
				t = 0
			}
			if t > float64(shape[i]-1) {
				t = float64(shape[i] - 1)
			}
			idx[i] = min(int(t), shape[i]-2)
			frac[i] = t - float64(idx[i])
			base += idx[i] * strides[i]
		}
		// The corners seed the deepest level and the axes combine from
		// the last one up: buf[l] takes the lo corner first, the hi
		// corner second, the operand order the recursive walk kept.
		for c := range corners {
			buf[c] = gv[base+flat[c]]
		}
		for i := dims - 1; i >= 0; i-- {
			w := frac[i]
			lo := 1 - w
			for l := range 1 << i {
				buf[l] = lo*buf[l] + w*buf[l+1<<i]
			}
		}
		out.floats[q] = buf[0]
	}
	return out, nil
}

// isFiniteStep guards the step validation against NaN and infinity.
func isFiniteStep(s float64) bool {
	return !math.IsNaN(s) && !math.IsInf(s, 0)
}