Field-Centric control code for a VEX IQ X-Drive robot

iq_fcc.py

# ------------------------------------------
# 
# 	Project:      VEX IQ Field-Centric Holonomic Drive
# 	Author:       John Holbrook
# 	Created:      May 2023
# 	Description:  Drive code for a VEX IQ X-drive robot with field-centric drive
#                 (i.e., pushing the stick up makes the robot drive in a constant
#                 direction regardless of heading).
# 
# ------------------------------------------

# Library imports
from vex import *
from math import pow, sin, cos, radians, degrees, atan, fabs, sqrt
import urandom

# Brain should be defined by default
brain=Brain()

# Robot configuration code
brain_inertial = Inertial()
motor_1 = Motor(Ports.PORT1, False)
motor_6 = Motor(Ports.PORT6, True)
motor_7 = Motor(Ports.PORT7, False)
motor_12 = Motor(Ports.PORT12, True)
controller = Controller()

def cube(axis):
    '''
    apply cubic scaling to a joystick axis
    '''
    return pow(axis, 3)/10000
    
def draw_line(start, end):
    '''
    Helper function to draw a line between two points, represented as tuples.
    '''
    brain.screen.draw_line(start[0], start[1], end[0], end[1])

def rotate(p, c, deg):
    '''
    Rotate point p about point c by the specified angle in degrees. Both p and c are tuples.
    '''
    rad = radians(deg)
    cx = c[0]
    cy = c[1] 
    return (cx + int(((p[0]-cx)*cos(rad)) - ((p[1]-cy)*sin(rad))),
            cy + int(((p[0]-cx)*sin(rad)) + ((p[1]-cy)*cos(rad))))

def draw_arrow(theta):
    '''
    Draw an arrow pointing in the specified direction (in degrees)
    '''
    brain.screen.clear_screen()
    brain.screen.set_pen_color(Color.WHITE)
    center = (80, 50)
    lines = [
        ((70,80), (90,80)),
        ((90, 80), (90, 50)),
        ((90, 50), (105, 50)),
        ((105, 50), (80, 20)),
        ((80, 20), (55, 50)),
        ((55, 50), (70, 50)),
        ((70, 50), (70, 80))
    ]
    for line in lines:
        draw_line(rotate(line[0], center, theta), rotate(line[1], center, theta))

def coords_to_angle(x,y):
    '''
    Given the x and y position of the left joystick, compute the angle of the stick
    (0 degrees is straight up, values increase clockwise)
    '''
    if x>0 and y>0:
        return int(degrees(atan(x/y)))
    elif x>0 and y<0:
        return 90+int(degrees(atan(fabs(y)/x)))
    elif x<0 and y<0:
        return 180+int(degrees(atan(fabs(x)/fabs(y))))
    elif x<0 and y>0:
        return 270+int(degrees(atan(y/fabs(x))))
    elif x==0 and y<0: return 180
    elif x==0 and y>0: return 0
    elif x<0 and y==0: return 270
    elif x>0 and y==0: return 90
    else:
        return 0

def magnitude(x,y):
    '''
    Given the x and y position of the left joystick, compute the distance from center to joystick position
    (full range is 100)
    '''
    return int(sqrt( pow(x,2) + pow(y,2) ))

def polar_to_xy(theta, r):
    '''
    Convert from a given angle and distance to equavalent joystick x/y coordinates
    '''
    # print("Theta: ", theta)
    if theta == 0:
        return (0, r)
    elif theta == 90:
        return (r, 0)
    elif theta == 180:
        return (0, -r)
    elif theta == 270:
        return (-r, 0)
    elif theta>0 and theta<90:
        theta = radians(theta)
        return (int(r*sin(theta)), int(r*cos(theta)))
    elif theta>90 and theta<180:
        theta = radians(theta-90)
        return (int(r*cos(theta)), -1*int(r*sin(theta)))
    elif (theta>180 and theta<270):
        theta = radians(theta-180)
        return (-1*int(r*sin(theta)), -1*int(r*cos(theta)))
    elif (theta>270):
        theta = radians(theta-270)
        return (-1*int(r*cos(theta)), int(r*sin(theta)))
    else: return ("ERROR", "ERROR")

# global to enable or disable field-centric control   
field_centric_enabled = True
def toggle_fcc():
    global field_centric_enabled
    field_centric_enabled = not field_centric_enabled
    for _ in range(2): draw_arrow(0)

# toggle FCC when F Up is pressed
controller.buttonFUp.pressed(toggle_fcc)

calibrating = False
def calibrate():
    global calibrating

    # stop all motors so the robot can't move during calibration
    motor_1.stop()
    motor_6.stop()
    motor_7.stop()
    motor_12.stop()

    calibrating = True

    brain.screen.clear_screen()
    brain.screen.set_fill_color(Color.RED)
    brain.screen.set_pen_color(Color.RED)
    brain.screen.draw_rectangle(0, 0, 160, 100)

    brain_inertial.calibrate()
    wait(4, SECONDS)

    brain.screen.set_fill_color(Color.GREEN)
    brain.screen.set_pen_color(Color.GREEN)
    brain.screen.draw_rectangle(0, 0, 160, 100)
    brain.play_note(2, 0, 500)
    brain.screen.clear_screen()

    calibrating = False

    # re-enable motors
    motor_1.spin(FORWARD)
    motor_6.spin(FORWARD)
    motor_7.spin(FORWARD)
    motor_12.spin(FORWARD)

# recalibrate when R3 is pressed (click on right stick)
controller.buttonR3.pressed(calibrate)

def main():
    global field_centric_enabled, calibrating

    calibrate()
    
    count = 0
    while True:
        # get robot heading
        robot_heading = brain_inertial.heading(DEGREES)
    
        # update arrow on screen (only every 5th cycle to reduce flickering)
        count = (count+1)%5
        if count==0 and field_centric_enabled and not calibrating:
            draw_arrow(360-robot_heading)

        # fetch raw joystick values and apply cubic scaling
        ax_a = cube(controller.axisA.position())
        ax_b = cube(controller.axisB.position())
        ax_c = cube(controller.axisC.position())

        if field_centric_enabled:
            # field-centric control

            # convert left stick values to polar coordinates
            stick_angle = coords_to_angle(ax_b, ax_a)
            stick_mag = magnitude(ax_b, ax_a)

            # calculate direction robot should move relative to itself, based on current heading and
            # left stick angle
            robot_local_dir = (360 - robot_heading + stick_angle) % 360

            # print(stick_angle, stick_mag, robot_heading, robot_local_dir)

            # convert robot direction back to joystick coordinates for math below
            (rbt_b, rbt_a) = polar_to_xy(robot_local_dir, stick_mag)
            
            # compute and set power for each motor
            motor_1.set_velocity(rbt_a + rbt_b + ax_c, PERCENT)
            motor_6.set_velocity(rbt_a - rbt_b - ax_c, PERCENT)
            motor_7.set_velocity(rbt_a - rbt_b + ax_c, PERCENT)
            motor_12.set_velocity(rbt_a + rbt_b - ax_c, PERCENT)

        else: 
            # normal control (robot-centric)
            motor_1.set_velocity(ax_a + ax_b + ax_c, PERCENT)
            motor_6.set_velocity(ax_a - ax_b - ax_c, PERCENT)
            motor_7.set_velocity(ax_a - ax_b + ax_c, PERCENT)
            motor_12.set_velocity(ax_a + ax_b - ax_c, PERCENT)

        # wait 10ms
        wait(10, MSEC)

main()