compLIB/compLib/Motor.py

import atexit
from enum import IntEnum

from compLib.LogstashLogging import Logging
from compLib.MetricsLogging import MetricsLogging
from compLib.Spi import Spi, Register

MOTOR_COUNT = 4
MAX_MOTOR_SPEED = 65535
MOTOR_PERCENTAGE_MULT = MAX_MOTOR_SPEED / 100.0
MOTOR_CURVE = [0.0, 0.0, 426.5, 692.0, 842.5, 953.5, 1032.5, 1090.5, 1135.5, 1171.0, 1203.5, 1230.0, 1249.5, 1268.0, 1283.0, 1298.5, 1308.0, 1320.0, 1332.0, 1339.5, 1352.5]

class MotorMode(IntEnum):
    COAST = 0,
    FORWARD = 1,
    BACKWARD = 2,
    BREAK = 3

class Motor(object):
    """Class used to control the motors
    """

    @staticmethod
    def pwm(port: int, pwm: int, mode: MotorMode):
        """Set specified motor to a specific pwm value and motor mode

        :param port: Port, which the motor is connected to. 1-4 allowed
        :param pwm: Raw PWM value which the motor should be set to. From 0 to 2^16 - 1
        :param mode: Motor mode. See enum MotorMode for more info
        :raises: IndexError
        """
        if port <= 0 or port > MOTOR_COUNT:
            raise IndexError("Invalid Motor port specified!")

        if port == 1:
            Spi.write(Register.MOTOR_1_PWM_H, 2, pwm)
            Spi.write(Register.PWM_1_CTRL, 1, int(mode))
        elif port == 2:
            Spi.write(Register.MOTOR_2_PWM_H, 2, pwm)
            Spi.write(Register.PWM_2_CTRL, 1, int(mode))
        elif port == 3:
            Spi.write(Register.MOTOR_3_PWM_H, 2, pwm)
            Spi.write(Register.PWM_3_CTRL, 1, int(mode))
        elif port == 4:
            Spi.write(Register.MOTOR_4_PWM_H, 2, pwm)
            Spi.write(Register.PWM_4_CTRL, 1, int(mode))

    @staticmethod
    def power_raw(port: int, percent: float, log_metric = True):
        """Set specified motor to percentage power

        :param port: Port, which the motor is connected to. 1-4
        :param percent: Percentage of max speed. between -100 and 100
        :raises: IndexError
        """

        if port <= 0 or port > MOTOR_COUNT:
            raise IndexError("Invalid Motor port specified!")

        if percent < -100 or percent > 100:
            raise IndexError("Invalid Motor speed specified! Speed is between -100 and 100 percent!")

        if log_metric:
            MetricsLogging.put("MotorRaw", percent, port)

        mode = MotorMode.COAST
        if percent < 0:
            percent = abs(percent)
            mode = MotorMode.BACKWARD
        elif percent > 0:
            mode = MotorMode.FORWARD

        pwm = percent * MOTOR_PERCENTAGE_MULT

        Motor.pwm(port, int(pwm), mode)

    @staticmethod
    def power(port: int, percent: float):
        """Set specified motor to percentage power, percentage is linearized
        so that it's roughly proportional to the speed

        :param port: Port, which the motor is connected to. 1-4
        :param percent: Percentage of max speed. between -100 and 100
        :raises: IndexError
        """
        raw_power = 0
        if percent > 0:
            raw_power = Motor.__linearizePower(MOTOR_CURVE, percent)
        elif percent < 0:
            raw_power =  -Motor.__linearizePower(MOTOR_CURVE, -percent)

        MetricsLogging.put("Motor", percent, port)
        Motor.power_raw(port, raw_power, False)

    @staticmethod
    def all_off():
        """
        Turns of all motors
        """
        Logging.get_logger().debug(f"Motor.all_off")

        for i in range(1, MOTOR_COUNT + 1):
            Motor.active_break(i)

    @staticmethod
    def active_break(port: int):
        """
        Actively break with a specific motor

        :param port: Port, which the motor is connected to. 1-4
        """
        Motor.pwm(port, 0, MotorMode.BREAK)

    @staticmethod
    def set_motor_curve(curve):
        """
        Set the global motor curve, must be a float array with exactly 21 elements.
        [0]  = x ticks/s  (0% power)
        [1]  = x ticks/s  (5% power)
        [2]  = x ticks/s  (10% power)
        ...
        [20] = x ticks/s  (100% power)

        :param curve: float array with 21 elements
        :raises: ValueError
        """
        if (len(curve) != 21):
            raise ValueError('The motor curve is invalid, check documentation for set_motor_curve()!')

        MOTOR_CURVE = curve

    @staticmethod
    def get_motor_curve():
        """
        Get the currently active motor curve. Check set_motor_curve() for more info.

        :return: current motor curve
        """
        return MOTOR_CURVE

    @staticmethod
    def __map(x, in_min, in_max, out_min, out_max):
        """
        Linear interpolation. Check https://www.arduino.cc/reference/en/language/functions/math/map/
        """
        return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min

    @staticmethod
    def __interpolateSpeed(curve, speed):
        """
        Interpolate a speed in the specified motor curve and return 
        the 'power percentage that is needed to reach that speed'
        """
        if speed < curve[0] or speed > curve[20]:
            raise ValueError(f'Speed out of range: {str(speed)} ticks/s')

        for index in range(20):  # There are 20 speed ranges
            if speed >= curve[index] and speed <= curve[index + 1] and curve[index] != curve[index + 1]:
                return Motor.__map(speed, curve[index], curve[index + 1], index * 5, index * 5 + 5)

        return 0

    @staticmethod
    def __linearizePower(curve, power):
        """
        Takes raw power and returns it corrected so that the motor speed would be roughly linear
        """
        if power < 0 or power > 100:
            raise ValueError(f'Power out of range: {str(power)}%')

        requiredSpeed = Motor.__map(power, 0, 100, curve[0], curve[20])
        return Motor.__interpolateSpeed(curve, requiredSpeed)


atexit.register(Motor.all_off)