Create trajectory/fit_trajectory.py

trajectory/fit_trajectory.py

@@ -0,0 +1,22 @@
+"""Polynomial fitting & projection to stumps."""
+
+import numpy as np
+from typing import List, Tuple
+
+
+class TrajectoryFitter:
+    def __init__(self, degree: int = 2):
+        self.deg = degree
+        self.coeffs = None
+
+    def fit(self, points: List[Tuple[float, float]]):
+        """Fit y = f(x) polynomial to ball path."""
+        pts = np.array(points)
+        x, y = pts[:, 0], pts[:, 1]
+        self.coeffs = np.polyfit(x, y, self.deg)
+
+    def project(self, x_vals: np.ndarray):
+        if self.coeffs is None:
+            raise RuntimeError("Call fit() before project()")
+        poly = np.poly1d(self.coeffs)
+        return poly(x_vals)