Init

rotorpy/trajectories/polynomial_traj.py

@@ -0,0 +1,107 @@
+import numpy as np
+
+class Polynomial(object):
+    """
+
+    """
+    def __init__(self, points, v_avg=1.2):
+        """
+        Inputs:
+            points, (N, 3) array of N waypoint coordinates in 3D
+            v_avg, the average speed between segments
+        """
+
+        def get_poly(xi, xf, T):
+            """
+            Return fully constrained polynomial coefficients from xi to xf in
+            time interval [0,T]. Low derivatives are all zero at the endpoints.
+            """
+            A = np.array([[      0,       0,      0,    0, 0, 1],
+                          [   T**5,    T**4,   T**3, T**2, T, 1],
+                          [      0,       0,      0,    0, 1, 0],
+                          [ 5*T**4,  4*T**3, 3*T**2,  2*T, 1, 0],
+                          [      0,       0,      0,    2, 0, 0],
+                          [20*T**3, 12*T**2,    6*T,    2, 0, 0]])
+            b = np.array([xi, xf, 0, 0, 0, 0])
+            poly = np.linalg.solve(A,b)
+            return poly
+
+        self.v_avg = v_avg
+
+        # Remove any sequential duplicate points; always keep the first point.
+        seg_dist = np.linalg.norm(np.diff(points, axis=0), axis=1)
+        seg_mask = np.append(True, seg_dist > 1e-3)
+        points = points[seg_mask,:]
+
+        # If at least two waypoints remain, calculate segment polynomials.
+        if points.shape[0] >= 2:
+            N = points.shape[0]-1
+            self.T = seg_dist/self.v_avg # segment duration
+            self.x_poly = np.zeros((N, 3, 6))
+            for i in range(N):
+                for j in range(3):
+                    self.x_poly[i,j,:] = get_poly(points[i,j], points[i+1,j], self.T[i])
+        # Otherwise, hard code constant polynomial at initial waypoint position.
+        else:
+            N = 1
+            self.T = np.zeros((N,))
+            self.x_poly = np.zeros((N, 3, 6))
+            self.x_poly[0,:,-1] = points[0,:]
+
+        # Calculate global start time of each segment.
+        self.t_start = np.concatenate(([0], np.cumsum(self.T[:-1])))
+
+        # Calculate derivative polynomials.
+        self.x_dot_poly    = np.zeros((N, 3, 5))
+        self.x_ddot_poly   = np.zeros((N, 3, 4))
+        self.x_dddot_poly  = np.zeros((N, 3, 3))
+        self.x_ddddot_poly = np.zeros((N, 3, 2))
+        for i in range(N):
+            for j in range(3):
+                self.x_dot_poly[i,j,:]    = np.polyder(self.x_poly[i,j,:], m=1)
+                self.x_ddot_poly[i,j,:]   = np.polyder(self.x_poly[i,j,:], m=2)
+                self.x_dddot_poly[i,j,:]  = np.polyder(self.x_poly[i,j,:], m=3)
+                self.x_ddddot_poly[i,j,:] = np.polyder(self.x_poly[i,j,:], m=4)
+
+    def update(self, t):
+        """
+        Given the present time, return the desired flat output and derivatives.
+
+        Inputs
+            t, time, s
+        Outputs
+            flat_output, a dict describing the present desired flat outputs with keys
+                x,        position, m
+                x_dot,    velocity, m/s
+                x_ddot,   acceleration, m/s**2
+                x_dddot,  jerk, m/s**3
+                x_ddddot, snap, m/s**4
+                yaw,      yaw angle, rad
+                yaw_dot,  yaw rate, rad/s
+        """
+        x        = np.zeros((3,))
+        x_dot    = np.zeros((3,))
+        x_ddot   = np.zeros((3,))
+        x_dddot  = np.zeros((3,))
+        x_ddddot = np.zeros((3,))
+        yaw = 0
+        yaw_dot = 0
+
+        # Find interval index i and time within interval t.
+        t = np.clip(t, self.t_start[0], self.t_start[-1]+self.T[-1])
+        for i in range(self.t_start.size):
+            if self.t_start[i] + self.T[i] >= t:
+                break
+        t = t - self.t_start[i]
+
+        # Evaluate polynomial.
+        for j in range(3):
+            x[j]        = np.polyval(       self.x_poly[i,j,:], t)
+            x_dot[j]    = np.polyval(   self.x_dot_poly[i,j,:], t)
+            x_ddot[j]   = np.polyval(  self.x_ddot_poly[i,j,:], t)
+            x_dddot[j]  = np.polyval( self.x_dddot_poly[i,j,:], t)
+            x_ddddot[j] = np.polyval(self.x_ddddot_poly[i,j,:], t)
+
+        flat_output = { 'x':x, 'x_dot':x_dot, 'x_ddot':x_ddot, 'x_dddot':x_dddot, 'x_ddddot':x_ddddot,
+                        'yaw':yaw, 'yaw_dot':yaw_dot}
+        return flat_output