Use SphericalCoordinates as a column instead of PointSourceDetectinons in Observations

Author: moeyensj

thor/config.py

@@ -1,9 +1,6 @@
 from dataclasses import dataclass
 from typing import Literal, Optional
 
-import numpy as np
-import numpy.typing as npt
-
 
 @dataclass
 class Config:

thor/observations/__init__.py

@@ -1,2 +1,3 @@
 # flake8: noqa: F401
 from .observations import Observations
+from .photometry import Photometry

thor/observations/filters.py

@@ -6,7 +6,7 @@
 import numpy as np
 import quivr as qv
 import ray
-from adam_core.observations import PointSourceDetections
+from adam_core.coordinates import SphericalCoordinates
 
 from ..orbit import TestOrbit, TestOrbitEphemeris
 
@@ -45,7 +45,7 @@ def TestOrbitRadiusObservationFilter_worker(
     # Select the ephemeris and observations for this state
     ephemeris_state = ephemeris.select("id", state_id)
     observations_state = observations.select("state_id", state_id)
-    detections_state = observations_state.detections
+    coordinates_state = observations_state.coordinates
 
     assert (
         len(ephemeris_state) == 1
@@ -56,7 +56,7 @@ def TestOrbitRadiusObservationFilter_worker(
 
     # Return the observations within the radius for this particular state
     return observations_state.apply_mask(
-        _within_radius(detections_state, ephem_ra, ephem_dec, radius)
+        _within_radius(coordinates_state, ephem_ra, ephem_dec, radius)
     )
 
 
@@ -199,7 +199,11 @@ def apply(
 
         observations_filtered = qv.concatenate(filtered_observations_list)
         observations_filtered = observations_filtered.sort_by(
-            ["detections.time.days", "detections.time.nanos", "observatory_code"]
+            [
+                "coordinates.time.days",
+                "coordinates.time.nanos",
+                "coordinates.origin.code",
+            ]
         )
 
         time_end = time.perf_counter()
@@ -213,19 +217,19 @@ def apply(
 
 
 def _within_radius(
-    detections: PointSourceDetections,
+    coords: SphericalCoordinates,
     ra: float,
     dec: float,
     radius: float,
 ) -> np.array:
     """
     Return a boolean mask that identifies which of
-    the detections are within a given radius of a given ra and dec.
+    the coords are within a given radius of a given ra and dec.
 
     Parameters
     ----------
-    detections : `~adam_core.observations.detections.PointSourceDetections`
-        The detections to filter.
+    coords : `~adam_core.coordinates.spherical.SphericalCoordinates`
+        The coords to filter.
     ra : float
         The right ascension of the center of the radius in degrees.
     dec : float
@@ -236,11 +240,11 @@ def _within_radius(
     Returns
     -------
     mask : `~numpy.ndarray`
-        A boolean mask that identifies which of the detections are within
+        A boolean mask that identifies which of the coords are within
         the radius.
     """
-    det_ra = np.deg2rad(detections.ra.to_numpy())
-    det_dec = np.deg2rad(detections.dec.to_numpy())
+    det_ra = np.deg2rad(coords.lon.to_numpy())
+    det_dec = np.deg2rad(coords.lat.to_numpy())
 
     center_ra = np.deg2rad(ra)
     center_dec = np.deg2rad(dec)

thor/observations/observations.py

@@ -9,6 +9,8 @@
 from adam_core.observers import Observers
 from adam_core.time import Timestamp
 
+from .photometry import Photometry
+
 
 class ObserversWithStates(qv.Table):
     state_id = qv.Int64Column()
@@ -26,19 +28,12 @@ class Observations(qv.Table):
     assign a unique state ID. If not using this constructor, please ensure that the detections are sorted
     by time and observatory code and that each unique combination of time and observatory code has a unique
     state ID.
-
-    Columns
-    -------
-    detections : `~adam_core.observations.detections.PointSourceDetections`
-        A table of point source detections.
-    observatory_code : `~qv.StringColumn`
-        The observatory code for each detection.
-    state_id : `~qv.Int64Column`
-        The state ID for each detection.
     """
 
-    detections = PointSourceDetections.as_column()
-    observatory_code = qv.StringColumn()
+    id = qv.StringColumn()
+    exposure_id = qv.StringColumn()
+    coordinates = SphericalCoordinates.as_column()
+    photometry = Photometry.as_column()
     state_id = qv.Int64Column()
 
     @classmethod
@@ -100,13 +95,15 @@ def from_detections_and_exposures(
         )
 
         # Extract the exposure IDs and the observatory codes from the exposures table
-        exposure_obscodes = pa.table(
-            [exposures.id, exposures.observatory_code],
-            names=["exposure_id", "observatory_code"],
+        exposure_filters_obscodes = pa.table(
+            [exposures.id, exposures.filter, exposures.observatory_code],
+            names=["exposure_id", "filter", "observatory_code"],
         )
 
         # Join the detection times and the exposure IDs so that each detection has an observatory code
-        obscode_times = detections_flattened.join(exposure_obscodes, ["exposure_id"])
+        obscode_times = detections_flattened.join(
+            exposure_filters_obscodes, ["exposure_id"]
+        )
 
         # Group the detections by the observatory code and the detection times and then grab the unique ones
         unique_obscode_times = obscode_times.group_by(
@@ -139,23 +136,34 @@ def from_detections_and_exposures(
             [("state_id", "ascending")]
         )
 
+        sigmas = np.zeros((len(detections_with_states), 6))
+        sigmas[:, 1] = detections_with_states["ra_sigma"].to_numpy(zero_copy_only=False)
+        sigmas[:, 2] = detections_with_states["dec_sigma"].to_numpy(
+            zero_copy_only=False
+        )
+
         return cls.from_kwargs(
-            detections=PointSourceDetections.from_kwargs(
-                id=detections_with_states["id"],
-                exposure_id=detections_with_states["exposure_id"],
+            id=detections_with_states["id"],
+            exposure_id=detections_with_states["exposure_id"],
+            coordinates=SphericalCoordinates.from_kwargs(
+                lon=detections_with_states["ra"],
+                lat=detections_with_states["dec"],
                 time=Timestamp.from_kwargs(
                     days=detections_with_states["days"],
                     nanos=detections_with_states["nanos"],
                     scale="utc",
                 ),
-                ra=detections_with_states["ra"],
-                ra_sigma=detections_with_states["ra_sigma"],
-                dec=detections_with_states["dec"],
-                dec_sigma=detections_with_states["dec_sigma"],
+                covariance=CoordinateCovariances.from_sigmas(sigmas),
+                origin=Origin.from_kwargs(
+                    code=detections_with_states["observatory_code"]
+                ),
+                frame="equatorial",
+            ),
+            photometry=Photometry.from_kwargs(
+                filter=detections_with_states["filter"],
                 mag=detections_with_states["mag"],
                 mag_sigma=detections_with_states["mag_sigma"],
             ),
-            observatory_code=detections_with_states["observatory_code"],
             state_id=detections_with_states["state_id"],
         )
 
@@ -174,7 +182,7 @@ def get_observers(self) -> ObserversWithStates:
         for code, observations_i in self.group_by_observatory_code():
             # Extract unique times and make sure they are sorted
             # by time in ascending order
-            unique_times = observations_i.detections.time.unique()
+            unique_times = observations_i.coordinates.time.unique()
             unique_times = unique_times.sort_by(["days", "nanos"])
 
             # States are defined by unique times and observatory codes and
@@ -194,28 +202,6 @@ def get_observers(self) -> ObserversWithStates:
         observers = qv.concatenate(observers_with_states)
         return observers.sort_by("state_id")
 
-    def to_spherical_coordinates(self) -> SphericalCoordinates:
-        """
-        Convert the observations to spherical coordinates which can be used
-        to calculate residuals.
-
-        Returns
-        -------
-        coordinates : `~adam_core.coordinates.spherical.SphericalCoordinates`
-            The detections represented as spherical coordinates.
-        """
-        sigmas = np.zeros((len(self), 6))
-        sigmas[:, 1] = self.detections.ra_sigma.to_numpy(zero_copy_only=False)
-        sigmas[:, 2] = self.detections.dec_sigma.to_numpy(zero_copy_only=False)
-        return SphericalCoordinates.from_kwargs(
-            lon=self.detections.ra,
-            lat=self.detections.dec,
-            time=self.detections.time,
-            covariance=CoordinateCovariances.from_sigmas(sigmas),
-            origin=Origin.from_kwargs(code=self.observatory_code),
-            frame="equatorial",
-        )
-
     def select_exposure(self, exposure_id: int) -> "Observations":
         """
         Select observations from a single exposure.
@@ -230,7 +216,7 @@ def select_exposure(self, exposure_id: int) -> "Observations":
         observations : `~Observations`
             Observations from the specified exposure.
         """
-        return self.apply_mask(pc.equal(self.detections.exposure_id, exposure_id))
+        return self.apply_mask(pc.equal(self.exposure_id, exposure_id))
 
     def group_by_exposure(self) -> Iterator[Tuple[str, "Observations"]]:
         """
@@ -242,7 +228,7 @@ def group_by_exposure(self) -> Iterator[Tuple[str, "Observations"]]:
         observations : Iterator[`~thor.observations.observations.Observations`]
             Observations belonging to individual exposures.
         """
-        exposure_ids = self.detections.exposure_id
+        exposure_ids = self.exposure_id
         for exposure_id in exposure_ids.unique().sort():
             yield exposure_id.as_py(), self.select_exposure(exposure_id)
 
@@ -260,7 +246,7 @@ def select_observatory_code(self, observatory_code) -> "Observations":
         observations : `~Observations`
             Observations from the specified observatory.
         """
-        return self.apply_mask(pc.equal(self.observatory_code, observatory_code))
+        return self.apply_mask(pc.equal(self.coordinates.origin.code, observatory_code))
 
     def group_by_observatory_code(self) -> Iterator[Tuple[str, "Observations"]]:
         """
@@ -271,7 +257,7 @@ def group_by_observatory_code(self) -> Iterator[Tuple[str, "Observations"]]:
         observations : Iterator[`~thor.observations.observations.Observations`]
             Observations belonging to individual observatories.
         """
-        observatory_codes = self.observatory_code
+        observatory_codes = self.coordinates.origin.code
         for observatory_code in observatory_codes.unique().sort():
             yield observatory_code.as_py(), self.select_observatory_code(
                 observatory_code

thor/orbit.py

@@ -76,48 +76,19 @@ def range_observations_worker(
     """
     observations_state = observations.select("state_id", state_id)
     ephemeris_state = ephemeris.select("id", state_id)
-    detections_state = observations_state.detections
 
     # Get the heliocentric position vector of the object at the time of the exposure
     r = ephemeris_state.ephemeris.aberrated_coordinates.r[0]
 
     # Get the observer's heliocentric coordinates
     observer_i = ephemeris_state.observer
 
-    # Create an array of observatory codes for the detections
-    num_detections = len(observations_state)
-    observatory_codes = np.repeat(
-        observations_state.observatory_code[0].as_py(), num_detections
-    )
-
-    # The following can be replaced with:
-    # coords = observations_state.to_spherical(observatory_codes)
-    # Start replacement:
-    sigma_data = np.vstack(
-        [
-            pa.nulls(num_detections, pa.float64()),
-            detections_state.ra_sigma.to_numpy(zero_copy_only=False),
-            detections_state.dec_sigma.to_numpy(zero_copy_only=False),
-            pa.nulls(num_detections, pa.float64()),
-            pa.nulls(num_detections, pa.float64()),
-            pa.nulls(num_detections, pa.float64()),
-        ]
-    ).T
-    coords = SphericalCoordinates.from_kwargs(
-        lon=detections_state.ra,
-        lat=detections_state.dec,
-        time=detections_state.time,
-        covariance=CoordinateCovariances.from_sigmas(sigma_data),
-        origin=Origin.from_kwargs(code=observatory_codes),
-        frame="equatorial",
-    )
-    # End replacement (only once
-    # https://github.com/B612-Asteroid-Institute/adam_core/pull/45 is merged)
-
     return RangedPointSourceDetections.from_kwargs(
-        id=detections_state.id,
-        exposure_id=detections_state.exposure_id,
-        coordinates=assume_heliocentric_distance(r, coords, observer_i.coordinates),
+        id=observations_state.id,
+        exposure_id=observations_state.exposure_id,
+        coordinates=assume_heliocentric_distance(
+            r, observations_state.coordinates, observer_i.coordinates
+        ),
         state_id=observations_state.state_id,
     )
 
@@ -205,9 +176,7 @@ def _is_cache_fresh(self, observations: Observations) -> bool:
         if self._cached_ephemeris is None or self._cached_observation_ids is None:
             return False
         elif pc.all(
-            pc.is_in(
-                observations.detections.id.sort(), self._cached_observation_ids.sort()
-            )
+            pc.is_in(observations.id.sort(), self._cached_observation_ids.sort())
         ).as_py():
             return True
         else:
@@ -231,7 +200,7 @@ def _cache_ephemeris(
         None
         """
         self._cached_ephemeris = ephemeris
-        self._cached_observation_ids = observations.detections.id
+        self._cached_observation_ids = observations.id
 
     def propagate(
         self,