Initial commit

yksa_orbital/backends/sgp4.py

@@ -0,0 +1,217 @@
+"""SGP4 propagation backend.
+
+Ported from the sibling ``yksa_tdas`` service's ``location/propagate.py`` (pure
+``sgp4`` + stdlib ``math``, no numpy). Builds a ``Satrec`` straight from the
+stored canonical OMM dict via :func:`sgp4.omm.initialize` -- no TLE-line
+round-trip -- propagates to the requested instant, and fills the TEME state plus
+Earth-fixed (ECEF) and WGS84 geodetic views on the :class:`StateVector`.
+
+The TEME->ECEF rotation uses GMST only (no polar motion / nutation); a future
+Orekit backend will provide rigorous frames. That approximation is fine for the
+sub-km display accuracy this service targets.
+"""
+
+from __future__ import annotations
+
+import math
+from datetime import datetime, timedelta, timezone
+
+from sgp4 import omm as sgp4_omm
+from sgp4.api import SGP4_ERRORS, Satrec, jday
+from sgp4.propagation import gstime
+
+from odm import PropagationError, PropagatorBackend, StateVector, parse_omm_epoch
+
+WGS84_A_KM = 6378.137
+WGS84_F = 1.0 / 298.257223563
+WGS84_E2 = WGS84_F * (2.0 - WGS84_F)
+EARTH_ROT_RAD_S = 7.292115e-5
+
+
+#: Frames the pure-Python backend can transform between (GMST rotation only).
+_SGP4_FRAMES = ("TEME", "ITRF")
+
+
+class Sgp4Backend(PropagatorBackend):
+    name = "sgp4"
+    # SGP4 natively yields TEME; the GMST rotation gives an Earth-fixed view.
+    display_frames = ("TEME", "ITRF")
+
+    def transform(self, states, frame_in, frame_out):
+        frame_in = frame_in.upper()
+        frame_out = frame_out.upper()
+        if frame_in == frame_out:
+            return list(states)
+        if frame_in not in _SGP4_FRAMES or frame_out not in _SGP4_FRAMES:
+            raise PropagationError(
+                f"sgp4 backend only transforms between {_SGP4_FRAMES}; "
+                f"got {frame_in!r}->{frame_out!r} (use the orekit backend for more)"
+            )
+        out = []
+        for sv in states:
+            if sv.epoch is None:
+                raise PropagationError("state needs an epoch to rotate frames")
+            jd_ut1 = _jd_of(sv.epoch)
+            if frame_in == "TEME":  # TEME -> ITRF
+                r, v = _teme_to_ecef(sv.r_km, sv.v_kms, jd_ut1)
+            else:  # ITRF -> TEME
+                r, v = _ecef_to_teme(sv.r_km, sv.v_kms, jd_ut1)
+            out.append(_replace_frame(sv, frame_out, r, v))
+        return out
+
+    def state_at(self, omm: dict, at: datetime | None = None) -> StateVector:
+        sat = _satrec_from_omm(omm)
+        element_epoch = _epoch_of(sat)
+        when = at if at is not None else element_epoch
+        if when is None:
+            raise PropagationError("OMM has no usable epoch and no target time given")
+        return _state(sat, when, element_epoch)
+
+    def ephemeris(
+        self, omm: dict, start: datetime, stop: datetime, step_s: float,
+    ) -> list[StateVector]:
+        if step_s <= 0:
+            raise PropagationError("step must be positive")
+        if stop < start:
+            raise PropagationError("stop must be on or after start")
+        sat = _satrec_from_omm(omm)
+        element_epoch = _epoch_of(sat)
+        out: list[StateVector] = []
+        t = start.astimezone(timezone.utc)
+        stop = stop.astimezone(timezone.utc)
+        step = timedelta(seconds=step_s)
+        # Guard against runaway loops; callers cap this via OEM_MAX_POINTS.
+        while t <= stop + timedelta(microseconds=1):
+            out.append(_state(sat, t, element_epoch))
+            t += step
+        return out
+
+
+def _satrec_from_omm(omm: dict) -> Satrec:
+    sat = Satrec()
+    sgp4_omm.initialize(sat, _normalise_for_sgp4(omm))
+    return sat
+
+
+def _normalise_for_sgp4(omm: dict) -> dict:
+    """Return a copy whose EPOCH matches sgp4.omm's strict ``...%S.%f`` parser.
+
+    Stored epochs occasionally lack fractional seconds (or carry a trailing
+    ``Z``); sgp4.omm.initialize only accepts ``%Y-%m-%dT%H:%M:%S.%f``.
+    """
+    fields = dict(omm)
+    epoch = fields.get("EPOCH")
+    dt = parse_omm_epoch(epoch)
+    if dt is not None:
+        fields["EPOCH"] = dt.astimezone(timezone.utc).replace(tzinfo=None).strftime(
+            "%Y-%m-%dT%H:%M:%S.%f"
+        )
+    return fields
+
+
+def _epoch_of(sat: Satrec) -> datetime | None:
+    try:
+        jd = sat.jdsatepoch + sat.jdsatepochF
+        unix = (jd - 2440587.5) * 86400.0
+        return datetime.fromtimestamp(unix, tz=timezone.utc)
+    except Exception:  # noqa: BLE001 -- defensive; a bad epoch just yields None
+        return None
+
+
+def _state(sat: Satrec, when: datetime, element_epoch: datetime | None) -> StateVector:
+    when = when.astimezone(timezone.utc)
+    jd, fr = jday(
+        when.year, when.month, when.day,
+        when.hour, when.minute, when.second + when.microsecond / 1e6,
+    )
+    err, r, v = sat.sgp4(jd, fr)
+    if err != 0:
+        raise PropagationError(SGP4_ERRORS.get(err, f"sgp4 error {err}"))
+
+    r_ecef, v_ecef = _teme_to_ecef(r, v, jd + fr)
+    geodetic = _ecef_to_geodetic(r_ecef)
+    return StateVector(
+        epoch=when,
+        frame="TEME",
+        r_km=(r[0], r[1], r[2]),
+        v_kms=(v[0], v[1], v[2]),
+        ecef_km=r_ecef,
+        ecef_v_kms=v_ecef,
+        geodetic=geodetic,
+        element_epoch=element_epoch,
+    )
+
+
+def _jd_of(when: datetime) -> float:
+    """Full Julian date (UT1≈UTC) for a UTC datetime."""
+    when = when.astimezone(timezone.utc)
+    jd, fr = jday(
+        when.year, when.month, when.day,
+        when.hour, when.minute, when.second + when.microsecond / 1e6,
+    )
+    return jd + fr
+
+
+def _replace_frame(sv: StateVector, frame: str, r, v) -> StateVector:
+    """Copy ``sv`` with a new frame label and position/velocity.
+
+    When the target is the Earth-fixed frame, mirror the rotated state into the
+    ECEF convenience fields so geodetic-dependent conversions keep working.
+    """
+    ecef_km = r if frame == "ITRF" else sv.ecef_km
+    ecef_v_kms = v if frame == "ITRF" else sv.ecef_v_kms
+    return StateVector(
+        epoch=sv.epoch,
+        frame=frame,
+        r_km=tuple(r),
+        v_kms=tuple(v),
+        ecef_km=ecef_km,
+        ecef_v_kms=ecef_v_kms,
+        geodetic=sv.geodetic,
+        element_epoch=sv.element_epoch,
+        warnings=list(sv.warnings),
+    )
+
+
+def _teme_to_ecef(r, v, jd_ut1):
+    """Rotate TEME position/velocity to Earth-fixed (ECEF) via GMST."""
+    theta = gstime(jd_ut1)
+    cos, sin = math.cos(theta), math.sin(theta)
+    x = cos * r[0] + sin * r[1]
+    y = -sin * r[0] + cos * r[1]
+    z = r[2]
+    vx = cos * v[0] + sin * v[1] + EARTH_ROT_RAD_S * y
+    vy = -sin * v[0] + cos * v[1] - EARTH_ROT_RAD_S * x
+    vz = v[2]
+    return (x, y, z), (vx, vy, vz)
+
+
+def _ecef_to_teme(r, v, jd_ut1):
+    """Inverse of :func:`_teme_to_ecef`: Earth-fixed (ECEF) -> TEME via GMST."""
+    theta = gstime(jd_ut1)
+    cos, sin = math.cos(theta), math.sin(theta)
+    x = cos * r[0] - sin * r[1]
+    y = sin * r[0] + cos * r[1]
+    z = r[2]
+    # Undo the Earth-rotation term, then the rotation, to recover TEME velocity.
+    a = v[0] - EARTH_ROT_RAD_S * r[1]
+    b = v[1] + EARTH_ROT_RAD_S * r[0]
+    vx = cos * a - sin * b
+    vy = sin * a + cos * b
+    vz = v[2]
+    return (x, y, z), (vx, vy, vz)
+
+
+def _ecef_to_geodetic(r) -> tuple[float, float, float]:
+    """ECEF (km) -> WGS84 geodetic latitude/longitude (deg) and altitude (km)."""
+    x, y, z = r
+    lon = math.atan2(y, x)
+    p = math.hypot(x, y)
+    lat = math.atan2(z, p * (1.0 - WGS84_E2))
+    alt = 0.0
+    for _ in range(8):
+        sin_lat = math.sin(lat)
+        n = WGS84_A_KM / math.sqrt(1.0 - WGS84_E2 * sin_lat * sin_lat)
+        alt = p / math.cos(lat) - n
+        lat = math.atan2(z, p * (1.0 - WGS84_E2 * n / (n + alt)))
+    return math.degrees(lat), math.degrees(lon), alt