ODrive/tools/odrive/utils.py

from __future__ import print_function

import sys
import time
import threading
import platform
import subprocess
import os
from fibre.utils import Event
from odrive.enums import *

try:
    if platform.system() == 'Windows':
        import win32console
        import colorama
        colorama.init()
except ImportError:
    print("Could not init terminal features.")
    print("Refer to install instructions at http://docs.odriverobotics.com/#downloading-and-installing-tools")
    sys.stdout.flush()
    pass

_VT100Colors = {
    'green': '\x1b[92;1m',
    'cyan': '\x1b[96;1m',
    'yellow': '\x1b[93;1m',
    'red': '\x1b[91;1m',
    'default': '\x1b[0m'
}

class OperationAbortedException(Exception):
    pass

def dump_errors(odrv, clear=False):
    axes = [(name, axis) for name, axis in odrv._remote_attributes.items() if 'axis' in name]
    axes.sort()
    for name, axis in axes:
        print(name)

        # Flatten axis and submodules
        # (name, remote_obj, errorcode)
        module_decode_map = [
            ('axis', axis, errors.axis),
            ('motor', axis.motor, errors.motor),
            ('encoder', axis.encoder, errors.encoder),
            ('controller', axis.controller, errors.controller),
        ]

        # Module error decode
        for name, remote_obj, errorcodes in module_decode_map:
            prefix = ' '*2 + name + ": "
            if (remote_obj.error != errorcodes.ERROR_NONE):
                print(prefix + _VT100Colors['red'] + "Error(s):" + _VT100Colors['default'])
                errorcodes_tup = [(name, val) for name, val in errorcodes.__dict__.items() if 'ERROR_' in name]
                for codename, codeval in errorcodes_tup:
                    if remote_obj.error & codeval != 0:
                        print("    " + codename)
                if clear:
                    remote_obj.error = errorcodes.ERROR_NONE
            else:
                print(prefix + _VT100Colors['green'] + "no error" + _VT100Colors['default'])

data_rate = 100
plot_rate = 10
num_samples = 1000
def start_liveplotter(get_var_callback):
    """
    Starts a liveplotter.
    The variable that is plotted is retrieved from get_var_callback.
    This function returns immediately and the liveplotter quits when
    the user closes it.
    """

    import matplotlib.pyplot as plt

    cancellation_token = Event()

    global vals
    vals = []
    def fetch_data():
        global vals
        while not cancellation_token.is_set():
            try:
                data = get_var_callback()
            except Exception as ex:
                print(str(ex))
                time.sleep(1)
                continue
            vals.append(data)
            if len(vals) > num_samples:
                vals = vals[-num_samples:]
            time.sleep(1/data_rate)

    # TODO: use animation for better UI performance, see:
    # https://matplotlib.org/examples/animation/simple_anim.html
    def plot_data():
        global vals

        plt.ion()

        # Make sure the script terminates when the user closes the plotter
        def did_close(evt):
            cancellation_token.set()
        fig = plt.figure()
        fig.canvas.mpl_connect('close_event', did_close)

        while not cancellation_token.is_set():
            plt.clf()
            plt.plot(vals)
            plt.legend(list(range(len(vals))))
            fig.canvas.draw()
            fig.canvas.start_event_loop(1/plot_rate)

    fetch_t = threading.Thread(target=fetch_data)
    fetch_t.daemon = True
    fetch_t.start()
    
    plot_t = threading.Thread(target=plot_data)
    plot_t.daemon = True
    plot_t.start()

    return cancellation_token;
    #plot_data()


class BulkCapture:
    '''
    Asynchronously captures a bulk set of data when instance is created.

    get_var_callback: a function that returns the data you want to collect (see the example below)
    data_rate: Rate in hz
    length: Length of time to capture in seconds

    Example Usage:
        capture = BulkCapture(lambda :[odrv0.axis0.encoder.pos_estimate, odrv0.axis0.controller.pos_setpoint])
        # Do stuff while capturing (like sending position commands)
        capture.event.wait() # When you're done doing stuff, wait for the capture to be completed.
        print(capture.data) # Do stuff with the data
        capture.plot_data() # Helper method to plot the data
    '''

    def __init__(self,
                 get_var_callback,
                 data_rate=500.0,
                 length=2.0):
        from threading import Event, Thread
        import pandas as pd

        self.event = Event()
        def loop():
            vals = []
            start_time = time.monotonic()
            total_samples = int(length * data_rate)
            period = 1.0 / data_rate
            for i in range(total_samples):
                try:
                    data = get_var_callback()
                except Exception as ex:
                    print(str(ex))
                    print("Waiting 1 second before next data point")
                    time.sleep(1)
                    continue
                relative_time = time.monotonic() - start_time
                vals.append([relative_time] + data)
                time.sleep(period - (relative_time % period)) # this ensures consistently timed samples
            self.data = pd.DataFrame(vals) # A lock is not really necessary due to the event
            print("Achieved average data rate: {}Hz".format(total_samples / self.data.iloc[-1, 0]))
            print("If this rate is significantly lower than what you specified, consider lowering it below the achieved value for more consistent sampling.")
            self.event.set() # tell the main thread that the bulk capture is complete
        Thread(target=loop, daemon=True).start()
    
    def plot_data(self):
        import matplotlib.pyplot as plt
        plt.plot(self.data[0], self.data.drop(0, axis=1))
        plt.xlabel("Time (seconds)")
        plt.ylabel("Counts")
        plt.show()


def step_and_plot(axis, step_size=100.0, settle_time=1.0, data_rate=500.0):
    initial_settle_time = 0.5
    axis.requested_state = AXIS_STATE_CLOSED_LOOP_CONTROL
    capture = BulkCapture(lambda :[axis.encoder.pos_estimate, axis.controller.pos_setpoint],
                          data_rate=data_rate,
                          length = settle_time + initial_settle_time)
    initial_setpoint = axis.encoder.pos_estimate
    axis.controller.pos_setpoint = initial_setpoint # set current position as setpoint
    time.sleep(initial_settle_time)
    axis.controller.pos_setpoint = initial_setpoint + step_size # relative/incremental movement
    capture.event.wait() # wait for Bulk Capture to be complete
    axis.requested_state = AXIS_STATE_IDLE
    capture.plot_data()


def print_drv_regs(name, motor):
    """
    Dumps the current gate driver regisers for the specified motor
    """
    fault = motor.gate_driver.drv_fault
    status_reg_1 = motor.gate_driver.status_reg_1
    status_reg_2 = motor.gate_driver.status_reg_2
    ctrl_reg_1 = motor.gate_driver.ctrl_reg_1
    ctrl_reg_2 = motor.gate_driver.ctrl_reg_2
    print(name + ": " + str(fault))
    print("DRV Fault Code: " + str(fault))
    print("Status Reg 1: " + str(status_reg_1) + " (" + format(status_reg_1, '#010b') + ")")
    print("Status Reg 2: " + str(status_reg_2) + " (" + format(status_reg_2, '#010b') + ")")
    print("Control Reg 1: " + str(ctrl_reg_1) + " (" + format(ctrl_reg_1, '#013b') + ")")
    print("Control Reg 2: " + str(ctrl_reg_2) + " (" + format(ctrl_reg_2, '#09b') + ")")

def show_oscilloscope(odrv):
    size = 18000
    values = []
    for i in range(size):
        values.append(odrv.get_oscilloscope_val(i))

    import matplotlib.pyplot as plt
    plt.plot(values)
    plt.show()

def rate_test(device):
    """
    Tests how many integers per second can be transmitted
    """

    # import matplotlib.pyplot as plt
    # plt.ion()

    print("reading 10000 values... new value")
    numFrames = 10000
    vals = []
    for _ in range(numFrames):
        vals.append(device.axis0.loop_counter)

    loopsPerFrame = (vals[-1] - vals[0])/numFrames
    loopsPerSec = (168000000/(6*3500))
    FramePerSec = loopsPerSec/loopsPerFrame
    print("Frames per second: " + str(FramePerSec))

    # plt.plot(vals)
    # plt.show(block=True)

def usb_burn_in_test(get_var_callback, cancellation_token):
    """
    Starts background threads that read a values form the USB device in a spin-loop
    """

    def fetch_data():
        global vals
        i = 0
        while not cancellation_token.is_set():
            try:
                get_var_callback()
                i += 1
            except Exception as ex:
                print(str(ex))
                time.sleep(1)
                i = 0
                continue
            if i % 1000 == 0:
                print("read {} values".format(i))
    threading.Thread(target=fetch_data, daemon=True).start()

def yes_no_prompt(question, default=None):
    if default is None:
        question += " [y/n] "
    elif default == True:
        question += " [Y/n] "
    elif default == False:
        question += " [y/N] "

    while True:
        print(question, end='')

        choice = input().lower()
        if choice in {'yes', 'y'}:
            return True
        elif choice in {'no', 'n'}:
            return False
        elif choice == '' and default is not None:
            return default