| from constants import * |
| import numpy as np |
| import scipy.optimize |
| from libspline import Spline, DistanceSpline, Trajectory |
| import copy |
| from dataclasses import dataclass |
| |
| |
| @dataclass |
| class ControlPointIndex: |
| """Class specifying the index of a control point""" |
| |
| # The index of the multispline in the list of multisplines |
| multispline_index: int |
| |
| # The index of the spline in the multispline |
| spline_index: int |
| |
| # The index of the control point in the spline [0-5] |
| control_point_index: int |
| |
| |
| class Multispline(): |
| def __init__(self): |
| self.staged_points = [] # Holds all points not yet in spline |
| self.libsplines = [] # Formatted for libspline library usage |
| self.splines = [] # Formatted for drawing |
| self.constraints = [ # default constraints |
| { |
| "constraint_type": "LONGITUDINAL_ACCELERATION", |
| "value": 3 |
| }, { |
| "constraint_type": "LATERAL_ACCELERATION", |
| "value": 2 |
| }, { |
| "constraint_type": "VOLTAGE", |
| "value": 10 |
| } |
| ] |
| |
| def __deepcopy__(self, memo): |
| new_copy = Multispline() |
| new_copy.staged_points = copy.deepcopy(self.staged_points, memo) |
| new_copy.splines = copy.deepcopy(self.splines, memo) |
| new_copy.constraints = copy.deepcopy(self.constraints, memo) |
| new_copy.update_lib_spline() |
| return new_copy |
| |
| def getLibsplines(self): |
| return self.libsplines |
| |
| def setConstraint(self, id, value): |
| for constraint in self.constraints: |
| if constraint["constraint_type"] == id: |
| constraint["value"] = value |
| break |
| |
| def getConstraint(self, id): |
| for constraint in self.constraints: |
| if constraint["constraint_type"] == id: |
| return constraint["value"] |
| |
| def addConstraintsToTrajectory(self, trajectory): |
| for constraint in self.constraints: |
| if constraint["constraint_type"] == "VOLTAGE": |
| trajectory.SetVoltageLimit(constraint["value"]) |
| elif constraint["constraint_type"] == "LATERAL_ACCELERATION": |
| trajectory.SetLateralAcceleration(constraint["value"]) |
| elif constraint["constraint_type"] == "LONGITUDINAL_ACCELERATION": |
| trajectory.SetLongitudinalAcceleration(constraint["value"]) |
| |
| @staticmethod |
| def splineExtrapolate(multisplines, index, *, snap=True): |
| multispline = multisplines[index.multispline_index] |
| spline = multispline.splines[index.spline_index] |
| |
| # find the adjacent splines that will be affected by this spline |
| |
| prev_multispline = None |
| prev_spline = None |
| if index.spline_index != 0: |
| prev_spline = multispline.splines[index.spline_index - 1] |
| elif index.multispline_index != 0: |
| prev_multispline = multisplines[index.multispline_index - 1] |
| prev_spline = prev_multispline.splines[-1] |
| |
| next_multispline = None |
| next_spline = None |
| if spline is not multispline.splines[-1]: |
| next_spline = multispline.splines[index.spline_index + 1] |
| elif multispline is not multisplines[-1]: |
| next_multispline = multisplines[index.multispline_index + 1] |
| if next_multispline.splines: |
| next_spline = next_multispline.splines[0] |
| |
| # adjust the adjacent splines according to our smoothness constraints |
| # the three points on either side are entirely constrained by this spline |
| |
| if next_spline is not None: |
| f = spline[5] # the end of the spline |
| e = spline[4] # determines the heading |
| d = spline[3] |
| if next_multispline is None: |
| next_spline[0] = f |
| next_spline[1] = f * 2 + e * -1 |
| next_spline[2] = d + f * 4 + e * -4 |
| else: |
| if snap: |
| Multispline.snapSplines(spline, next_spline, match_first_to_second=False) |
| |
| next_spline[0] = f |
| next_multispline.update_lib_spline() |
| |
| if prev_spline is not None: |
| a = spline[0] |
| b = spline[1] |
| c = spline[2] |
| if prev_multispline is None: |
| prev_spline[5] = a |
| prev_spline[4] = a * 2 + b * -1 |
| prev_spline[3] = c + a * 4 + b * -4 |
| else: |
| if snap: |
| Multispline.snapSplines(prev_spline, spline, match_first_to_second=True) |
| |
| prev_spline[5] = a |
| prev_multispline.update_lib_spline() |
| |
| def snapSplines(first_spline, second_spline, *, match_first_to_second): |
| """Snaps two adjacent splines together, preserving the heading from one to another. |
| |
| The end of `first_spline` connects to the beginning of `second_spline` |
| The user will have their mouse one one of these splines, so we |
| only want to snap the spline that they're not holding. |
| |
| They can have the same heading, or they can have opposite headings |
| which represents the robot driving that spline backwards. |
| """ |
| |
| # Represent the position of the second control point (controls the heading) |
| # as a vector from the end of the spline to the control point |
| first_vector = first_spline[-2] - first_spline[-1] |
| second_vector = second_spline[1] - second_spline[0] |
| |
| # we want to preserve the distance |
| first_magnitude = np.linalg.norm(first_vector, ord=2) |
| second_magnitude = np.linalg.norm(second_vector, ord=2) |
| |
| normalized_first = first_vector / first_magnitude |
| normalized_second = second_vector / second_magnitude |
| |
| # the proposed new vector if we were to make them point the same direction |
| swapped_second = normalized_first * second_magnitude |
| swapped_first = normalized_second * first_magnitude |
| |
| # they were pointing in opposite directions |
| if np.dot(first_vector, second_vector) < 0: |
| |
| # rotate the new vectors 180 degrees |
| # to keep them pointing in opposite directions |
| swapped_first = -swapped_first |
| swapped_second = -swapped_second |
| |
| # Calculate how far we ended up moving the second control point |
| # so we can move the third control point with it |
| change_in_first = swapped_first - first_vector |
| change_in_second = swapped_second - second_vector |
| |
| # apply the changes and discard the other proposed snap |
| if match_first_to_second: |
| first_spline[-2] = swapped_first + first_spline[-1] |
| first_spline[-3] += change_in_first |
| # swapped_second doesn't get used |
| else: |
| second_spline[1] = swapped_second + second_spline[0] |
| second_spline[2] += change_in_second |
| # swapped_first doesn't get used |
| |
| def updates_for_mouse_move(self, multisplines, index, mouse): |
| """Moves the control point and adjacent points to follow the mouse""" |
| if index == None: return |
| spline_edit = index.spline_index |
| index_of_edit = index.control_point_index |
| |
| spline = self.splines[spline_edit] |
| |
| # we want to move it to be on the mouse |
| diffs = mouse - spline[index_of_edit] |
| |
| spline[index_of_edit] = mouse |
| |
| |
| # all three points move together with the endpoint |
| if index_of_edit == 5: |
| spline[3] += diffs |
| spline[4] += diffs |
| |
| # check if the next multispline exists and has a spline |
| if index.spline_index == len( |
| self.splines) - 1 and not index.multispline_index == len( |
| multisplines) - 1 and len( |
| multisplines[index.multispline_index + |
| 1].splines) > 0: |
| # move the points that lay across the multispline boundary |
| other_spline = multisplines[index.multispline_index + |
| 1].splines[0] |
| |
| other_spline[1] += diffs |
| other_spline[2] += diffs |
| |
| if index_of_edit == 0: |
| spline[2] += diffs |
| spline[1] += diffs |
| |
| # check if previous multispline exists |
| if index.spline_index == 0 and not index.multispline_index == 0: |
| other_spline = multisplines[index.multispline_index - |
| 1].splines[-1] |
| |
| other_spline[3] += diffs |
| other_spline[4] += diffs |
| |
| # the third point moves with the second point |
| if index_of_edit == 4: |
| spline[3] += diffs |
| |
| if index_of_edit == 1: |
| spline[2] += diffs |
| |
| Multispline.splineExtrapolate(multisplines, index, snap=False) |
| |
| def update_lib_spline(self): |
| self.libsplines = [] |
| array = np.zeros(shape=(6, 2), dtype=float) |
| for points in self.splines: |
| for j, point in enumerate(points): |
| array[j, 0] = point[0] |
| array[j, 1] = point[1] |
| spline = Spline(np.ascontiguousarray(np.transpose(array))) |
| self.libsplines.append(spline) |
| |
| @staticmethod |
| def nearest_distance(multisplines, point): |
| """Finds the spot along the multisplines that is closest to a |
| given point on the field |
| |
| Returns the closest multispline and the distance along that multispline |
| """ |
| def distance(t, distance_spline, point): |
| return np.sum((distance_spline.XY(t) - point)**2) |
| |
| # We know the derivative of the function, |
| # so scipy doesn't need to compute it every time |
| def ddistance(t, distance_spline, point): |
| return np.sum(2 * (distance_spline.XY(t) - point) * |
| distance_spline.DXY(t)) |
| |
| best_result = None |
| best_multispline = None |
| |
| for multispline_index, multispline in enumerate(multisplines): |
| distance_spline = DistanceSpline(multispline.getLibsplines()) |
| |
| # The optimizer finds local minima that often aren't what we want, |
| # so try from multiple locations to find a better minimum. |
| guess_points = np.linspace(0, distance_spline.Length(), num=5) |
| |
| for guess in guess_points: |
| result = scipy.optimize.minimize( |
| distance, |
| guess, |
| args=(distance_spline, point), |
| bounds=((0, distance_spline.Length()), ), |
| jac=ddistance, |
| ) |
| |
| if result.success and (best_result == None |
| or result.fun < best_result.fun): |
| best_result = result |
| best_multispline = multispline |
| |
| return (best_multispline, best_result) |
| |
| def toJsonObject(self): |
| multi_spline = { |
| "spline_count": 0, |
| "spline_x": [], |
| "spline_y": [], |
| "constraints": self.constraints, |
| } |
| for points in self.splines: |
| multi_spline["spline_count"] += 1 |
| for j, point in enumerate(points): |
| if j == 0 and multi_spline["spline_count"] > 1: |
| continue # skip overlapping points |
| multi_spline["spline_x"].append(point[0]) |
| multi_spline["spline_y"].append(point[1]) |
| return multi_spline |
| |
| @staticmethod |
| def fromJsonObject(multi_spline): |
| multispline = Multispline() |
| multispline.constraints = multi_spline["constraints"] |
| multispline.splines = [] |
| multispline.staged_points = [] |
| |
| i = 0 |
| for j in range(multi_spline["spline_count"]): |
| # get the last point of the last spline |
| # and read in another 6 points |
| for i in range(i, i + 6): |
| multispline.staged_points.append( |
| [multi_spline["spline_x"][i], multi_spline["spline_y"][i]]) |
| multispline.splines.append(np.array(multispline.staged_points)) |
| multispline.staged_points = [] |
| multispline.update_lib_spline() |
| |
| return multispline |
| |
| def getSplines(self): |
| return self.splines |
| |
| def setControlPoint(self, index, x, y): |
| self.splines[index.spline_index][index.control_point_index] = [x, y] |
| |
| def addPoint(self, x, y): |
| if (len(self.staged_points) < 6): |
| self.staged_points.append([x, y]) |
| if (len(self.staged_points) == 6): |
| self.splines.append(np.array(self.staged_points)) |
| self.staged_points = [] |
| self.update_lib_spline() |
| return True |
| |
| def extrapolate(self, spline): |
| """Stages 3 points extrapolated from the end of the multispline""" |
| |
| self.staged_points = [] |
| |
| point1 = spline[5] |
| point2 = spline[4] |
| point3 = spline[3] |
| |
| self.staged_points.append(point1) |
| self.staged_points.append(point1 * 2 - point2) |
| self.staged_points.append(point3 + point1 * 4 - point2 * 4) |