predictor/internal/numerics/grid.go

129 lines
3.6 KiB
Go

package numerics
import "fmt"
// Axis describes a regularly-spaced grid axis with N grid points,
// values left, left+step, left+2*step, ..., left+(N-1)*step.
//
// If Wrap is true, the axis is periodic with period N*step (e.g. longitude).
// A query value at left+N*step wraps to the value at left+0*step. Locate
// returns Hi = 0 in that case.
type Axis struct {
Left float64
Step float64
N int
Wrap bool
Name string
}
// AxisError is returned by Axis.Locate when value lies outside a non-wrapping axis.
type AxisError struct {
Axis string
Value float64
}
func (e *AxisError) Error() string {
return fmt.Sprintf("%s=%v out of range", e.Axis, e.Value)
}
// Bracket holds the two surrounding grid indices and the fractional position
// of a value within an axis. The weight at Lo is (1 - Frac); the weight at Hi
// is Frac. Frac lies in [0, 1).
type Bracket struct {
Lo, Hi int
Frac float64
}
// Locate returns the bracket containing value within the axis.
// For a non-wrapping axis, value must lie in [Left, Left + (N-1)*Step);
// for a wrapping axis, value must lie in [Left, Left + N*Step).
func (a Axis) Locate(value float64) (Bracket, error) {
pos := (value - a.Left) / a.Step
lo := int(pos) // truncates toward zero; pos is non-negative for valid inputs
maxLo := a.N - 2
if a.Wrap {
maxLo = a.N - 1
}
if lo < 0 || lo > maxLo {
return Bracket{}, &AxisError{Axis: a.Name, Value: value}
}
hi := lo + 1
if a.Wrap && hi == a.N {
hi = 0
}
return Bracket{Lo: lo, Hi: hi, Frac: pos - float64(lo)}, nil
}
// TrilinearWeights returns the eight corner weights for a (axis0, axis1,
// axis2) bracket triple, in the canonical visiting order
//
// (0,0,0) (0,0,1) (0,1,0) (0,1,1) (1,0,0) (1,0,1) (1,1,0) (1,1,1)
//
// where the bit triple selects Lo (0) or Hi (1) on each axis. The weights sum
// to 1. Pair this with Dot8 over corner values fetched in the same order.
func TrilinearWeights(b3 [3]Bracket) [8]float64 {
wa0, wa1 := 1-b3[0].Frac, b3[0].Frac
wb0, wb1 := 1-b3[1].Frac, b3[1].Frac
wc0, wc1 := 1-b3[2].Frac, b3[2].Frac
wa0wb0 := wa0 * wb0
wa0wb1 := wa0 * wb1
wa1wb0 := wa1 * wb0
wa1wb1 := wa1 * wb1
return [8]float64{
wa0wb0 * wc0,
wa0wb0 * wc1,
wa0wb1 * wc0,
wa0wb1 * wc1,
wa1wb0 * wc0,
wa1wb0 * wc1,
wa1wb1 * wc0,
wa1wb1 * wc1,
}
}
// Dot8 returns the multiply-accumulate sum w[0]*v[0] + ... + w[7]*v[7].
//
// The fixed length and straight-line accumulation are written so the Go
// compiler can keep the values in registers and a future hand-vectorised
// port can replace the body with a single SIMD MAC. The accumulation order
// is fixed (ascending index) so results are reproducible.
func Dot8(w, v *[8]float64) float64 {
acc := w[0] * v[0]
acc = w[1]*v[1] + acc
acc = w[2]*v[2] + acc
acc = w[3]*v[3] + acc
acc = w[4]*v[4] + acc
acc = w[5]*v[5] + acc
acc = w[6]*v[6] + acc
acc = w[7]*v[7] + acc
return acc
}
// EvalTrilinear samples a 3D field via f at the eight corners defined by b3
// and returns the trilinearly interpolated value.
//
// Corners are visited in the canonical order documented on TrilinearWeights.
// With f(i,j,k) = a*i + b*j + c*k + d this returns a*pos0 + b*pos1 + c*pos2
// + d, modulo floating-point rounding. For the hot path prefer precomputing
// weights once via TrilinearWeights and reducing with Dot8.
func EvalTrilinear(b3 [3]Bracket, f func(i, j, k int) float64) float64 {
w := TrilinearWeights(b3)
a0, a1 := b3[0].Lo, b3[0].Hi
b0, b1 := b3[1].Lo, b3[1].Hi
c0, c1 := b3[2].Lo, b3[2].Hi
v := [8]float64{
f(a0, b0, c0),
f(a0, b0, c1),
f(a0, b1, c0),
f(a0, b1, c1),
f(a1, b0, c0),
f(a1, b0, c1),
f(a1, b1, c0),
f(a1, b1, c1),
}
return Dot8(&w, &v)
}