trajectory/fit_trajectory.py

"""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)