Merge pull request #706 from dirac-institute/generator_config_cleanup

Cleanup trajectory generation configuration

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                   |
@@ -149,17 +141,12 @@ This document serves to provide a quick overview of the existing parameters and
 |                        |                             | if:                                    |
 |                        |                             | * ``sum`` - (default) Per pixel sum    |
 |                        |                             | * ``median`` - A per pixel median      |
-|                        |                             | * ``mean`` - A per pixel mean          |\
+|                        |                             | * ``mean`` - A per pixel mean          |
 +------------------------+-----------------------------+----------------------------------------+
 | ``track_filtered``     | False                       | A Boolean indicating whether to track  |
 |                        |                             | 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,
             "ind_output_files": True,
             "im_filepath": None,
@@ -52,7 +57,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/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"]:
@@ -288,14 +287,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)