Clean up trajectory interface

src/kbmod/run_search.py

@@ -347,9 +347,11 @@ def _count_known_matches(self, result_list, search):
         # Get the pixel positions of results
         ps_list = []
 
+        times = search.stack.build_zeroed_times()
         for row in result_list.results:
-            pix_pos_objs = search.get_trajectory_positions(row.trajectory)
-            PixelPositions = list(map(lambda p: [p.x, p.y], pix_pos_objs))
+            trj = row.trajectory
+            PixelPositions = [[trj.get_x_pos(t), trj.get_y_pos(t)] for t in times]
+
             ps = koffi.PotentialSource()
             ps.build_from_images_and_xy_positions(PixelPositions, metadata)
             ps_list.append(ps)

src/kbmod/search/common.h

@@ -21,6 +21,18 @@ constexpr float NO_DATA = -9999.0;
 
 enum StampType { STAMP_SUM = 0, STAMP_MEAN, STAMP_MEDIAN };
 
+// The position (in pixels) of a trajectory.
+struct PixelPos {
+    float x;
+    float y;
+
+    const std::string to_string() const { return "x: " + std::to_string(x) + " y: " + std::to_string(y); }
+
+    const std::string to_yaml() const {
+        return "{x: " + std::to_string(x) + ", y: " + std::to_string(y) + "}";
+    }
+};
+
 /*
  * Data structure to represent an objects trajectory
  * through a stack of images
@@ -39,6 +51,11 @@ struct Trajectory {
     // Number of images summed
     short obs_count;
 
+    // Get pixel positions from a zero-shifted time.
+    float get_x_pos(float time) const { return x + time * vx; }
+    float get_y_pos(float time) const { return y + time * vy; }
+    PixelPos get_pos(float time) const { return {x + time * vx, y + time * vy}; }
+
     // I can't believe string::format is not a thing until C++ 20
     const std::string to_string() const {
         return "lh: " + std::to_string(lh) + " flux: " + std::to_string(flux) + " x: " + std::to_string(x) +
@@ -54,18 +71,6 @@ struct Trajectory {
     }
 };
 
-// The position (in pixels) of a trajectory.
-struct PixelPos {
-    float x;
-    float y;
-
-    const std::string to_string() const { return "x: " + std::to_string(x) + " y: " + std::to_string(y); }
-
-    const std::string to_yaml() const {
-        return "{x: " + std::to_string(x) + ", y: " + std::to_string(y) + "}";
-    }
-};
-
 /* The parameters to use for the on device search. */
 
 struct SearchParameters {
@@ -152,6 +157,9 @@ static void trajectory_bindings(py::module &m) {
             .def_readwrite("x", &tj::x)
             .def_readwrite("y", &tj::y)
             .def_readwrite("obs_count", &tj::obs_count)
+            .def("get_x_pos", &tj::get_x_pos, pydocs::DOC_Trajectory_get_x_pos)
+            .def("get_y_pos", &tj::get_y_pos, pydocs::DOC_Trajectory_get_y_pos)
+            .def("get_pos", &tj::get_pos, pydocs::DOC_Trajectory_get_pos)
             .def("__repr__", [](const tj &t) { return "Trajectory(" + t.to_string() + ")"; })
             .def("__str__", &tj::to_string)
             .def(py::pickle(

src/kbmod/search/pydocs/common_docs.h

@@ -25,6 +25,48 @@ namespace pydocs {
         Number of observations trajectory was seen in.
   )doc";
 
+  static const auto DOC_Trajectory_get_x_pos = R"doc(
+  Returns the predicted x position of the trajectory.
+
+  Parameters
+  ----------
+  time : `float`
+      A zero shifted time.
+
+  Returns
+  -------
+  `float`
+     The predicted x position (in pixels).
+  )doc";
+
+  static const auto DOC_Trajectory_get_y_pos = R"doc(
+  Returns the predicted y position of the trajectory.
+
+  Parameters
+  ----------
+  time : `float`
+      A zero shifted time.
+
+  Returns
+  -------
+  `float`
+     The predicted y position (in pixels).
+  )doc";
+
+  static const auto DOC_Trajectory_get_pos = R"doc(
+  Returns the predicted (x, y) position of the trajectory.
+
+  Parameters
+  ----------
+  time : `float`
+      A zero shifted time.
+
+  Returns
+  -------
+  `PixelPos`
+     The predicted (x, y) position (in pixels).
+  )doc";
+
   static const auto DOC_PixelPos = R"doc(
   A coordinate pair to store a location on an image.
 

src/kbmod/search/stack_search.cpp

@@ -406,21 +406,6 @@ std::vector<RawImage> StackSearch::get_coadded_stamps_gpu(std::vector<Trajectory
     return results;
 }
 
-PixelPos StackSearch::get_trajectory_position(const Trajectory& t, int i) const {
-    float time = stack.get_zeroed_time(i);
-    return {t.x + time * t.vx, t.y + time * t.vy};
-}
-
-std::vector<PixelPos> StackSearch::get_trajectory_positions(Trajectory& t) const {
-    std::vector<PixelPos> results;
-    int num_times = stack.img_count();
-    for (int i = 0; i < num_times; ++i) {
-        PixelPos pos = get_trajectory_position(t, i);
-        results.push_back(pos);
-    }
-    return results;
-}
-
 std::vector<float> StackSearch::create_curves(Trajectory t, const std::vector<RawImage>& imgs) {
     /*Create a lightcurve from an image along a trajectory
      *
@@ -538,8 +523,6 @@ static void stack_search_bindings(py::module& m) {
             .def("filter_stamp", &ks::filter_stamp, pydocs::DOC_StackSearch_filter_stamp)
             .def("get_trajectory_position", &ks::get_trajectory_position,
                  pydocs::DOC_StackSearch_get_trajectory_position)
-            .def("get_trajectory_positions", &ks::get_trajectory_positions,
-                 pydocs::DOC_StackSearch_get_trajectory_positions)
             .def("get_psi_curves", (std::vector<float>(ks::*)(tj&)) & ks::get_psi_curves,
                  pydocs::DOC_StackSearch_get_psi_curves)
             .def("get_phi_curves", (std::vector<float>(ks::*)(tj&)) & ks::get_phi_curves,

src/kbmod/search/stack_search.h

@@ -44,8 +44,10 @@ class StackSearch {
     std::vector<Trajectory> get_results(int start, int end);
 
     // Get the predicted (pixel) positions for a given trajectory.
-    PixelPos get_trajectory_position(const Trajectory& t, int i) const;
-    std::vector<PixelPos> get_trajectory_positions(Trajectory& t) const;
+    PixelPos get_trajectory_position(const Trajectory& t, int i) const {
+        float time = stack.get_zeroed_time(i);
+        return t.get_pos(time);
+    }
 
     // Filters the results based on various parameters.
     void filter_results(int min_observations);