Spaces:
Sleeping
Sleeping
File size: 647 Bytes
a6a8a0d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
"""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)
|