Add a neighborhood search to TrajectoryExplorer

notebooks/TrajectoryExplorer.ipynb

@@ -244,6 +244,30 @@
     "for i in range(len(fake_times)):\n",
     "    print(f\"Time {i} is valid = {result['obs_valid'][0][i]}.\")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Neighborhood Searches\n",
+    "\n",
+    "The `TrajectoryExplorer` can also be used to perform a hyper-localized search.  This search effectively uses a small neighborhood around a given trajectory (both in terms of starting pixel and velocities) and returns all results from this neighborhood. This localized set can be used to:\n",
+    "1) refine trajectories by searching a finer parameter space around the best results found by the initial search, or\n",
+    "2) collect a distribution of trajectories and their likelihoods around a single result.\n",
+    "For this search, the `TrajectoryExplorer` does not perform any filtering, so it will return all trajectories and their likelihoods (even one <= -1.0)\n",
+    "\n",
+    "Only basic statistics, such as likelihood and flux, are returned from the search. Stamps and lightcurves are not computed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "samples = explorer.evaluate_around_linear_trajectory(50, 60, 5.0, -2.0, pixel_radius=5)\n",
+    "print(samples[0:10])"
+   ]
   }
  ],
  "metadata": {

src/kbmod/run_search.py

@@ -19,6 +19,54 @@
 logger = kb.Logging.getLogger(__name__)
 
 
+def configure_kb_search_stack(search, config):
+    """Configure the kbmod SearchStack object from a search configuration.
+
+    Parameters
+    ----------
+    search : `kb.StackSearch`
+        The SearchStack object.
+    config : `SearchConfiguration`
+        The configuration parameters
+    """
+    width = search.get_image_width()
+    height = search.get_image_height()
+
+    # Set the search bounds.
+    if config["x_pixel_bounds"] and len(config["x_pixel_bounds"]) == 2:
+        search.set_start_bounds_x(config["x_pixel_bounds"][0], config["x_pixel_bounds"][1])
+    elif config["x_pixel_buffer"] and config["x_pixel_buffer"] > 0:
+        search.set_start_bounds_x(-config["x_pixel_buffer"], width + config["x_pixel_buffer"])
+
+    if config["y_pixel_bounds"] and len(config["y_pixel_bounds"]) == 2:
+        search.set_start_bounds_y(config["y_pixel_bounds"][0], config["y_pixel_bounds"][1])
+    elif config["y_pixel_buffer"] and config["y_pixel_buffer"] > 0:
+        search.set_start_bounds_y(-config["y_pixel_buffer"], height + config["y_pixel_buffer"])
+
+    # Set the results per pixel.
+    search.set_results_per_pixel(config["results_per_pixel"])
+
+    # If we are using gpu_filtering, enable it and set the parameters.
+    if config["gpu_filter"]:
+        logger.debug("Using in-line GPU sigmaG filtering methods")
+        coeff = SigmaGClipping.find_sigma_g_coeff(
+            config["sigmaG_lims"][0],
+            config["sigmaG_lims"][1],
+        )
+        search.enable_gpu_sigmag_filter(
+            np.array(config["sigmaG_lims"]) / 100.0,
+            coeff,
+            config["lh_level"],
+        )
+    else:
+        search.disable_gpu_sigmag_filter()
+
+    # If we are using an encoded image representation on GPU, enable it and
+    # set the parameters.
+    if config["encode_num_bytes"] > 0:
+        search.enable_gpu_encoding(config["encode_num_bytes"])
+
+
 class SearchRunner:
     """A class to run the KBMOD grid search."""
 
@@ -139,45 +187,11 @@ def do_gpu_search(self, config, stack, trj_generator):
         """
         # Create the search object which will hold intermediate data and results.
         search = kb.StackSearch(stack)
-
-        width = search.get_image_width()
-        height = search.get_image_height()
-
-        # Set the search bounds.
-        if config["x_pixel_bounds"] and len(config["x_pixel_bounds"]) == 2:
-            search.set_start_bounds_x(config["x_pixel_bounds"][0], config["x_pixel_bounds"][1])
-        elif config["x_pixel_buffer"] and config["x_pixel_buffer"] > 0:
-            search.set_start_bounds_x(-config["x_pixel_buffer"], width + config["x_pixel_buffer"])
-
-        if config["y_pixel_bounds"] and len(config["y_pixel_bounds"]) == 2:
-            search.set_start_bounds_y(config["y_pixel_bounds"][0], config["y_pixel_bounds"][1])
-        elif config["y_pixel_buffer"] and config["y_pixel_buffer"] > 0:
-            search.set_start_bounds_y(-config["y_pixel_buffer"], height + config["y_pixel_buffer"])
-
-        # Set the results per pixel.
-        search.set_results_per_pixel(config["results_per_pixel"])
+        configure_kb_search_stack(search, config)
 
         search_timer = kb.DebugTimer("grid search", logger)
         logger.debug(f"{trj_generator}")
 
-        # If we are using gpu_filtering, enable it and set the parameters.
-        if config["gpu_filter"]:
-            logger.debug("Using in-line GPU sigmaG filtering methods")
-            coeff = SigmaGClipping.find_sigma_g_coeff(
-                config["sigmaG_lims"][0],
-                config["sigmaG_lims"][1],
-            )
-            search.enable_gpu_sigmag_filter(
-                np.array(config["sigmaG_lims"]) / 100.0,
-                coeff,
-                config["lh_level"],
-            )
-
-        # If we are using an encoded image representation on GPU, enable it and
-        # set the parameters.
-        if config["encode_num_bytes"] > 0:
-            search.enable_gpu_encoding(config["encode_num_bytes"])
-
         # Do the actual search.
         candidates = [trj for trj in trj_generator]
         search.search_all(candidates, int(config["num_obs"]))

src/kbmod/search/kernels/kernels.cu

@@ -251,7 +251,7 @@ __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_v
     for (int r = 0; r < params.results_per_pixel; ++r) {
         results[base_index + r].x = x;
         results[base_index + r].y = y;
-        results[base_index + r].lh = -1.0;
+        results[base_index + r].lh = -FLT_MAX;
     }
 
     // For each trajectory we'd like to search
@@ -274,10 +274,9 @@ __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_v
             continue;
 
         // Insert the new trajectory into the sorted list of final results.
-        // Only sort the values with valid likelihoods.
         Trajectory temp;
         for (unsigned int r = 0; r < params.results_per_pixel; ++r) {
-            if (curr_trj.lh > results[base_index + r].lh && curr_trj.lh > -1.0) {
+            if (curr_trj.lh > results[base_index + r].lh) {
                 temp = results[base_index + r];
                 results[base_index + r] = curr_trj;
                 curr_trj = temp;

src/kbmod/search/pydocs/stack_search_docs.h

@@ -30,6 +30,10 @@ static const auto DOC_StackSearch_set_min_lh = R"doc(
       The minimum likelihood value for a trajectory to be returned.
   )doc";
 
+static const auto DOC_StackSearch_disable_gpu_sigmag_filter = R"doc(
+  Turns off the on-GPU sigma-G filtering.
+  )doc";
+
 static const auto DOC_StackSearch_enable_gpu_sigmag_filter = R"doc(
   Enable on-GPU sigma-G filtering.
 

src/kbmod/search/stack_search.cpp

@@ -66,6 +66,10 @@ void StackSearch::enable_gpu_sigmag_filter(std::vector<float> percentiles, float
     params.sigmag_coeff = sigmag_coeff;
     params.min_lh = min_lh;
 }
+
+void StackSearch::disable_gpu_sigmag_filter() {
+    params.do_sigmag_filter = false;
+}
 
 void StackSearch::enable_gpu_encoding(int encode_num_bytes) {
     // If changing a setting that would impact the search data encoding, clear the cached values.
@@ -310,6 +314,8 @@ static void stack_search_bindings(py::module& m) {
             .def("set_min_lh", &ks::set_min_lh, pydocs::DOC_StackSearch_set_min_lh)
             .def("set_results_per_pixel", &ks::set_results_per_pixel,
                  pydocs::DOC_StackSearch_set_results_per_pixel)
+            .def("disable_gpu_sigmag_filter", &ks::disable_gpu_sigmag_filter,
+                 pydocs::DOC_StackSearch_disable_gpu_sigmag_filter)
             .def("enable_gpu_sigmag_filter", &ks::enable_gpu_sigmag_filter,
                  pydocs::DOC_StackSearch_enable_gpu_sigmag_filter)
             .def("enable_gpu_encoding", &ks::enable_gpu_encoding, pydocs::DOC_StackSearch_enable_gpu_encoding)

src/kbmod/search/stack_search.h

@@ -40,6 +40,7 @@ class StackSearch {
     // Parameter setters used to control the searches.
     void set_min_obs(int new_value);
     void set_min_lh(float new_value);
+    void disable_gpu_sigmag_filter();
     void enable_gpu_sigmag_filter(std::vector<float> percentiles, float sigmag_coeff, float min_lh);
     void enable_gpu_encoding(int num_bytes);
     void set_start_bounds_x(int x_min, int x_max);

src/kbmod/trajectory_explorer.py

@@ -3,8 +3,10 @@
 from kbmod.configuration import SearchConfiguration
 from kbmod.filters.sigma_g_filter import apply_clipped_sigma_g, SigmaGClipping
 from kbmod.results import Results
-from kbmod.search import StackSearch, StampCreator, Logging
+from kbmod.run_search import configure_kb_search_stack
+from kbmod.search import DebugTimer, StackSearch, StampCreator, Logging
 from kbmod.filters.stamp_filters import append_all_stamps, append_coadds
+from kbmod.trajectory_generator import PencilSearch
 from kbmod.trajectory_utils import make_trajectory_from_ra_dec
 
 
@@ -42,9 +44,24 @@ def __init__(self, img_stack, config=None):
         # Allocate and configure the StackSearch object.
         self.search = None
 
-    def initialize_data(self):
-        """Perform any needed initialization and preprocessing on the images."""
+    def initialize_data(self, config=None):
+        """Initialize the data, including applying the configuration parameters.
+
+        Parameters
+        ----------
+        config : `SearchConfiguration`, optional
+            Any custom configuration parameters to use for this run.
+            If ``None`` uses the default configuration parameters.
+        """
+        if config is None:
+            config = self.config
+
         if self._data_initalized:
+            # Always reapply the configuration parameters if in case we used custom
+            # ones on a previous search.
+            configure_kb_search_stack(self.search, config)
+
+            # Nothing else to do
             return
 
         # If we are using an encoded image representation on GPU, enable it and
@@ -55,6 +72,7 @@ def initialize_data(self):
 
         # Allocate the search structure.
         self.search = StackSearch(self.im_stack)
+        configure_kb_search_stack(self.search, config)
 
         self._data_initalized = True
 
@@ -121,6 +139,82 @@ def evaluate_angle_trajectory(self, ra, dec, v_ra, v_dec, wcs):
         trj = make_trajectory_from_ra_dec(ra, dec, v_ra, v_dec, wcs)
         return self.evaluate_linear_trajectory(trj.x, trj.y, trj.vx, trj.vy)
 
+    def evaluate_around_linear_trajectory(
+        self,
+        x,
+        y,
+        vx,
+        vy,
+        pixel_radius=5,
+        max_ang_offset=0.2618,
+        ang_step=0.035,
+        max_vel_offset=10.0,
+        vel_step=0.5,
+    ):
+        """Evaluate all the trajectories within a local neighborhood of the given trajectory.
+        No filtering is done at all.
+
+        Parameters
+        ----------
+        x : `int`
+            The starting x pixel of the trajectory.
+        y : `int`
+            The starting y pixel of the trajectory.
+        vx : `float`
+            The x velocity of the trajectory in pixels per day.
+        vy : `float`
+            The y velocity of the trajectory in pixels per day.
+        pixel_radius : `int`
+            The number of pixels to evaluate to each side of the Trajectory's starting pixel.
+        max_ang_offset : `float`
+            The maximum offset of a candidate trajectory from the original (in radians)
+        ang_step : `float`
+            The step size to explore for each angle (in radians)
+        max_vel_offset : `float`
+            The maximum offset of the velocity's magnitude from the original (in pixels per day)
+        vel_step : `float`
+            The step size to explore for each velocity magnitude (in pixels per day)
+
+        Returns
+        -------
+        result : `Results`
+            The results table with a single row and all the columns filled out.
+        """
+        if pixel_radius < 0:
+            raise ValueError(f"Pixel radius must be >= 0. Got {pixel_radius}")
+        num_pixels = (2 * pixel_radius + 1) * (2 * pixel_radius + 1)
+        logger.debug(f"Testing {num_pixels} starting pixels.")
+
+        # Create a pencil search around the given trajectory.
+        trj_generator = PencilSearch(vx, vy, max_ang_offset, ang_step, max_vel_offset, vel_step)
+        num_trj = len(trj_generator)
+        logger.debug(f"Exploring {num_trj} trajectories per starting pixel.")
+
+        # Set the search bounds to right around the trajectory's starting position and
+        # turn off all filtering.
+        reduced_config = self.config.copy()
+        reduced_config.set("x_pixel_bounds", [x - pixel_radius, x + pixel_radius + 1])
+        reduced_config.set("y_pixel_bounds", [y - pixel_radius, y + pixel_radius + 1])
+        reduced_config.set("results_per_pixel", min(num_trj, 10_000))
+        reduced_config.set("gpu_filter", False)
+        reduced_config.set("num_obs", 1)
+        reduced_config.set("max_lh", 1e25)
+        reduced_config.set("lh_level", -1e25)
+        self.initialize_data(config=reduced_config)
+
+        # Do the actual search.
+        search_timer = DebugTimer("grid search", logger)
+        candidates = [trj for trj in trj_generator]
+        self.search.search_all(candidates, int(reduced_config["num_obs"]))
+        search_timer.stop()
+
+        # Load all of the results without any filtering.
+        logger.debug(f"Loading {num_pixels * num_trj} results.")
+        trjs = self.search.get_results(0, num_pixels * num_trj)
+        results = Results.from_trajectories(trjs)
+
+        return results
+
     def apply_sigma_g(self, result):
         """Apply sigma G clipping to a single ResultRow. Modifies the row in-place.
 

tests/test_trajectory_explorer.py

@@ -33,7 +33,7 @@ def setUp(self):
         )
         fake_ds.insert_object(self.trj)
 
-        # Remove at least observation from the trajectory.
+        # Remove at least one observation from the trajectory.
         pred_x = self.trj.get_x_index(fake_times[10])
         pred_y = self.trj.get_y_index(fake_times[10])
         sci_t10 = fake_ds.stack.get_single_image(10).get_science()
@@ -80,6 +80,40 @@ def test_evaluate_trajectory(self):
         self.explorer.apply_sigma_g(result)
         self.assertFalse(result["obs_valid"][0][10])
 
+    @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")
+    def test_evaluate_around_linear_trajectory(self):
+        radius = 3
+        edge_length = 2 * radius + 1
+        num_pixels = edge_length * edge_length
+
+        results = self.explorer.evaluate_around_linear_trajectory(
+            self.x0,
+            self.y0,
+            self.vx,
+            self.vy,
+            pixel_radius=radius,
+            max_ang_offset=0.2618,
+            ang_step=0.035,
+            max_vel_offset=10.0,
+            vel_step=0.5,
+        )
+
+        # Using the above settings should provide 615 trajectories per starting pixel.
+        self.assertEqual(len(results), num_pixels * 615)
+
+        # Count the number of results we have per starting pixel.
+        counts = np.zeros((edge_length, edge_length))
+        for row in range(len(results)):
+            self.assertGreaterEqual(results["x"][row], self.x0 - 3)
+            self.assertLessEqual(results["x"][row], self.x0 + 3)
+            self.assertGreaterEqual(results["y"][row], self.y0 - 3)
+            self.assertLessEqual(results["y"][row], self.y0 + 3)
+
+            x = results["x"][row] - self.x0 + 3
+            y = results["y"][row] - self.y0 + 3
+            counts[y, x] += 1
+        self.assertTrue(np.all(counts == 615))
+
 
 if __name__ == "__main__":
     unittest.main()