feat: refactor

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
+}