From 6dadfd164bbba160df0c75a877413b6e40567bbf Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 15 Nov 2024 14:35:15 -0500
Subject: [PATCH] Add a pencil search TrajectoryGenerator

---
 src/kbmod/trajectory_generator.py  | 82 ++++++++++++++++++++++++++++++
 tests/test_trajectory_generator.py | 24 +++++++++
 2 files changed, 106 insertions(+)

diff --git a/src/kbmod/trajectory_generator.py b/src/kbmod/trajectory_generator.py
index 7dc012f46..b57632daf 100644
--- a/src/kbmod/trajectory_generator.py
+++ b/src/kbmod/trajectory_generator.py
@@ -2,6 +2,7 @@
 import copy
 import logging
 import math
+import numpy as np
 import random
 
 import astropy.units as u
@@ -234,6 +235,87 @@ def generate(self, *args, **kwargs):
                 yield Trajectory(vx=vx, vy=vy)
 
 
+class PencilSearch(TrajectoryGenerator):
+    """A search in a small cone around a given velocity.
+
+    The search varies the given velocity's angle and magnitude.
+    The angle includes the range: original angle +/- max_ang_offset.
+    The velcoty magnitude includes the range: original magnitude +/- max_vel_offset
+
+    Parameters
+    ----------
+    vx : `float`
+        The center velocity in the x dimension (in pixels per day)
+    vy : `float`
+        The center velocity in the y dimension (in pixels per day)
+    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)
+    """
+
+    def __init__(
+        self,
+        vx,
+        vy,
+        max_ang_offset=0.2618,
+        ang_step=0.035,
+        max_vel_offset=10.0,
+        vel_step=0.5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if ang_step <= 0.0 or vel_step <= 0.0 or max_ang_offset <= 0.0 or max_vel_offset <= 0.0:
+            raise ValueError("Invalid parameters. All ranges and step sizes must be > 0.0.")
+
+        self.center_vx = vx
+        self.center_vy = vy
+
+        self.center_ang = np.arctan2(vy, vx)
+        self.min_ang = self.center_ang - max_ang_offset
+        self.max_ang = self.center_ang + max_ang_offset
+        self.ang_step = ang_step
+
+        self.center_vel = np.sqrt(vx * vx + vy * vy)
+        self.min_vel = np.max([self.center_vel - max_vel_offset, 0.0])
+        self.max_vel = self.center_vel + max_vel_offset
+        self.vel_step = vel_step
+
+    def __repr__(self):
+        return (
+            "PencilSearch:"
+            f" v=[{self.min_vel}, {self.max_vel}), {self.vel_step}"
+            f" a=[{self.min_ang}, {self.max_ang}), {self.ang_step}"
+        )
+
+    def __str__(self):
+        return (
+            "PencilSearch:\n"
+            f"    Vel: [{self.min_vel}, {self.max_vel}) in {self.vel_step} sized steps.\n"
+            f"    Ang: [{self.min_ang}, {self.max_ang}) in {self.ang_step} sized steps."
+        )
+
+    def generate(self, *args, **kwargs):
+        """Produces a single candidate trajectory to test.
+
+        Returns
+        -------
+        candidate : `Trajectory`
+            A ``Trajectory`` to test at each pixel.
+        """
+        for ang in np.arange(self.min_ang, self.max_ang + 1e-8, self.ang_step):
+            for vel in np.arange(self.min_vel, self.max_vel + 1e-8, self.vel_step):
+                vx = np.cos(ang) * vel
+                vy = np.sin(ang) * vel
+
+                yield Trajectory(vx=vx, vy=vy)
+
+
 class KBMODV1Search(TrajectoryGenerator):
     """Search a grid defined by velocities and angles."""
 
diff --git a/tests/test_trajectory_generator.py b/tests/test_trajectory_generator.py
index f08fbd9e1..7a26a1897 100644
--- a/tests/test_trajectory_generator.py
+++ b/tests/test_trajectory_generator.py
@@ -9,6 +9,7 @@
     KBMODV1Search,
     KBMODV1SearchConfig,
     EclipticCenteredSearch,
+    PencilSearch,
     SingleVelocitySearch,
     RandomVelocitySearch,
     VelocityGridSearch,
@@ -49,6 +50,29 @@ def test_VelocityGridSearch(self):
         self.assertRaises(ValueError, VelocityGridSearch, 3, 0.0, 2.0, 3, 0.25, -0.25)
         self.assertRaises(ValueError, VelocityGridSearch, 3, 2.0, 0.0, 3, -0.25, 0.25)
 
+    def test_PencilSearch(self):
+        gen = PencilSearch(
+            10.0,
+            20.0,
+            max_ang_offset=0.1,
+            ang_step=0.05,
+            max_vel_offset=5.0,
+            vel_step=2.5,
+        )
+
+        trjs = [trj for trj in gen]
+        self.assertEqual(len(trjs), 25)
+
+        expected_angs = np.arctan2(20.0, 10.0) + np.array([-0.1, -0.05, 0.0, 0.05, 0.1])
+        expected_vels = np.sqrt(500.0) + np.array([-5.0, -2.5, 0.0, 2.5, 5.0])
+        for a_i in range(5):
+            for v_i in range(5):
+                trj = trjs[5 * a_i + v_i]
+                ang = np.arctan2(trj.vy, trj.vx)
+                vel = np.sqrt(trj.vx * trj.vx + trj.vy * trj.vy)
+                self.assertAlmostEqual(ang, expected_angs[a_i], delta=1e-5)
+                self.assertAlmostEqual(vel, expected_vels[v_i], delta=1e-5)
+
     def test_KBMODV1Search(self):
         # Note that KBMOD v1's legacy search will never include the upper bound of angle or velocity.
         gen = KBMODV1Search(3, 0.0, 3.0, 2, -0.25, 0.25)