Blending of motion parabolas

blend.py

import numpy as np
import matplotlib.pyplot as plt
import dataclasses


@dataclasses.dataclass
class MotionEstimate:
    coeffs: np.ndarray  # 3x1
    t0: float
    degree: int = dataclasses.field(init=False)

    def __post_init__(self):
        self.degree = len(self.coeffs) - 1

    def __call__(self, t: np.ndarray):
        T = np.vander(t - self.t0, self.degree + 1)  # Nx3
        return (T @ self.coeffs).reshape(-1)

    def at(self, t: float):
        t = t - self.t0
        return np.dot(self._v(t), self.coeffs)

    def dx_at(self, t: float):
        t = t - self.t0
        return np.dot(self._vprime(t), self.coeffs[:-1])

    def _v(self, t: float):
        return np.array([t ** i for i in reversed(range(self.degree + 1))])

    def _vprime(self, t: float):
        return np.array([i * t ** (i - 1) for i in reversed(range(1, self.degree + 1))])


def blend(
    m1: MotionEstimate, m2: MotionEstimate, tnow: float, h: float
) -> MotionEstimate:
    A = np.zeros((4, 4))
    b = np.zeros(4)

    # Position at beginning (tnow) should match
    A[0, 0] = 0
    A[0, 1] = 0
    A[0, 2] = 0
    A[0, 3] = 1
    b[0] = m1.at(tnow)

    # Position at end of horizon should match
    A[1, 0] = h ** 3
    A[1, 1] = h ** 2
    A[1, 2] = h
    A[1, 3] = 1
    b[1] = m2.at(tnow + h)

    # at beginning and end

    A[2, 0] = 0
    A[2, 1] = 0
    A[2, 2] = 1
    A[2, 3] = 0
    b[2] = m1.dx_at(tnow)

    A[3, 0] = 3 * h ** 2
    A[3, 1] = 2 * h
    A[3, 2] = 1
    A[3, 3] = 0
    b[3] = m2.dx_at(tnow + h)

    coeffs = np.linalg.solve(A, b)
    return MotionEstimate(coeffs, tnow)


class BlendedMotionEstimate:
    def __init__(
        self, m1: MotionEstimate, m2: MotionEstimate, tnow: float, h: float
    ) -> None:
        self.b = blend(m1, m2, tnow, h)
        self.m1 = m1
        self.m2 = m2
        self.h = h
        self.tnow = tnow
        self.t0 = self.b.t0
        self.degree = self.b.degree

    def __call__(self, t: np.ndarray):
        x = self.b(t)
        mask = t < self.b.t0
        x[mask] = self.m1(t[mask])
        mask = t > self.tnow + self.h
        x[mask] = self.m2(t[mask])
        return x

    def at(self, t: float):
        return self._estimator(t).at(t)

    def dx_at(self, t: float):
        return self._estimator(t).dx_at(t)

    def _estimator(self, t: float) -> "MotionEstimate":
        if t < self.b.t0:
            return self.m1
        elif t > self.tnow + self.h:
            return self.m2
        else:
            return self.b


def main():
    m1 = MotionEstimate(coeffs=[-0.8, 1.0, 0.5], t0=0)
    m2 = MotionEstimate(coeffs=[0, 3.0, 5.0], t0=1.0)
    m3 = MotionEstimate(coeffs=[1.2, 5.0, 7.0], t0=3.0)

    tnow = 2.5
    h = 1.0
    mb1 = BlendedMotionEstimate(m1, m2, tnow, h)

    tnow = 3.5
    h = 1.0
    mb2 = BlendedMotionEstimate(mb1, m3, tnow, h)
    print(mb2)

    t = np.linspace(0, 10, 100)
    fig, ax = plt.subplots()
    ax.plot(t[t >= m1.t0], m1(t[t >= m1.t0]), label="m1")
    ax.plot(t[t >= m2.t0], m2(t[t >= m2.t0]), label="m2")
    ax.plot(t[t >= m3.t0], m3(t[t >= m3.t0]), label="m3")
    ax.plot(t[t >= mb1.t0], mb1(t[t >= mb1.t0]), label="b1")
    ax.plot(t[t >= mb2.t0], mb2(t[t >= mb2.t0]), label="b2")
    ax.axvline(tnow, linestyle="--")
    ax.axvline(tnow + h, linestyle="--")
    plt.legend()
    plt.show()

    pass


if __name__ == "__main__":
    main()