Bezier S-curve math

# Trying to understand the math involved in
#   https://github.com/synthetos/TinyG/blob/master/firmware/tinyg/plan_exec.c
#   https://github.com/synthetos/g2/blob/master/g2core/plan_exec.cpp#L693
#   https://github.com/synthetos/g2/blob/master/g2core/planner.h
#   https://github.com/synthetos/g2/blob/92e455e965e2a68741218ca9b10162b70e45ffa2/g2core/plan_zoid.cpp

import sympy as sp
import sympy.abc as abc


def get_bernstein_basis_polynomial_coefficient(i, n, t):
    return sp.simplify(sp.binomial(n, i) * t**i * (1-t)**(n-i))


def get_bernstein_basis_polynomial(n, x, coefficient_formater='b_{i}'):
    assert isinstance(x, sp.Symbol)
    return sp.simplify(sum([get_bernstein_basis_polynomial_coefficient(i, n, x)
                            * sp.Symbol(coefficient_formater.format(i=i)) for i in range(n+1)]))


class PopControl:
    NAMES = ['position', 'velocity', 'acceleration', 'jerk', 'snap', 'crackle', 'pop']

    @property
    def bezier_order(self):
        return len(self.NAMES) - 2  # We parameterize in velocity and oder is one lower for linear NAMES[-1]

    def derive_math(self):
        # We are using a quintic (fifth-degree) Bezier polynomial for the velocity curve.
        # This gives us a "linear pop" velocity curve; with pop being the sixth derivative of position:
        # velocity - 1st, acceleration - 2nd, jerk - 3rd, snap - 4th, crackle - 5th, pop - 6th

        # The Bezier curve takes the form:
        #
        # V(t) = P_0 * B_0(t) + P_1 * B_1(t) + P_2 * B_2(t) + P_3 * B_3(t) + P_4 * B_4(t) + P_5 * B_5(t)
        #
        # Where 0 <= t <= 1, and V(t) is the velocity. P_0 through P_5 are the control points, and B_0(t)
        # through B_5(t) are the Bernstein basis.

        bezier_d0 = get_bernstein_basis_polynomial(self.bezier_order, abc.t, coefficient_formater='P_{i}')
        bezier_poly = sp.Poly(bezier_d0, abc.t)

        # We use forward-differencing to calculate each position through the curve.
        #   This requires a formula of the form:
        #
        #       V_f(t) = A*t^5 + B*t^4 + C*t^3 + D*t^2 + E*t + F
        #
        # Looking at the above B_0(t) through B_5(t) expanded forms, if we take the coefficients of t^5
        # through t of the Bezier form of V(t), we can determine that:

        # for label, coeff in zip('ABCDEF', bezier_poly.all_coeffs()):
        #     print(f'{label} = {coeff}')

        # A = -P_0 + 5*P_1 - 10*P_2 + 10*P_3 - 5*P_4 + P_5
        # B = 5*P_0 - 20*P_1 + 30*P_2 - 20*P_3 + 5*P_4
        # C = -10*P_0 + 30*P_1 - 30*P_2 + 10*P_3
        # D = 10*P_0 - 20*P_1 + 10*P_2
        # E = -5*P_0 + 5*P_1
        # F = P_0

        # Now, since we will (currently) *always* want the initial acceleration and jerk values to be 0.

        bezier_d1 = sp.diff(bezier_d0, abc.t)
        bezier_d2 = sp.diff(bezier_d1, abc.t)

        constraints = sp.solve([sp.Eq(bezier_d1.subs(abc.t, 0), 0),     # Initial acceleration = 0
                                sp.Eq(bezier_d1.subs(abc.t, 1), 0),     # Final acceleration = 0

                                sp.Eq(bezier_d2.subs(abc.t, 0), 0),     # Initial jerk = 0
                                sp.Eq(bezier_d2.subs(abc.t, 1), 0),     # Final jerk = 0
                                ], bezier_poly.free_symbols)

        # We set P_i = P_0 = P_1 = P_2 (initial velocity), and P_t = P_3 = P_4 = P_5 (target velocity)
        bezier_d0 = bezier_d0.subs(constraints)

        P_i, P_t = sp.Symbol('P_i'), sp.Symbol('P_t')
        bezier_d0 = bezier_d0.subs({bezier_d0.subs(abc.t, 0): P_i,
                                    bezier_d0.subs(abc.t, 1): P_t})

        bezier_poly = sp.Poly(bezier_d0, abc.t)
        bezier_d1 = sp.diff(bezier_d0, abc.t)
        bezier_d2 = sp.diff(bezier_d1, abc.t)

        #  which, after simplification, resolves to:

        # for label, coeff in zip('ABCDEF', bezier_poly.all_coeffs()):
        #     print(f'{label} = {coeff}')

        # A = -6*P_i + 6*P_t
        # B = 15*P_i - 15*P_t
        # C = -10*P_i + 10*P_t
        # D = 0
        # E = 0
        # F = P_i

        if True:
            # Let's visualize this:
            import matplotlib.pyplot as plt
            import numpy as np
            plt.figure()
            ts = np.linspace(0, 1, 200)
            param = {P_i: 0, P_t:-5}
            for order in range(self.bezier_order):
                expr = sp.diff(bezier_d0.subs(param), abc.t, order)
                fun = sp.lambdify(abc.t, expr, 'numpy')
                plt.plot(ts, fun(ts), label=self.NAMES[order+1])
            plt.legend()
            plt.show()



        # Note, that due to construction, jerk is always zero at 0, 1/2, 1 and hence acceleration in maximum at 1/2
        # print(sp.solve(bezier_d2, abc.t))

        # We can hence easily calculate the maximum acceleration
        a_max = sp.simplify(bezier_d1.subs(abc.t, sp.Rational(1, 2)))
        print(f'a_max: {a_max}')

        # Note, that due to construction, jerk is maximum at +/- sqrt(3)/6 + 1/2
        # print(sp.solve(sp.diff(bezier_d0, abc.t, 3), abc.t))
        j_max = sp.simplify(bezier_d2.subs(abc.t, sp.solve(sp.diff(bezier_d0, abc.t, 3), abc.t)[1]))
        print(f'j_max: {j_max}')



def main():
    pop_control = PopControl()
    pop_control.derive_math()

if __name__ == '__main__':
    main()