frc971/control_loops/python/libspline.py

#!/usr/bin/python3
"""Wrapper around spline.h/cc through spline_array.cc."""

__author__ = 'Alex Perry (alex.perry96@gmail.com)'

import ctypes as ct
import numpy as np
import os

libSpline = None
for path in os.environ.get('PYTHONPATH').split(':'):
    try:
        libSpline = ct.cdll.LoadLibrary(
            os.path.join(path, 'frc971/control_loops/drivetrain/spline.so'))
    except (OSError):
        pass

# Define required output types.
libSpline.NewSpline.restype = ct.c_void_p
libSpline.SplineTheta.restype = ct.c_double
libSpline.SplineDTheta.restype = ct.c_double
libSpline.SplineDDTheta.restype = ct.c_double
libSpline.NewDistanceSpline.restype = ct.c_void_p
libSpline.DistanceSplineTheta.restype = ct.c_double
libSpline.DistanceSplineDTheta.restype = ct.c_double
libSpline.DistanceSplineDThetaDt.restype = ct.c_double
libSpline.DistanceSplineDDTheta.restype = ct.c_double
libSpline.DistanceSplineLength.restype = ct.c_double
libSpline.NewTrajectory.restype = ct.c_void_p
libSpline.TrajectoryLength.restype = ct.c_double
libSpline.TrajectoryDistance.restype = ct.c_double
libSpline.TrajectoryGetPlanXVAPtr.restype = ct.c_void_p

# Required for trajectory
libSpline.SetUpLogging()


class Spline:
    """
    A wrapper around spline.h/cc through libspline.cc.
    The functions return values parameterized by alpha, a number that varies
    between 0 and 1 along the length of the spline.
    """

    def __init__(self, points):
        assert points.shape == (2, 6)
        self.__points = points
        self.__spline = ct.c_void_p(
            libSpline.NewSpline(np.ctypeslib.as_ctypes(self.__points[0]),
                                np.ctypeslib.as_ctypes(self.__points[1])))

    def __del__(self):
        libSpline.deleteSpline(self.__spline)

    def setPoint(self, index, x, y):
        self.__points[0, index] = x
        self.__points[1, index] = y
        libSpline.deleteSpline(self.__spline)
        self.__spline = ct.c_void_p(
            libSpline.newSpline(np.ctypeslib.as_ctypes(self.__points[0]),
                                np.ctypeslib.as_ctypes(self.__points[1])))

    def Point(self, alpha):
        result = np.zeros(2)
        libSpline.SplinePoint(self.__spline, ct.c_double(alpha),
                              np.ctypeslib.as_ctypes(result))
        return result

    def DPoint(self, alpha):
        result = np.zeros(2)
        libSpline.SplineDPoint(self.__spline, ct.c_double(alpha),
                               np.ctypeslib.as_ctypes(result))
        return result

    def DDPoint(self, alpha):
        result = np.zeros(2)
        libSpline.SplineDDPoint(self.__spline, ct.c_double(alpha),
                                np.ctypeslib.as_ctypes(result))
        return result

    def DDDPoint(self, alpha):
        result = np.zeros(2)
        libSpline.SplineDDDPoint(self.__spline, ct.c_double(alpha),
                                 np.ctypeslib.as_ctypes(result))
        return result

    def Theta(self, alpha):
        return libSpline.SplineTheta(self.__spline, ct.c_double(alpha))

    def DTheta(self, alpha):
        return libSpline.SplineDTheta(self.__spline, ct.c_double(alpha))

    def DDTheta(self, alpha):
        return libSpline.SplineDDTheta(self.__spline, ct.c_double(alpha))

    def ControlPoints(self):
        return self.__points

    def GetSplinePtr(self):
        return self.__spline


class DistanceSpline:
    """
    A wrapper around distance_spline.h/cc through libspline.cc.
    The functions return values parameterized by the distance along the spline,
    starting at 0 and and ending at the value returned by Length().
    """

    def __init__(self, splines):
        self.__spline = None
        spline_ptrs = []
        for spline in splines:
            spline_ptrs.append(spline.GetSplinePtr().value)
        spline_ptrs = np.array(spline_ptrs)

        spline_array = np.ctypeslib.as_ctypes(spline_ptrs)
        self.__spline = ct.c_void_p(
            libSpline.NewDistanceSpline(ct.byref(spline_array), len(splines)))

    def __del__(self):
        libSpline.deleteDistanceSpline(self.__spline)

    def XY(self, distance):
        result = np.zeros(2)
        libSpline.DistanceSplineXY(self.__spline, ct.c_double(distance),
                                   np.ctypeslib.as_ctypes(result))
        return result

    def DXY(self, distance):
        result = np.zeros(2)
        libSpline.DistanceSplineDXY(self.__spline, ct.c_double(distance),
                                    np.ctypeslib.as_ctypes(result))
        return result

    def DDXY(self, distance):
        result = np.zeros(2)
        libSpline.DistanceSplineDDXY(self.__spline, ct.c_double(distance),
                                     np.ctypeslib.as_ctypes(result))
        return result

    def Theta(self, distance):
        return libSpline.DistanceSplineTheta(self.__spline,
                                             ct.c_double(distance))

    def DTheta(self, distance):
        return libSpline.DistanceSplineDTheta(self.__spline,
                                              ct.c_double(distance))

    def DThetaDt(self, distance, velocity):
        return libSpline.DistanceSplineDThetaDt(self.__spline,
                                                ct.c_double(distance),
                                                ct.c_double(velocity))

    def DDTheta(self, distance):
        return libSpline.DistanceSplineDDTheta(self.__spline,
                                               ct.c_double(distance))

    def Length(self):
        return libSpline.DistanceSplineLength(self.__spline)

    def GetSplinePtr(self):
        return self.__spline


class Trajectory:
    """A wrapper around trajectory.h/cc through libspline.cc."""

    def __init__(self, distance_spline, vmax=10, num_distance=0):
        self.__trajectory = ct.c_void_p(
            libSpline.NewTrajectory(distance_spline.GetSplinePtr(),
                                    ct.c_double(vmax), num_distance))

    def __del__(self):
        libSpline.deleteTrajectory(self.__trajectory)

    def SetLongitudinalAcceleration(self, accel):
        libSpline.TrajectorySetLongitudinalAcceleration(
            self.__trajectory, ct.c_double(accel))

    def SetLateralAcceleration(self, accel):
        libSpline.TrajectorySetLateralAcceleration(self.__trajectory,
                                                   ct.c_double(accel))

    def SetVoltageLimit(self, limit):
        libSpline.TrajectorySetVoltageLimit(self.__trajectory,
                                            ct.c_double(limit))

    def LimitVelocity(self, start_distance, end_distance, max_vel):
        libSpline.TrajectoryLimitVelocity(self.__trajectory,
                                          ct.c_double(start_distance),
                                          ct.c_double(end_distance),
                                          ct.c_double(max_vel))

    def Plan(self):
        """
        Call this to compute the plan, if any of the limits change, a new plan
        must be generated.
        """
        libSpline.TrajectoryPlan(self.__trajectory)

    def Voltage(self, distance):
        """
        Returns a pair of voltages for a given distance.
        Order is left-right.
        """
        result = np.zeros(2)
        libSpline.TrajectoryVoltage(self.__trajectory, ct.c_double(distance),
                                    np.ctypeslib.as_ctypes(result))
        return result

    def Length(self):
        return libSpline.TrajectoryLength(self.__trajectory)

    def Distance(self, index):
        return libSpline.TrajectoryDistance(self.__trajectory, index)

    def Distances(self):
        """
        Returns an array of distances used to compute the plan. The linear
        distance between each distance is equal.
        """
        path_length = libSpline.TrajectoryGetPathLength(self.__trajectory)
        distances = numpy.zeros(path_length)
        libSpline.TrajectoryDistances(self.__trajectory,
                                      np.ctypeslib.as_ctypes(distances))
        return distances

    def GetPlan(self):
        """
        Returns the plan as an array of velocities matched to each distance
        from Distances.
        """
        path_length = libSpline.TrajectoryGetPathLength(self.__trajectory)
        velocities = numpy.zeros(path_length)
        libSpline.TrajectoryGetPlan(self.__trajectory,
                                    np.ctypeslib.as_ctypes(velocities))
        return velocities

    def GetPlanXVA(self, dt):
        """
        dt is in seconds
        Returns the position, velocity, and acceleration as a function of time.
        This is returned as a 3xN numpy array where N is proportional to how
        long it takes to run the path.
        This is slow so don't call more than once with the same data.
        """
        XVAPtr = ct.c_void_p(
            libSpline.TrajectoryGetPlanXVAPtr(self.__trajectory,
                                              int(dt * 1e9)))
        XVALength = libSpline.TrajectoryGetVectorLength(XVAPtr)
        if XVALength == 0:
            libSpline.TrajectoryDeleteVector(XVAPtr)
            return None
        X = np.zeros(XVALength)
        V = np.zeros(XVALength)
        A = np.zeros(XVALength)
        libSpline.TrajectoryGetPlanXVA(XVAPtr, np.ctypeslib.as_ctypes(X),
                                       np.ctypeslib.as_ctypes(V),
                                       np.ctypeslib.as_ctypes(A))
        libSpline.TrajectoryDeleteVector(XVAPtr)
        return np.vstack([X, V, A])