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

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import functools 

5from dataclasses import dataclass 

6from datetime import UTC, datetime 

7from typing import TYPE_CHECKING 

8 

9import numpy as np 

10 

11from sensorkit.astro.common import ReferenceFrame 

12 

13if TYPE_CHECKING: 

14 from astropy.coordinates import SkyCoord 

15 from astropy.units import Unit 

16 

17 

18@functools.lru_cache 

19def astropy_unit(spec: str | Unit) -> Unit: 

20 """Resolve a unit spec to an astropy Unit, memoizing the parse. 

21 

22 Passing an already-resolved Unit is a cheap identity check; string specs are 

23 parsed once and cached, which matters for compound units like "m/s" whose 

24 grammar parse is otherwise paid on every call. 

25 """ 

26 from astropy.units import Unit 

27 

28 return Unit(spec) 

29 

30 

31type Coordinates = Horizontal | Equatorial | Geodetic | Cartesian 

32 

33 

34@dataclass(frozen=True, slots=True) 

35class Horizontal: 

36 """Horizontal coordinates.""" 

37 az: float 

38 alt: float 

39 

40 def to_astropy( 

41 self, 

42 units: str | Unit = "deg", 

43 **kwargs, 

44 ): 

45 """Convert to an astropy SkyCoord in the AltAz frame.""" 

46 from astropy import coordinates as ac 

47 

48 units = astropy_unit(units) 

49 return ac.SkyCoord( 

50 az=self.az * units, 

51 alt=self.alt * units, 

52 frame=ReferenceFrame.ALTAZ, 

53 **kwargs, 

54 ) 

55 

56 

57@dataclass(frozen=True, slots=True) 

58class Equatorial: 

59 """Equatorial coordinates.""" 

60 ra: float 

61 dec: float 

62 

63 def to_astropy( 

64 self, 

65 units: str | Unit = "deg", 

66 ra_units: str | Unit | None = None, 

67 frame: ReferenceFrame = ReferenceFrame.ICRF, 

68 **kwargs, 

69 ): 

70 """Convert to an astropy SkyCoord in the given equatorial frame.""" 

71 from astropy import coordinates as ac 

72 

73 dec_unit = astropy_unit(units) 

74 ra_unit = astropy_unit(ra_units) if ra_units is not None else dec_unit 

75 return ac.SkyCoord( 

76 ra=self.ra * ra_unit, 

77 dec=self.dec * dec_unit, 

78 frame=frame.to_astropy(), 

79 **kwargs, 

80 ) 

81 

82 @property 

83 def ra_hms(self): 

84 """Right ascension formatted as a `"H M S"` string.""" 

85 ra = self.ra / 15 

86 hr = int(ra) 

87 fractional = abs(ra - hr) 

88 min = int(fractional * 60) 

89 sec = (fractional * 60 - min) * 60 

90 return f"{hr} {min} {sec}" 

91 

92 @property 

93 def dec_dms(self): 

94 """Declination formatted as a `"D M S"` string.""" 

95 deg = int(self.dec) 

96 fractional = abs(self.dec - deg) 

97 arcmin = int(fractional * 60) 

98 arcsec = (fractional * 60 - arcmin) * 60 

99 return f"{deg} {arcmin} {arcsec}" 

100 

101 

102@dataclass(frozen=True, slots=True) 

103class Geodetic: 

104 """Geodetic coordinates.""" 

105 lon: float 

106 lat: float 

107 elev: float 

108 

109 @functools.cache 

110 def to_astropy(self, angle_units: str | Unit = "deg", distance_units: str | Unit = "m"): 

111 """Convert to an astropy EarthLocation (result is cached).""" 

112 from astropy import coordinates as ac 

113 

114 angle = astropy_unit(angle_units) 

115 distance = astropy_unit(distance_units) 

116 return ac.EarthLocation( 

117 lon=self.lon * angle, 

118 lat=self.lat * angle, 

119 height=self.elev * distance, 

120 ) 

121 

122 

123@dataclass(frozen=True, slots=True) 

124class Cartesian: 

125 """A 3-D Cartesian vector (x, y, z).""" 

126 x: float 

127 y: float 

128 z: float 

129 

130 def __mul__(self, other): 

131 match other: 

132 case float() | int(): 

133 return Cartesian(self.x * other, self.y * other, self.z * other) 

134 case Cartesian(): 

135 return Cartesian(self.x * other.x, self.y * other.y, self.z * other.z) 

136 case _: 

137 raise RuntimeError(f"cannot multiply Cartesian by {type(other)}") 

138 

139 

140@dataclass(frozen=True, slots=True) 

141class StateVector: 

142 """Representation of a state vector.""" 

143 t: datetime 

144 r: Cartesian 

145 v: Cartesian 

146 

147 def to_numpy(self): 

148 """Return a 6-element numpy array `[rx, ry, rz, vx, vy, vz]`.""" 

149 return np.array([self.r.x, self.r.y, self.r.z, self.v.x, self.v.y, self.v.z]) 

150 

151 @classmethod 

152 def from_astropy( 

153 cls, 

154 coord: SkyCoord, 

155 position_units: str | Unit = "m", 

156 velocity_units: str | Unit = "m/s", 

157 ): 

158 """Construct a StateVector from an astropy SkyCoord with Cartesian position and velocity.""" 

159 pos = astropy_unit(position_units) 

160 vel = astropy_unit(velocity_units) 

161 r = coord.cartesian 

162 v = coord.velocity 

163 return cls( 

164 coord.obstime.to_datetime(UTC), 

165 Cartesian( 

166 r.x.to_value(pos), 

167 r.y.to_value(pos), 

168 r.z.to_value(pos), 

169 ), 

170 Cartesian( 

171 v.d_x.to_value(vel), 

172 v.d_y.to_value(vel), 

173 v.d_z.to_value(vel), 

174 ), 

175 ) 

176 

177 def to_astropy( 

178 self, 

179 position_units: str | Unit = "m", 

180 velocity_units: str | Unit = "m/s", 

181 frame: ReferenceFrame = ReferenceFrame.ICRF, 

182 ): 

183 """Convert to an astropy SkyCoord with Cartesian representation and differential.""" 

184 from astropy import coordinates as ac 

185 

186 pos = astropy_unit(position_units) 

187 vel = astropy_unit(velocity_units) 

188 return ac.SkyCoord( 

189 x=self.r.x * pos, 

190 y=self.r.y * pos, 

191 z=self.r.z * pos, 

192 v_x=self.v.x * vel, 

193 v_y=self.v.y * vel, 

194 v_z=self.v.z * vel, 

195 frame=frame.to_astropy(), 

196 representation_type="cartesian", 

197 differential_type="cartesian", 

198 obstime=self.t, 

199 )