feat: refactor

internal/prediction/interpolate.go

@@ -0,0 +1,153 @@
+package prediction
+
+import (
+	"fmt"
+
+	"predictor-refactored/internal/dataset"
+)
+
+// Exact port of the reference interpolation logic (interpolate.pyx).
+// 4D interpolation: time, latitude, longitude, altitude (via geopotential height).
+
+// lerp1 holds an index and interpolation weight for one axis.
+type lerp1 struct {
+	index int
+	lerp  float64
+}
+
+// lerp3 holds indices and a combined weight for the (hour, lat, lon) axes.
+type lerp3 struct {
+	hour, lat, lng int
+	lerp           float64
+}
+
+// RangeError indicates a coordinate is outside the dataset bounds.
+type RangeError struct {
+	Variable string
+	Value    float64
+}
+
+func (e *RangeError) Error() string {
+	return fmt.Sprintf("%s=%f out of range", e.Variable, e.Value)
+}
+
+// pick computes interpolation indices and weights for a single axis.
+// left: axis start, step: axis spacing, n: number of points, value: query value.
+// Returns two lerp1 values (lower and upper bracket).
+func pick(left, step float64, n int, value float64, variableName string) ([2]lerp1, error) {
+	a := (value - left) / step
+	b := int(a) // truncation toward zero, same as Cython <long> cast
+	if b < 0 || b >= n-1 {
+		return [2]lerp1{}, &RangeError{Variable: variableName, Value: value}
+	}
+	l := a - float64(b)
+	return [2]lerp1{
+		{index: b, lerp: 1 - l},
+		{index: b + 1, lerp: l},
+	}, nil
+}
+
+// pick3 computes 8 trilinear interpolation weights for (hour, lat, lng).
+func pick3(hour, lat, lng float64) ([8]lerp3, error) {
+	lhour, err := pick(0, 3, 65, hour, "hour")
+	if err != nil {
+		return [8]lerp3{}, err
+	}
+	llat, err := pick(-90, 0.5, 361, lat, "lat")
+	if err != nil {
+		return [8]lerp3{}, err
+	}
+	// Longitude wraps: tell pick the axis is one larger, then wrap index 720 → 0
+	llng, err := pick(0, 0.5, 720+1, lng, "lng")
+	if err != nil {
+		return [8]lerp3{}, err
+	}
+	if llng[1].index == 720 {
+		llng[1].index = 0
+	}
+
+	var out [8]lerp3
+	i := 0
+	for _, a := range lhour {
+		for _, b := range llat {
+			for _, c := range llng {
+				out[i] = lerp3{
+					hour: a.index,
+					lat:  b.index,
+					lng:  c.index,
+					lerp: a.lerp * b.lerp * c.lerp,
+				}
+				i++
+			}
+		}
+	}
+	return out, nil
+}
+
+// interp3 performs 8-point weighted interpolation at a given variable and pressure level.
+func interp3(ds *dataset.File, lerps [8]lerp3, variable, level int) float64 {
+	var r float64
+	for i := 0; i < 8; i++ {
+		v := ds.Val(lerps[i].hour, level, variable, lerps[i].lat, lerps[i].lng)
+		r += float64(v) * lerps[i].lerp
+	}
+	return r
+}
+
+// search finds the largest pressure level index where interpolated geopotential
+// height is less than the target altitude. Searches levels 0..45 (excludes topmost).
+func search(ds *dataset.File, lerps [8]lerp3, target float64) int {
+	lower, upper := 0, 45
+
+	for lower < upper {
+		mid := (lower + upper + 1) / 2
+		test := interp3(ds, lerps, dataset.VarHeight, mid)
+		if target <= test {
+			upper = mid - 1
+		} else {
+			lower = mid
+		}
+	}
+
+	return lower
+}
+
+// interp4 performs altitude-interpolated wind lookup using two bracketing levels.
+func interp4(ds *dataset.File, lerps [8]lerp3, altLerp lerp1, variable int) float64 {
+	lower := interp3(ds, lerps, variable, altLerp.index)
+	upper := interp3(ds, lerps, variable, altLerp.index+1)
+	return lower*altLerp.lerp + upper*(1-altLerp.lerp)
+}
+
+// GetWind returns interpolated (u, v) wind components for the given position.
+// hour: fractional hours since dataset start.
+// lat: latitude in degrees (-90 to +90).
+// lng: longitude in degrees (0 to 360).
+// alt: altitude in metres above sea level.
+func GetWind(ds *dataset.File, warnings *Warnings, hour, lat, lng, alt float64) (u, v float64, err error) {
+	lerps, err := pick3(hour, lat, lng)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	altidx := search(ds, lerps, alt)
+	lower := interp3(ds, lerps, dataset.VarHeight, altidx)
+	upper := interp3(ds, lerps, dataset.VarHeight, altidx+1)
+
+	var altLerp float64
+	if lower != upper {
+		altLerp = (upper - alt) / (upper - lower)
+	} else {
+		altLerp = 0.5
+	}
+
+	if altLerp < 0 {
+		warnings.AltitudeTooHigh.Add(1)
+	}
+
+	alt1 := lerp1{index: altidx, lerp: altLerp}
+	u = interp4(ds, lerps, alt1, dataset.VarWindU)
+	v = interp4(ds, lerps, alt1, dataset.VarWindV)
+
+	return u, v, nil
+}

internal/prediction/models.go

@@ -0,0 +1,188 @@
+package prediction
+
+import (
+	"math"
+	"time"
+
+	"predictor-refactored/internal/dataset"
+	"predictor-refactored/internal/elevation"
+)
+
+// Exact port of the reference flight models (models.py).
+
+const (
+	pi180 = math.Pi / 180.0
+	_180pi = 180.0 / math.Pi
+)
+
+// --- Up/Down Models ---
+
+// ConstantAscent returns a model with constant vertical velocity (m/s).
+func ConstantAscent(ascentRate float64) Model {
+	return func(t, lat, lng, alt float64) (dlat, dlng, dalt float64) {
+		return 0, 0, ascentRate
+	}
+}
+
+// DragDescent returns a descent-under-parachute model.
+// seaLevelDescentRate is the descent rate at sea level (m/s, positive value).
+// Uses the NASA atmosphere model for density at altitude.
+func DragDescent(seaLevelDescentRate float64) Model {
+	dragCoefficient := seaLevelDescentRate * 1.1045
+
+	return func(t, lat, lng, alt float64) (dlat, dlng, dalt float64) {
+		return 0, 0, -dragCoefficient / math.Sqrt(nasaDensity(alt))
+	}
+}
+
+// nasaDensity computes air density using the NASA atmosphere model.
+// Reference: http://www.grc.nasa.gov/WWW/K-12/airplane/atmosmet.html
+func nasaDensity(alt float64) float64 {
+	var temp, pressure float64
+
+	switch {
+	case alt > 25000:
+		temp = -131.21 + 0.00299*alt
+		pressure = 2.488 * math.Pow((temp+273.1)/216.6, -11.388)
+	case alt > 11000:
+		temp = -56.46
+		pressure = 22.65 * math.Exp(1.73-0.000157*alt)
+	default:
+		temp = 15.04 - 0.00649*alt
+		pressure = 101.29 * math.Pow((temp+273.1)/288.08, 5.256)
+	}
+
+	return pressure / (0.2869 * (temp + 273.1))
+}
+
+// --- Sideways Models ---
+
+// WindVelocity returns a model that gives lateral movement at the wind velocity.
+// ds is the wind dataset, dsEpoch is the dataset start time as UNIX timestamp.
+func WindVelocity(ds *dataset.File, dsEpoch float64, warnings *Warnings) Model {
+	return func(t, lat, lng, alt float64) (dlat, dlng, dalt float64) {
+		tHours := (t - dsEpoch) / 3600.0
+		u, v, err := GetWind(ds, warnings, tHours, lat, lng, alt)
+		if err != nil {
+			return 0, 0, 0
+		}
+
+		R := 6371009.0 + alt
+		dlat = _180pi * v / R
+		dlng = _180pi * u / (R * math.Cos(lat*pi180))
+		return dlat, dlng, 0
+	}
+}
+
+// --- Model Combinations ---
+
+// LinearModel returns a model that sums all component models.
+func LinearModel(models ...Model) Model {
+	return func(t, lat, lng, alt float64) (dlat, dlng, dalt float64) {
+		for _, m := range models {
+			d1, d2, d3 := m(t, lat, lng, alt)
+			dlat += d1
+			dlng += d2
+			dalt += d3
+		}
+		return
+	}
+}
+
+// --- Termination Criteria ---
+
+// BurstTermination returns a terminator that fires when altitude >= burstAltitude.
+func BurstTermination(burstAltitude float64) Terminator {
+	return func(t, lat, lng, alt float64) bool {
+		return alt >= burstAltitude
+	}
+}
+
+// SeaLevelTermination fires when altitude <= 0.
+func SeaLevelTermination(t, lat, lng, alt float64) bool {
+	return alt <= 0
+}
+
+// TimeTermination returns a terminator that fires when t > maxTime.
+func TimeTermination(maxTime float64) Terminator {
+	return func(t, lat, lng, alt float64) bool {
+		return t > maxTime
+	}
+}
+
+// ElevationTermination returns a terminator that fires when alt < ground level.
+// Uses ruaumoko-compatible elevation data. Longitude is normalised internally.
+func ElevationTermination(elev *elevation.Dataset) Terminator {
+	return func(t, lat, lng, alt float64) bool {
+		return elev.Get(lat, lng) > alt
+	}
+}
+
+// --- Pre-Defined Profiles ---
+
+// Stage pairs a model with its termination criterion.
+type Stage struct {
+	Model      Model
+	Terminator Terminator
+}
+
+// StandardProfile creates the chain for a standard high-altitude balloon flight:
+// ascent at constant rate → burst → descent under parachute.
+// If elev is non-nil, descent terminates at ground level; otherwise at sea level.
+func StandardProfile(ascentRate, burstAltitude, descentRate float64,
+	ds *dataset.File, dsEpoch float64, warnings *Warnings,
+	elev *elevation.Dataset) []Stage {
+
+	wind := WindVelocity(ds, dsEpoch, warnings)
+
+	modelUp := LinearModel(ConstantAscent(ascentRate), wind)
+	termUp := BurstTermination(burstAltitude)
+
+	modelDown := LinearModel(DragDescent(descentRate), wind)
+	var termDown Terminator
+	if elev != nil {
+		termDown = ElevationTermination(elev)
+	} else {
+		termDown = Terminator(SeaLevelTermination)
+	}
+
+	return []Stage{
+		{Model: modelUp, Terminator: termUp},
+		{Model: modelDown, Terminator: termDown},
+	}
+}
+
+// FloatProfile creates the chain for a floating balloon flight:
+// ascent to float altitude → float until stop time.
+func FloatProfile(ascentRate, floatAltitude float64, stopTime time.Time,
+	ds *dataset.File, dsEpoch float64, warnings *Warnings) []Stage {
+
+	wind := WindVelocity(ds, dsEpoch, warnings)
+
+	modelUp := LinearModel(ConstantAscent(ascentRate), wind)
+	termUp := BurstTermination(floatAltitude)
+
+	modelFloat := wind
+	termFloat := TimeTermination(float64(stopTime.Unix()))
+
+	return []Stage{
+		{Model: modelUp, Terminator: termUp},
+		{Model: modelFloat, Terminator: termFloat},
+	}
+}
+
+// RunPrediction runs a prediction with the given profile stages.
+// launchTime is a UNIX timestamp.
+func RunPrediction(launchTime float64, lat, lng, alt float64, stages []Stage) []StageResult {
+	chain := make([]struct {
+		Model      Model
+		Terminator Terminator
+	}, len(stages))
+
+	for i, s := range stages {
+		chain[i].Model = s.Model
+		chain[i].Terminator = s.Terminator
+	}
+
+	return Solve(launchTime, lat, lng, alt, chain)
+}

internal/prediction/solver.go

@@ -0,0 +1,180 @@
+package prediction
+
+import "math"
+
+// Exact port of the reference RK4 solver (solver.pyx).
+// Integrates balloon state using RK4 with dt=60 seconds.
+// Termination uses binary search refinement (tolerance 0.01).
+
+// Vec holds the balloon state: latitude, longitude, altitude.
+type Vec struct {
+	Lat float64
+	Lng float64
+	Alt float64
+}
+
+// Model is a function that returns (dlat/dt, dlng/dt, dalt/dt) given state.
+// t is UNIX timestamp, lat/lng in degrees, alt in metres.
+type Model func(t float64, lat, lng, alt float64) (dlat, dlng, dalt float64)
+
+// Terminator returns true when integration should stop.
+type Terminator func(t float64, lat, lng, alt float64) bool
+
+// StageResult holds the trajectory points for one flight stage.
+type StageResult struct {
+	Points []TrajectoryPoint
+}
+
+// TrajectoryPoint is a single point in a trajectory (used by solver).
+type TrajectoryPoint struct {
+	T   float64 // UNIX timestamp
+	Lat float64
+	Lng float64
+	Alt float64
+}
+
+// pymod returns a % b with Python semantics (always non-negative when b > 0).
+func pymod(a, b float64) float64 {
+	r := math.Mod(a, b)
+	if r < 0 {
+		r += b
+	}
+	return r
+}
+
+// vecadd returns a + k*b, with lng wrapped to [0, 360).
+func vecadd(a Vec, k float64, b Vec) Vec {
+	return Vec{
+		Lat: a.Lat + k*b.Lat,
+		Lng: pymod(a.Lng+k*b.Lng, 360.0),
+		Alt: a.Alt + k*b.Alt,
+	}
+}
+
+// scalarLerp returns (1-l)*a + l*b.
+func scalarLerp(a, b, l float64) float64 {
+	return (1-l)*a + l*b
+}
+
+// lngLerp interpolates longitude handling the 0/360 wrap-around.
+func lngLerp(a, b, l float64) float64 {
+	l2 := 1 - l
+
+	if a > b {
+		a, b = b, a
+		l, l2 = l2, l
+	}
+
+	// distance round one way:  b - a
+	// distance around other:   (a + 360) - b
+	if b-a < 180.0 {
+		return l2*a + l*b
+	}
+	return pymod(l2*(a+360)+l*b, 360.0)
+}
+
+// vecLerp returns (1-l)*a + l*b with proper longitude wrapping.
+func vecLerp(a, b Vec, l float64) Vec {
+	return Vec{
+		Lat: scalarLerp(a.Lat, b.Lat, l),
+		Lng: lngLerp(a.Lng, b.Lng, l),
+		Alt: scalarLerp(a.Alt, b.Alt, l),
+	}
+}
+
+// rk4 integrates from initial conditions using RK4.
+// dt=60.0 seconds, terminationTolerance=0.01.
+func rk4(t float64, lat, lng, alt float64, model Model, terminator Terminator) []TrajectoryPoint {
+	const dt = 60.0
+	const terminationTolerance = 0.01
+
+	y := Vec{Lat: lat, Lng: lng, Alt: alt}
+	result := []TrajectoryPoint{{T: t, Lat: y.Lat, Lng: y.Lng, Alt: y.Alt}}
+
+	for {
+		// Evaluate model at 4 points (standard RK4)
+		k1lat, k1lng, k1alt := model(t, y.Lat, y.Lng, y.Alt)
+		k1 := Vec{Lat: k1lat, Lng: k1lng, Alt: k1alt}
+
+		mid1 := vecadd(y, dt/2, k1)
+		k2lat, k2lng, k2alt := model(t+dt/2, mid1.Lat, mid1.Lng, mid1.Alt)
+		k2 := Vec{Lat: k2lat, Lng: k2lng, Alt: k2alt}
+
+		mid2 := vecadd(y, dt/2, k2)
+		k3lat, k3lng, k3alt := model(t+dt/2, mid2.Lat, mid2.Lng, mid2.Alt)
+		k3 := Vec{Lat: k3lat, Lng: k3lng, Alt: k3alt}
+
+		end := vecadd(y, dt, k3)
+		k4lat, k4lng, k4alt := model(t+dt, end.Lat, end.Lng, end.Alt)
+		k4 := Vec{Lat: k4lat, Lng: k4lng, Alt: k4alt}
+
+		// y2 = y + dt/6*k1 + dt/3*k2 + dt/3*k3 + dt/6*k4
+		y2 := y
+		y2 = vecadd(y2, dt/6, k1)
+		y2 = vecadd(y2, dt/3, k2)
+		y2 = vecadd(y2, dt/3, k3)
+		y2 = vecadd(y2, dt/6, k4)
+
+		t2 := t + dt
+
+		if terminator(t2, y2.Lat, y2.Lng, y2.Alt) {
+			// Binary search to refine the termination point.
+			// Find l in [0, 1] such that (t3, y3) = lerp((t, y), (t2, y2), l)
+			// is near where the terminator fires.
+			left := 0.0
+			right := 1.0
+
+			var t3 float64
+			var y3 Vec
+			t3 = t2
+			y3 = y2
+
+			for right-left > terminationTolerance {
+				mid := (left + right) / 2
+				t3 = scalarLerp(t, t2, mid)
+				y3 = vecLerp(y, y2, mid)
+
+				if terminator(t3, y3.Lat, y3.Lng, y3.Alt) {
+					right = mid
+				} else {
+					left = mid
+				}
+			}
+
+			result = append(result, TrajectoryPoint{T: t3, Lat: y3.Lat, Lng: y3.Lng, Alt: y3.Alt})
+			break
+		}
+
+		// Update current state
+		t = t2
+		y = y2
+		result = append(result, TrajectoryPoint{T: t, Lat: y.Lat, Lng: y.Lng, Alt: y.Alt})
+	}
+
+	return result
+}
+
+// Solve runs through a chain of (model, terminator) stages.
+// Returns one StageResult per stage.
+func Solve(t, lat, lng, alt float64, chain []struct {
+	Model      Model
+	Terminator Terminator
+}) []StageResult {
+	var results []StageResult
+
+	for _, stage := range chain {
+		points := rk4(t, lat, lng, alt, stage.Model, stage.Terminator)
+		results = append(results, StageResult{Points: points})
+
+		// Next stage starts where this one ended
+		if len(points) > 0 {
+			last := points[len(points)-1]
+			t = last.T
+			lat = last.Lat
+			lng = last.Lng
+			alt = last.Alt
+		}
+	}
+
+	return results
+}

internal/prediction/warnings.go

@@ -0,0 +1,21 @@
+package prediction
+
+import "sync/atomic"
+
+// Warnings tracks warning conditions during a prediction run.
+type Warnings struct {
+	AltitudeTooHigh atomic.Int64
+}
+
+// ToMap returns warnings as a map suitable for JSON serialization.
+// Only includes warnings that have fired.
+func (w *Warnings) ToMap() map[string]any {
+	result := make(map[string]any)
+	if n := w.AltitudeTooHigh.Load(); n > 0 {
+		result["altitude_too_high"] = map[string]any{
+			"count":       n,
+			"description": "The altitude went too high, above the max forecast wind. Wind data will be unreliable",
+		}
+	}
+	return result
+}