Add helper functions to update trajectory statistics.

src/kbmod/trajectory_utils.py

@@ -221,3 +221,60 @@ def trajectory_to_yaml(trj):
         "valid": trj.valid,
     }
     return dump(yaml_dict)
+
+
+def update_trajectory_from_psi_phi(trj, psi_curve, phi_curve, index_valid=None, in_place=True):
+    """Update the trajectory's statistic information from a psi_curve and
+    phi_curve. Uses an optional index_valid mask (True/False) to mask out
+    pixels.
+    Parameters
+    ----------
+    trj : `Trajectory`
+        The trajectory to update.
+    psi_curve : `numpy.ndarray`
+        The float psi values at each time step.
+    phi_curve : `numpy.ndarray`
+        The float phi values at each time step.
+    index_valid : `numpy.ndarray`, optional
+        An array of Booleans indicating whether the time step is valid.
+    in_place : `bool`
+        Update the input trajectory in-place.
+    Returns
+    -------
+    result : `Trajectory`
+        The updated trajectory. May be the same as trj if in_place=True.
+    Raises
+    ------
+    Raises a ValueError if the input arrays are not the same size.
+    """
+    if len(psi_curve) != len(phi_curve):
+        raise ValueError("Mismatched psi and phi curve lengths.")
+
+    # Compute the sums of the (masked) arrays.
+    if index_valid is None:
+        psi_sum = np.sum(psi_curve)
+        phi_sum = np.sum(phi_curve)
+        num_obs = len(psi_curve)
+    else:
+        if len(psi_curve) != len(index_valid):
+            raise ValueError("Mismatched psi/phi curve and index_valid lengths.")
+        psi_sum = np.sum(psi_curve[index_valid])
+        phi_sum = np.sum(phi_curve[index_valid])
+        num_obs = len(psi_curve[index_valid])
+
+    # Create a copy of the trajectory if we are not modifying in-place.
+    if in_place:
+        result = trj
+    else:
+        result = make_trajectory(x=trj.x, y=trj.y, vx=trj.vx, vy=trj.vy)
+
+    # Update the statistics information (avoiding divide by zero).
+    if phi_sum <= 0.0:
+        result.lh = 0.0
+        result.flux = 0.0
+    else:
+        result.lh = psi_sum / np.sqrt(phi_sum)
+        result.flux = psi_sum / phi_sum
+    result.obs_count = num_obs
+
+    return result

tests/test_trajectory_utils.py

@@ -112,6 +112,38 @@ def test_trajectory_yaml(self):
         self.assertEqual(new_trj.lh, 6.0)
         self.assertEqual(new_trj.obs_count, 7)
 
+    def test_update_trajectory_from_psi_phi(self):
+        trj = make_trajectory(x=0, y=10, vx=-1.0, vy=2.0)
+        self.assertEqual(trj.lh, 0.0)
+        self.assertEqual(trj.flux, 0.0)
+        self.assertEqual(trj.obs_count, 0)
+
+        # Non-in-place update
+        psi = np.array([1.0, 1.1, 1.2, 1.3])
+        phi = np.array([1.0, 1.0, 0.0, 2.0])
+        trj2 = update_trajectory_from_psi_phi(trj, psi, phi, in_place=False)
+        self.assertEqual(trj2.obs_count, 4)
+        self.assertAlmostEqual(trj2.flux, 1.15)
+        self.assertAlmostEqual(trj2.lh, 2.3)
+
+        # Original Trajectory is unchanged.
+        self.assertEqual(trj.lh, 0.0)
+        self.assertEqual(trj.flux, 0.0)
+        self.assertEqual(trj.obs_count, 0)
+
+        # Check the original is modified in the in-place update.
+        trj3 = update_trajectory_from_psi_phi(trj, psi, phi, in_place=True)
+        self.assertEqual(trj.obs_count, 4)
+        self.assertAlmostEqual(trj.flux, 1.15)
+        self.assertAlmostEqual(trj.lh, 2.3)
+
+        # Mark index 1 invalid.
+        index_valid = np.array([True, False, True, True])
+        trj4 = update_trajectory_from_psi_phi(trj, psi, phi, index_valid=index_valid, in_place=False)
+        self.assertEqual(trj4.obs_count, 3)
+        self.assertAlmostEqual(trj4.flux, 1.1666667, delta=1e-5)
+        self.assertAlmostEqual(trj4.lh, 2.020725, delta=1e-5)
+
 
 if __name__ == "__main__":
     unittest.main()