Merge pull request #458 from dirac-institute/stack_search_single

Expands StackSearch to support single searches

src/kbmod/filters/sigma_g_filter.py

@@ -44,6 +44,27 @@ def __init__(self, low_bnd=25, high_bnd=75, n_sigma=2, clip_negative=False):
 
     @staticmethod
     def find_sigma_g_coeff(low_bnd, high_bnd):
+        """Compute the sigma G coefficient from the upper and lower bounds
+        of the percentiles.
+
+        Parameters
+        ----------
+        low_bnd : `float`
+            The lower bound of the percentile on a scale of [0, 100].
+        high_bnd : `float`
+            The lower bound of the percentile on a scale of [0, 100].
+
+        Returns
+        -------
+        result : `float`
+            The corresponding sigma G coefficient.
+
+        Raises
+        ------
+        Raises a ``ValueError`` is the bounds are invalid.
+        """
+        if (high_bnd <= low_bnd) or (low_bnd < 0) or (high_bnd > 100):
+            raise ValueError(f"Invalid percentiles for sigma G coefficient [{low_bnd}, {high_bnd}]")
         x1 = SigmaGClipping.invert_gauss_cdf(low_bnd / 100.0)
         x2 = SigmaGClipping.invert_gauss_cdf(high_bnd / 100.0)
         return 1 / (x2 - x1)

src/kbmod/search/kernels.cu

@@ -24,6 +24,9 @@
 namespace search {
 
 extern "C" void device_allocate_psi_phi_arrays(PsiPhiArray *data) {
+    if (data == nullptr) {
+        throw std::runtime_error("No data given.");
+    }
     if (!data->cpu_array_allocated() || !data->cpu_time_array_allocated()) {
         throw std::runtime_error("CPU data is not allocated.");
     }
@@ -48,6 +51,9 @@ extern "C" void device_allocate_psi_phi_arrays(PsiPhiArray *data) {
 }
 
 extern "C" void device_free_psi_phi_arrays(PsiPhiArray *data) {
+    if (data == nullptr) {
+        throw std::runtime_error("No data given.");
+    }
     if (data->gpu_array_allocated()) {
         checkCudaErrors(cudaFree(data->get_gpu_array_ptr()));
         data->set_gpu_array_ptr(nullptr);
@@ -61,7 +67,8 @@ extern "C" void device_free_psi_phi_arrays(PsiPhiArray *data) {
 __host__ __device__ PsiPhi read_encoded_psi_phi(PsiPhiArrayMeta &params, void *psi_phi_vect, int time,
                                                 int row, int col) {
     // Bounds checking.
-    if ((row < 0) || (col < 0) || (row >= params.height) || (col >= params.width)) {
+    if ((row < 0) || (col < 0) || (row >= params.height) || (col >= params.width) ||
+        (psi_phi_vect == nullptr)) {
         return {NO_DATA, NO_DATA};
     }
 
@@ -92,6 +99,9 @@ extern "C" __device__ __host__ void SigmaGFilteredIndicesCU(float *values, int n
                                                             float sgl1, float sigmag_coeff, float width,
                                                             int *idx_array, int *min_keep_idx,
                                                             int *max_keep_idx) {
+    // Basic data checking.
+    assert(idx_array != nullptr && min_keep_idx != nullptr && max_keep_idx != nullptr);
+
     // Clip the percentiles to [0.01, 99.99] to avoid invalid array accesses.
     if (sgl0 < 0.0001) sgl0 = 0.0001;
     if (sgl1 > 0.9999) sgl1 = 0.9999;
@@ -147,6 +157,9 @@ extern "C" __device__ __host__ void SigmaGFilteredIndicesCU(float *values, int n
 extern "C" __device__ __host__ void evaluateTrajectory(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_vect,
                                                        float *image_times, SearchParameters params,
                                                        Trajectory *candidate) {
+    // Basic data validity check.
+    assert(psi_phi_vect != nullptr && image_times != nullptr && candidate != nullptr);
+
     // Data structures used for filtering. We fill in only what we need.
     float psi_array[MAX_NUM_IMAGES];
     float phi_array[MAX_NUM_IMAGES];
@@ -226,6 +239,10 @@ extern "C" __device__ __host__ void evaluateTrajectory(PsiPhiArrayMeta psi_phi_m
 __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_vect, float *image_times,
                                    SearchParameters params, int num_trajectories, Trajectory *trajectories,
                                    Trajectory *results) {
+    // Basic data validity check.
+    assert(psi_phi_vect != nullptr && image_times != nullptr && trajectories != nullptr &&
+           results != nullptr);
+
     // Get the x and y coordinates within the search space.
     const int x_i = blockIdx.x * THREAD_DIM_X + threadIdx.x;
     const int y_i = blockIdx.y * THREAD_DIM_Y + threadIdx.y;
@@ -294,6 +311,9 @@ __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_v
 
 extern "C" void deviceSearchFilter(PsiPhiArray &psi_phi_array, SearchParameters params, int num_trajectories,
                                    Trajectory *trj_to_search, int num_results, Trajectory *best_results) {
+    // Basic data validity check.
+    assert(trj_to_search != nullptr && best_results != nullptr);
+
     // Allocate Device memory
     Trajectory *device_tests;
     Trajectory *device_search_results;

src/kbmod/search/pydocs/stack_search_docs.h

@@ -10,8 +10,41 @@ static const auto DOC_StackSearch_search = R"doc(
   todo
   )doc";
 
+static const auto DOC_StackSearch_set_min_obs = R"doc(
+  Sets the minimum number of observations for valid result.
+
+  Parameters
+  ----------
+  new_value : `int`
+      The minimum number of observations for a trajectory to be returned.
+  )doc";
+
+static const auto DOC_StackSearch_set_min_lh = R"doc(
+  Sets the minimum likelihood for valid result.
+
+  Parameters
+  ----------
+  new_value : `float`
+      The minimum likelihood value for a trajectory to be returned.
+  )doc";
+
+
 static const auto DOC_StackSearch_enable_gpu_sigmag_filter = R"doc(
-  todo
+  Enable on-GPU sigma-G filtering.
+
+  Parameters
+  ----------
+  percentiles : `list`
+      A length 2 list of percentiles (between 0.0 and 1.0). Example [0.25, 0.75].
+  sigmag_coeff : `float`
+      The sigma-G coefficient corresponding to the percentiles. This can
+      be computed via SigmaGClipping.find_sigma_g_coeff().
+  min_lh : `float`
+      The minimum likelihood for a result to be accepted.
+
+  Raises
+  ------
+  Raises a ``RunTimeError`` if invalid values are provided.
   )doc";
 
 static const auto DOC_StackSearch_enable_gpu_encoding = R"doc(
@@ -36,6 +69,10 @@ static const auto DOC_StackSearch_set_start_bounds_x = R"doc(
       The inclusive lower bound of the search.
   x_max : `int`
       The exclusive upper bound of the search.
+
+  Raises
+  ------
+  Raises a ``RunTimeError`` if invalid bounds are provided (x_max > x_min).
   )doc";
 
 static const auto DOC_StackSearch_set_start_bounds_y = R"doc(
@@ -48,6 +85,10 @@ static const auto DOC_StackSearch_set_start_bounds_y = R"doc(
       The inclusive lower bound of the search.
   y_max : `int`
       The exclusive upper bound of the search.
+
+  Raises
+  ------
+  Raises a ``RunTimeError`` if invalid bounds are provided (x_max > x_min).
   )doc";
 
 static const auto DOC_StackSearch_set_debug = R"doc(
@@ -107,6 +148,10 @@ static const auto DOC_StackSearch_get_phi_curves = R"doc(
      The phi values at each time step with NO_DATA replaced by 0.0.
   )doc";
 
+static const auto DOC_StackSearch_clear_psi_phi = R"doc(
+  Clear the pre-computed psi and phi data.
+  )doc";
+
 static const auto DOC_StackSearch_prepare_psi_phi = R"doc(
   Compute the cached psi and phi data.
   )doc";
@@ -119,5 +164,43 @@ static const auto DOC_StackSearch_set_results = R"doc(
   todo
   )doc";
 
+static const auto DOC_StackSearch_evaluate_single_trajectory = R"doc(
+  Performs the evaluation of a single Trajectory object. Modifies the object
+  in-place.
+
+  Note
+  ----
+  Runs on the CPU, but requires CUDA compiler.
+
+  Parameters
+  ----------
+  trj : `kb.Trajectory`
+      The trjactory to evaluate.
+   )doc";
+
+static const auto DOC_StackSearch_search_linear_trajectory = R"doc(
+  Performs the evaluation of a linear trajectory in pixel space.
+
+  Note
+  ----
+  Runs on the CPU, but requires CUDA compiler.
+
+  Parameters
+  ----------
+  x : `short`
+      The starting x pixel of the trajectory.
+  y : `short`
+      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.
+
+  Returns
+  -------
+  result : `kb.Trajectory`
+      The trajectory object with statistics set.
+   )doc";
+
 }  // namespace pydocs
 #endif /* STACKSEARCH_DOCS */