predictor/internal/prediction/solver.go

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
}