Add magnitude to Ephemeris, fix type error in test code

B612-Asteroid-Institute · Apr 14, 2023 · bd9942c · bd9942c
1 parent b1d5cd4
commit bd9942c
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 10 deletions.
diff --git a/precovery/orbit.py b/precovery/orbit.py
@@ -388,13 +388,20 @@ def precover_remote(
 
 class Ephemeris:
     def __init__(
-        self, mjd: float, ra: float, dec: float, ra_velocity: float, dec_velocity: float
+        self,
+        mjd: float,
+        ra: float,
+        dec: float,
+        ra_velocity: float,
+        dec_velocity: float,
+        mag: float,
     ):
         self.mjd = mjd
         self.ra = ra
         self.dec = dec
         self.ra_velocity = ra_velocity
         self.dec_velocity = dec_velocity
+        self.mag = mag
 
     @classmethod
     def from_pyoorb_vector(cls, raw_data: npt.NDArray[np.float64]):
@@ -404,7 +411,8 @@ def from_pyoorb_vector(cls, raw_data: npt.NDArray[np.float64]):
         # oorb returns vracos(dec), so lets remove the cos(dec) term
         ra_velocity = raw_data[3] / np.cos(np.radians(dec))  # deg per day
         dec_velocity = raw_data[4]  # deg per day
-        return cls(mjd, ra, dec, ra_velocity, dec_velocity)
+        mag = raw_data[9]
+        return cls(mjd, ra, dec, ra_velocity, dec_velocity, mag)
 
     def __str__(self):
         return f"<Ephemeris ra={self.ra:.4f} dec={self.dec:.4f} mjd={self.mjd:.6f}>"

diff --git a/tests/make_observations.py b/tests/make_observations.py
@@ -209,10 +209,10 @@ def make_observations(
                 # geocentric distance (au)
                 # predicted apparent V-band magnitude
                 # true anomaly (deg)
-                ephemeris_dict["mjd"].append(eph_i._raw_data[0])
-                ephemeris_dict["ra"].append(eph_i._raw_data[1])
-                ephemeris_dict["dec"].append(eph_i._raw_data[2])
-                ephemeris_dict["mag"].append(eph_i._raw_data[9])
+                ephemeris_dict["mjd"].append(eph_i.mjd)
+                ephemeris_dict["ra"].append(eph_i.ra)
+                ephemeris_dict["dec"].append(eph_i.dec)
+                ephemeris_dict["mag"].append(eph_i.mag)
 
             ephemeris_df = pd.DataFrame(ephemeris_dict)
             ephemeris_df.insert(0, "object_id", orbit_ids[orbit.orbit_id])

diff --git a/tests/test_orbit.py b/tests/test_orbit.py
@@ -94,7 +94,9 @@ def test_orbit_ephemeris_computation():
 
 
 def test_ephemeris_distance():
-    ephemeris = orbit.Ephemeris(mjd=1, ra=2, dec=3, ra_velocity=4, dec_velocity=5)
+    ephemeris = orbit.Ephemeris(
+        mjd=1, ra=2, dec=3, ra_velocity=4, dec_velocity=5, mag=6
+    )
 
     observations = ObservationArray(
         [
@@ -105,7 +107,7 @@ def test_ephemeris_distance():
                 id=b"exact",
                 ra_sigma=0,
                 dec_sigma=0,
-                mag=0,
+                mag=6,
                 mag_sigma=0,
             ),
             Observation(
@@ -115,7 +117,7 @@ def test_ephemeris_distance():
                 id=b"one_degree_dec",
                 ra_sigma=0,
                 dec_sigma=0,
-                mag=0,
+                mag=6,
                 mag_sigma=0,
             ),
             Observation(
@@ -125,7 +127,7 @@ def test_ephemeris_distance():
                 id=b"wraparound",
                 ra_sigma=0,
                 dec_sigma=0,
-                mag=0,
+                mag=6,
                 mag_sigma=0,
             ),
         ]