Coverage for core / src / sensorkit / astro / trajectory.py: 100%

1# SPDX-License-Identifier: Apache-2.0 

2import asyncio 

3from abc import ABC, abstractmethod 

4from collections.abc import Sequence 

5from datetime import datetime, timedelta 

6from typing import Self 

7 

8import numpy as np 

9import satkit 

10from astropy.coordinates import GCRS 

11from astropy.time import Time 

12 

13from sensorkit.astro.common import TLE 

14from sensorkit.astro.coords import Cartesian, StateVector, astropy_unit 

15 

16 

17def _gcrs_cartesian(pos: np.ndarray, vel: np.ndarray, jds: np.ndarray | float) -> GCRS: 

18 """Build a vectorized GCRS coordinate from position/velocity arrays and Julian dates. 

19 

20 `pos` and `vel` are shaped `(..., 3)` in meters and meters per second; `jds` 

21 broadcasts against the leading axis. 

22 """ 

23 m = astropy_unit("m") 

24 mps = astropy_unit("m/s") 

25 return GCRS( 

26 x=pos[..., 0] * m, 

27 y=pos[..., 1] * m, 

28 z=pos[..., 2] * m, 

29 v_x=vel[..., 0] * mps, 

30 v_y=vel[..., 1] * mps, 

31 v_z=vel[..., 2] * mps, 

32 obstime=Time(jds, format="jd"), 

33 representation_type="cartesian", 

34 ) 

35 

36 

37class Trajectory(ABC): 

38 """Abstract orbital trajectory that can be propagated and sampled.""" 

39 

40 @abstractmethod 

41 async def propagate(self, when: datetime | timedelta) -> Self: 

42 """Propagate the trajectory to the given point in time.""" 

43 

44 @abstractmethod 

45 def sample(self, epochs: Sequence[datetime] | None = None) -> GCRS: 

46 """Interpolate the state at each epoch. 

47 

48 With no epochs, samples once at the current time and returns a scalar 

49 coordinate. A sequence yields one array-valued GCRS spanning the epochs. 

50 """ 

51 

52 

53class OrbitalTrajectory(Trajectory): 

54 """Trajectory derived from numerical orbital propagation of a state vector via satkit.""" 

55 

56 def __init__(self, sv: StateVector, *, _result: satkit.propresult | None = None): 

57 self.sv = sv 

58 self._result = _result 

59 

60 async def propagate(self, when: datetime | timedelta) -> Self: 

61 result = await asyncio.to_thread( 

62 satkit.propagate, 

63 self.sv.to_numpy(), 

64 satkit.time.from_datetime(self.sv.t), 

65 satkit.time.from_datetime(when), 

66 ) 

67 return OrbitalTrajectory( 

68 StateVector( 

69 result.time.as_datetime(), 

70 Cartesian(*result.pos), 

71 Cartesian(*result.vel), 

72 ), 

73 _result=result, 

74 ) 

75 

76 def sample(self, epochs: Sequence[datetime] | None = None) -> GCRS: 

77 times = ( 

78 [satkit.time.now()] 

79 if epochs is None 

80 else [satkit.time.from_datetime(e) for e in epochs] 

81 ) 

82 vecs = np.array([self._result.interp(t) for t in times]) 

83 jds = np.array([t.as_jd() for t in times]) 

84 gcrs = _gcrs_cartesian(vecs[:, :3], vecs[:, 3:], jds) 

85 return gcrs[0] if epochs is None else gcrs 

86 

87 

88class TLETrajectory(Trajectory): 

89 """Trajectory derived from SGP4 propagation of a Two-Line Element set.""" 

90 

91 def __init__(self, tle: TLE): 

92 self.tle = satkit.TLE.from_lines(tle.to_list()) 

93 

94 async def propagate(self, when: datetime | timedelta) -> Self: 

95 return self 

96 

97 def sample(self, epochs: Sequence[datetime] | None = None) -> GCRS: 

98 times = ( 

99 [satkit.time.now()] 

100 if epochs is None 

101 else [satkit.time.from_datetime(e) for e in epochs] 

102 ) 

103 pos, vel, jds = [], [], [] 

104 

105 for t in times: 

106 teme_p, teme_v = satkit.sgp4(self.tle, t) 

107 q = satkit.frametransform.qteme2gcrf(t) 

108 pos.append(q * teme_p) 

109 vel.append(q * teme_v) 

110 jds.append(t.as_jd()) 

111 

112 gcrs = _gcrs_cartesian(np.array(pos), np.array(vel), np.array(jds)) 

113 return gcrs[0] if epochs is None else gcrs