Merge branch 'main' into output_changes

docs/source/user_manual/search_params.rst

@@ -8,14 +8,6 @@ This document serves to provide a quick overview of the existing parameters and
 +------------------------+-----------------------------+----------------------------------------+
 | **Parameter**          | **Default Value**           | **Interpretation**                     |
 +------------------------+-----------------------------+----------------------------------------+
-| ``ang_arr``            | [np.pi/15, np.pi/15, 128]   | Minimum, maximum and number of angles  |
-|                        |                             | to search through (in radians)         |
-+------------------------+-----------------------------+----------------------------------------+
-| ``average_angle``      | None                        | Overrides the ecliptic angle           |
-|                        |                             | calculation and instead centers the    |
-|                        |                             | average search around average_angle    |
-|                        |                             | (in radians).                          |
-+------------------------+-----------------------------+----------------------------------------+
 | ``center_thresh``      | 0.00                        | The minimum fraction of total flux     |
 |                        |                             | within a stamp that must be contained  |
 |                        |                             | in the central pixel                   |
@@ -150,11 +142,6 @@ This document serves to provide a quick overview of the existing parameters and
 |                        |                             | the filtered trajectories. Warning     |
 |                        |                             | can use a lot of memory.               |
 +------------------------+-----------------------------+----------------------------------------+
-| ``v_arr``              | [92.0, 526.0, 256]          | Minimum, maximum and number of         |
-|                        |                             | velocities to search through.  The     |
-|                        |                             | minimum and maximum velocities are     |
-|                        |                             | specified in pixels per day.           |
-+------------------------+-----------------------------+----------------------------------------+
 | ``x_pixel_bounds``     | None                        | A length two list giving the starting  |
 |                        |                             | and ending x pixels to use for the     |
 |                        |                             | search. `None` uses the image bounds.  |

docs/source/user_manual/search_space.rst

@@ -33,7 +33,8 @@ Choosing Velocities
 Perhaps the most complex aspect of the KBMOD algorithm is how it defines the grid of search velocities. KBMOD allows you to define custom search strategies to best match the data. These include:
 * ``SingleVelocitySearch`` - A single predefined x and y velocity
 * ``VelocityGridSearch`` - An evenly spaced grid of x and y velocities
-* ``KBMODV1SearchConfig`` - An evenly spaced grid of velocity magnitudes and angles (this was the only option in v1.1 and before)
+* ``EclipticCenteredSearch`` - An evenly spaced grid of velocity magnitudes and angles (using a current parameterization) centered on a given or computed ecliptic angle.
+* ``KBMODV1SearchConfig`` - An evenly spaced grid of velocity magnitudes and angles (using the legacy parameterization).
 * ``RandomVelocitySearch`` - Randomly sampled x and y velocities
 Additional search strategies can be defined by overriding the ``TrajectoryGenerator`` class in trajectory_generator.py.
 
@@ -77,6 +78,52 @@ The ``VelocityGridSearch`` strategy searches a uniform grid of x and y velocitie
 | ``max_vy``             | The maximum velocity in the y-dimension (pixels per day). |
 +------------------------+-----------------------------------------------------------+
 
+EclipticCenteredSearch
+----------------------
+
+The grid is defined by two sets of parameters: a sampling of absolute velocities in pixels per day and a sampling of the velocities' angles in degrees or radians. Each sampling consists of values defining the range and number of sampling steps. 
+
+Given the linear sampling for both velocities and angles, the full set of candidate trajectories is computed as::
+
+
+    for (int a = 0; a < angleSteps; ++a) {
+        for (int v = 0; v < velocitySteps; ++v) {
+            searchList[a * velocitySteps + v].xVel = cos(sampled_angles[a]) * sampled_velocities[v];
+            searchList[a * velocitySteps + v].yVel = sin(sampled_angles[a]) * sampled_velocities[v];
+        }
+    }
+
+where ``sampled_angles`` contains the list of angles to test and ``sampled_velocities`` contains the list of velocities. 
+
+The list of velocities is created from the given bounds list ``velocities=[min_vel, max_vel, vel_steps]``. The range is inclusive of both bounds.
+
+Each angle in the list is computed as an **offset** from the ecliptic angle. KBMOD uses the following ordering for extracting the ecliptic.
+1. If ``given_ecliptic`` is provided (is not ``None``) in the generator’s configuration that value is used directly.
+2. If the first image has a WCS, the ecliptic is estimated from that WCS.
+3. A default ecliptic of 0.0 is used.
+The angles used are defined from the list ``angles=[min_offset, max_offset, angle_steps]`` and will span ``[ecliptic + min_offset, ecliptic + max_offset]`` inclusive of both bounds. Angles can be specified in degrees or radians (as noted by the ``angle_units`` parameter) but must be consistent among all angles.
+
+
++------------------------+------------------------------------------------------+
+| **Parameter**          | **Interpretation**                                   |
++------------------------+------------------------------------------------------+
+| ``angles``             | A length 3 list with the minimum angle offset,       |
+|                        | the maximum offset, and the number of angles to      |
+|                        | to search through (angles specified in units given   |
+|                        | by ``angle_units``).                                 |
++------------------------+------------------------------------------------------+
+| ``angle_units``        | The units to use for angles, such as "rad" or "deg". |
++------------------------+------------------------------------------------------+
+| ``given_ecliptic``     | The given value of the ecliptic angle (specified in  |
+|                        | units given by ``angle_units``).                     |
++------------------------+------------------------------------------------------+
+| ``velocities``         | A length 3 list with the minimum velocity (in        |
+|                        | pixels per day), the maximum velocity (in pixels     |
+|                        | per day), and number of velocities to test.          |
++------------------------+------------------------------------------------------+
+| ``velocity_units``     | The units to use for velocities (e.g. "pix / d")     |
++------------------------+------------------------------------------------------+
+
 
 KBMODV1SearchConfig
 -------------------

notebooks/KBMOD_Demo.ipynb

@@ -122,9 +122,15 @@
     "ang_arr = [ang_below, ang_above, ang_steps]\n",
     "\n",
     "input_parameters = {\n",
-    "    # Grid search parameters\n",
-    "    \"v_arr\": v_arr,\n",
-    "    \"ang_arr\": ang_arr,\n",
+    "    # Use search parameters (including a force ecliptic angle of 0.0)\n",
+    "    # to match what we know is in the demo data.\n",
+    "    \"generator_config\": {\n",
+    "        \"name\": \"EclipticCenteredSearch\",\n",
+    "        \"angles\": [-0.5, 0.5, 11],\n",
+    "        \"velocities\": [0.0, 20.0, 21],\n",
+    "        \"angle_units\": \"radian\",\n",
+    "        \"force_ecliptic\": 0.0,\n",
+    "    },\n",
     "    # Output parameters\n",
     "    \"res_filepath\": res_filepath,\n",
     "    \"output_suffix\": results_suffix,\n",
@@ -142,9 +148,6 @@
     "    # Some basic stamp filtering limits.\n",
     "    \"mom_lims\": [37.5, 37.5, 1.5, 1.0, 1.0],\n",
     "    \"peak_offset\": [3.0, 3.0],\n",
-    "    # Override the ecliptic angle for the demo data since we\n",
-    "    # know the true angle in pixel space.\n",
-    "    \"average_angle\": 0.0,\n",
     "}\n",
     "config = SearchConfiguration.from_dict(input_parameters)"
    ]

notebooks/create_fake_data.ipynb

@@ -149,9 +149,13 @@
     "\n",
     "settings = {\n",
     "    # Override the search data to match the known object.\n",
-    "    \"average_angle\": 0.0,\n",
-    "    \"v_arr\": [0, 20, 20],\n",
-    "    \"ang_arr\": [0.5, 0.5, 10],\n",
+    "    \"generator_config\": {\n",
+    "        \"name\": \"EclipticCenteredSearch\",\n",
+    "        \"angles\": [-0.5, 0.5, 11],\n",
+    "        \"velocities\": [0.0, 20.0, 21],\n",
+    "        \"angle_units\": \"radian\",\n",
+    "        \"force_ecliptic\": 0.0,\n",
+    "    },\n",
     "    # Loosen the other filtering parameters.\n",
     "    \"clip_negative\": True,\n",
     "    \"sigmaG_lims\": [15, 60],\n",

src/kbmod/configuration.py

@@ -19,8 +19,6 @@ def __init__(self):
         self._required_params = set()
 
         self._params = {
-            "ang_arr": [math.pi / 15, math.pi / 15, 128],
-            "average_angle": None,
             "center_thresh": 0.00,
             "chunk_size": 500000,
             "clip_negative": False,
@@ -33,7 +31,14 @@ def __init__(self):
             "do_mask": True,
             "do_stamp_filter": True,
             "encode_num_bytes": -1,
-            "generator_config": None,
+            "generator_config": {
+                "name": "EclipticCenteredSearch",
+                "velocity": [92.0, 526.0, 257],
+                "angles": [-math.pi / 15, math.pi / 15, 129],
+                "angle_units": "radian",
+                "velocity_units": "pix / d",
+                "given_ecliptic": None,
+            },
             "gpu_filter": False,
             "im_filepath": None,
             "legacy_filename": None,
@@ -50,7 +55,6 @@ def __init__(self):
             "stamp_radius": 10,
             "stamp_type": "sum",
             "track_filtered": False,
-            "v_arr": [92.0, 526.0, 256],
             "x_pixel_bounds": None,
             "x_pixel_buffer": None,
             "y_pixel_bounds": None,

src/kbmod/fake_data/demo_helper.py

@@ -38,9 +38,13 @@ def make_demo_data(filename=None):
     # Create configuraiton settings that match the object inserted.
     settings = {
         # Override the search data to match the known object.
-        "average_angle": 0.0,
-        "v_arr": [0, 20, 20],
-        "ang_arr": [0.5, 0.5, 10],
+        "generator_config": {
+            "name": "EclipticCenteredSearch",
+            "velocities": [0, 20.0, 21],
+            "angles": [-0.5, 0.5, 11],
+            "angle_units": "radian",
+            "given_ecliptic": 0.0,
+        },
         # Loosen the other filtering parameters.
         "clip_negative": True,
         "sigmaG_lims": [15, 60],

src/kbmod/filters/sigma_g_filter.py

@@ -10,6 +10,7 @@
 from scipy.special import erfinv
 
 from kbmod.results import Results
+from kbmod.search import DebugTimer
 
 logger = logging.getLogger(__name__)
 
@@ -177,7 +178,9 @@ def apply_clipped_sigma_g(clipper, result_data):
         logger.info("SigmaG Clipping : skipping, nothing to filter.")
         return
 
+    filter_timer = DebugTimer("sigma-g filtering", logger)
     lh = result_data.compute_likelihood_curves(filter_obs=True, mask_value=np.nan)
     obs_valid = clipper.compute_clipped_sigma_g_matrix(lh)
     result_data.update_obs_valid(obs_valid)
+    filter_timer.stop()
     return

src/kbmod/filters/stamp_filters.py

@@ -188,13 +188,16 @@ def append_coadds(result_data, im_stack, coadd_types, radius, chunk_size=100_000
     """
     if radius <= 0:
         raise ValueError(f"Invalid stamp radius {radius}")
+    stamp_timer = DebugTimer("computing extra coadds", logger)
 
     params = StampParameters()
     params.radius = radius
     params.do_filtering = False
 
     # Loop through all the coadd types in the list, generating a corresponding stamp.
     for coadd_type in coadd_types:
+        logger.info(f"Adding coadd={coadd_type} for all results.")
+
         if coadd_type == "median":
             params.stamp_type = StampType.STAMP_MEDIAN
         elif coadd_type == "mean":
@@ -212,6 +215,7 @@ def append_coadds(result_data, im_stack, coadd_types, radius, chunk_size=100_000
             chunk_size=chunk_size,
             colname=f"coadd_{coadd_type}",
         )
+    stamp_timer.stop()
 
 
 def append_all_stamps(result_data, im_stack, stamp_radius):

src/kbmod/results.py

@@ -139,9 +139,8 @@ def from_trajectories(cls, trajectories, track_filtered=False):
         input_d = {}
         for col in cls.required_cols:
             input_d[col[0]] = []
-        valid_mask = []
 
-        # Add the valid trajectories to the table.
+        # Add the trajectories to the table.
         for trj in trajectories:
             input_d["x"].append(trj.x)
             input_d["y"].append(trj.y)
@@ -150,17 +149,14 @@ def from_trajectories(cls, trajectories, track_filtered=False):
             input_d["likelihood"].append(trj.lh)
             input_d["flux"].append(trj.flux)
             input_d["obs_count"].append(trj.obs_count)
-            valid_mask.append(trj.valid)
 
         # Check for any missing columns and fill in the default value.
         for col in cls.required_cols:
             if col[0] not in input_d:
-                input_d[col[0]] = [col[2]] * num_valid
-                invalid_d[col[0]] = [col[2]] * num_invalid
+                input_d[col[0]] = [col[2]] * len(trajectories)
 
         # Create the table and add the unfiltered (and filtered) results.
         results = Results(input_d, track_filtered=track_filtered)
-        results.filter_rows(np.array(valid_mask, dtype=bool), "invalid_trajectory")
         return results
 
     @classmethod

src/kbmod/run_search.py

@@ -11,8 +11,7 @@
 from .filters.stamp_filters import append_all_stamps, append_coadds, get_coadds_and_filter_results
 
 from .results import Results
-from .trajectory_generator import create_trajectory_generator, KBMODV1SearchConfig
-from .wcs_utils import calc_ecliptic_angle
+from .trajectory_generator import create_trajectory_generator
 from .work_unit import WorkUnit
 
 
@@ -200,7 +199,7 @@ def run_search(self, config, stack, trj_generator=None):
             The stack before the masks have been applied. Modified in-place.
         trj_generator : `TrajectoryGenerator`, optional
             The object to generate the candidate trajectories for each pixel.
-            If None uses the default KBMODv1 grid search
+            If None uses the default EclipticCenteredSearch
 
         Returns
         -------
@@ -219,7 +218,7 @@ def run_search(self, config, stack, trj_generator=None):
 
         # Perform the actual search.
         if trj_generator is None:
-            trj_generator = create_trajectory_generator(config)
+            trj_generator = create_trajectory_generator(config, work_unit=None)
         keep = self.do_gpu_search(config, stack, trj_generator)
 
         if config["do_stamp_filter"]:
@@ -279,14 +278,7 @@ def run_search_from_work_unit(self, work):
         keep : `Results`
             The results.
         """
-        # Set the average angle if it is not set.
-        if work.config["average_angle"] is None:
-            center_pixel = (work.im_stack.get_width() / 2, work.im_stack.get_height() / 2)
-            if work.get_wcs(0) is not None:
-                work.config.set("average_angle", calc_ecliptic_angle(work.get_wcs(0), center_pixel))
-            else:
-                logger.warning("Average angle not set and no WCS provided. Setting average_angle=0.0")
-                work.config.set("average_angle", 0.0)
+        trj_generator = create_trajectory_generator(work.config, work_unit=work)
 
         # Run the search.
         return self.run_search(work.config, work.im_stack)

src/kbmod/search/common.h

@@ -55,8 +55,6 @@ struct Trajectory {
     int y = 0;
     // Number of images summed
     int obs_count;
-    // Whether the trajectory is valid. Used for on-GPU filtering.
-    bool valid = true;
 
     // Get pixel positions from a zero-shifted time. Centered indicates whether
     // the prediction starts from the center of the pixel (which it does in the search)
@@ -79,14 +77,13 @@ struct Trajectory {
     const std::string to_string() const {
         return "lh: " + std::to_string(lh) + " flux: " + std::to_string(flux) + " x: " + std::to_string(x) +
                " y: " + std::to_string(y) + " vx: " + std::to_string(vx) + " vy: " + std::to_string(vy) +
-               " obs_count: " + std::to_string(obs_count) + " valid: " + std::to_string(valid);
+               " obs_count: " + std::to_string(obs_count);
     }
 
     // This is a hack to provide a constructor with non-default arguments in Python. If we include
     // the constructor as a method in the Trajectory struct CUDA will complain when creating new objects
     // because it cannot call out to a host function.
-    static Trajectory make_trajectory(int x, int y, float vx, float vy, float flux, float lh, int obs_count,
-                                      bool valid) {
+    static Trajectory make_trajectory(int x, int y, float vx, float vy, float flux, float lh, int obs_count) {
         Trajectory trj;
         trj.x = x;
         trj.y = y;
@@ -95,7 +92,6 @@ struct Trajectory {
         trj.flux = flux;
         trj.lh = lh;
         trj.obs_count = obs_count;
-        trj.valid = valid;
         return trj;
     }
 };
@@ -199,16 +195,14 @@ static void trajectory_bindings(py::module &m) {
 
     py::class_<tj>(m, "Trajectory", pydocs::DOC_Trajectory)
             .def(py::init(&tj::make_trajectory), py::arg("x") = 0, py::arg("y") = 0, py::arg("vx") = 0.0f,
-                 py::arg("vy") = 0.0f, py::arg("flux") = 0.0f, py::arg("lh") = 0.0f, py::arg("obs_count") = 0,
-                 py::arg("valid") = true)
+                 py::arg("vy") = 0.0f, py::arg("flux") = 0.0f, py::arg("lh") = 0.0f, py::arg("obs_count") = 0)
             .def_readwrite("vx", &tj::vx)
             .def_readwrite("vy", &tj::vy)
             .def_readwrite("lh", &tj::lh)
             .def_readwrite("flux", &tj::flux)
             .def_readwrite("x", &tj::x)
             .def_readwrite("y", &tj::y)
             .def_readwrite("obs_count", &tj::obs_count)
-            .def_readwrite("valid", &tj::valid)
             .def("get_x_pos", &tj::get_x_pos, py::arg("time"), py::arg("centered") = true,
                  pydocs::DOC_Trajectory_get_x_pos)
             .def("get_y_pos", &tj::get_y_pos, py::arg("time"), py::arg("centered") = true,
@@ -220,13 +214,13 @@ static void trajectory_bindings(py::module &m) {
             .def("__str__", &tj::to_string)
             .def(py::pickle(
                     [](const tj &p) {  // __getstate__
-                        return py::make_tuple(p.vx, p.vy, p.lh, p.flux, p.x, p.y, p.obs_count, p.valid);
+                        return py::make_tuple(p.vx, p.vy, p.lh, p.flux, p.x, p.y, p.obs_count);
                     },
                     [](py::tuple t) {  // __setstate__
                         if (t.size() != 8) throw std::runtime_error("Invalid state!");
                         tj trj = {t[0].cast<float>(), t[1].cast<float>(), t[2].cast<float>(),
                                   t[3].cast<float>(), t[4].cast<int>(),   t[5].cast<int>(),
-                                  t[6].cast<int>(),   t[7].cast<bool>()};
+                                  t[6].cast<int>()};
                         return trj;
                     }));
 }

src/kbmod/search/pydocs/common_docs.h

@@ -22,9 +22,6 @@ static const auto DOC_Trajectory = R"doc(
     flux : `float`
         Flux (accumulated?)
     obs_count : `int`
-        Number of observations trajectory was seen in.
-    valid : `bool`
-        Whether the trajectory is valid. Used for filtering.
   )doc";
 
 static const auto DOC_Trajectory_get_x_pos = R"doc(