feat: refactor

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