001// Copyright (c) Choreo contributors 002 003package choreo.trajectory; 004 005import choreo.util.ChoreoAllianceFlipUtil; 006import choreo.util.ChoreoArrayUtil; 007import edu.wpi.first.math.MathUtil; 008import edu.wpi.first.math.geometry.Pose2d; 009import edu.wpi.first.math.geometry.Rotation2d; 010import edu.wpi.first.math.kinematics.ChassisSpeeds; 011import edu.wpi.first.util.struct.Struct; 012import java.nio.ByteBuffer; 013import java.util.Arrays; 014 015/** A single swerve robot sample in a Trajectory. */ 016public class SwerveSample implements TrajectorySample<SwerveSample> { 017 private static final double[] EMPTY_MODULE_FORCES = new double[] {0, 0, 0, 0}; 018 019 /** The timestamp of this sample, relative to the beginning of the trajectory. */ 020 public final double t; 021 022 /** The X position of the sample relative to the blue alliance wall origin in meters. */ 023 public final double x; 024 025 /** The Y position of the sample relative to the blue alliance wall origin in meters. */ 026 public final double y; 027 028 /** The heading of the sample in radians, with 0 being in the +X direction. */ 029 public final double heading; 030 031 /** The velocity of the sample in the X direction in m/s. */ 032 public final double vx; 033 034 /** The velocity of the sample in the Y direction in m/s. */ 035 public final double vy; 036 037 /** The angular velocity of the sample in rad/s. */ 038 public final double omega; 039 040 /** The acceleration of the sample in the X direction in m/s². */ 041 public final double ax; 042 043 /** The acceleration of the sample in the Y direction in m/s². */ 044 public final double ay; 045 046 /** The angular acceleration of the sample in rad/s². */ 047 public final double alpha; 048 049 /** 050 * The force on each swerve module in the X direction in Newtons. Module forces appear in the 051 * following order: [FL, FR, BL, BR]. 052 */ 053 private final double[] fx; 054 055 /** 056 * The force on each swerve module in the Y direction in Newtons Module forces appear in the 057 * following order: [FL, FR, BL, BR]. 058 */ 059 private final double[] fy; 060 061 /** 062 * Constructs a SwerveSample with the specified parameters. 063 * 064 * @param t The timestamp of this sample, relative to the beginning of the trajectory. 065 * @param x The X position of the sample in meters. 066 * @param y The Y position of the sample in meters. 067 * @param heading The heading of the sample in radians, with 0 being in the +X direction. 068 * @param vx The velocity of the sample in the X direction in m/s. 069 * @param vy The velocity of the sample in the Y direction in m/s. 070 * @param omega The angular velocity of the sample in rad/s. 071 * @param ax The acceleration of the sample in the X direction in m/s². 072 * @param ay The acceleration of the sample in the Y direction in m/s². 073 * @param alpha The angular acceleration of the sample in rad/s². 074 * @param moduleForcesX The force on each swerve module in the X direction in Newtons. Module 075 * forces appear in the following order: [FL, FR, BL, BR]. 076 * @param moduleForcesY The force on each swerve module in the Y direction in Newtons. Module 077 * forces appear in the following order: [FL, FR, BL, BR]. 078 */ 079 public SwerveSample( 080 double t, 081 double x, 082 double y, 083 double heading, 084 double vx, 085 double vy, 086 double omega, 087 double ax, 088 double ay, 089 double alpha, 090 double[] moduleForcesX, 091 double[] moduleForcesY) { 092 this.t = t; 093 this.x = x; 094 this.y = y; 095 this.heading = heading; 096 this.vx = vx; 097 this.vy = vy; 098 this.omega = omega; 099 this.ax = ax; 100 this.ay = ay; 101 this.alpha = alpha; 102 this.fx = moduleForcesX; 103 this.fy = moduleForcesY; 104 } 105 106 /** 107 * A null safe getter for the module forces in the X direction. 108 * 109 * @return The module forces in the X direction. 110 */ 111 public double[] moduleForcesX() { 112 if (fx == null || fx.length != 4) { 113 return EMPTY_MODULE_FORCES; 114 } 115 return fx; 116 } 117 118 /** 119 * A null safe getter for the module forces in the Y direction. 120 * 121 * @return The module forces in the Y direction. 122 */ 123 public double[] moduleForcesY() { 124 if (fy == null || fy.length != 4) { 125 return EMPTY_MODULE_FORCES; 126 } 127 return fy; 128 } 129 130 @Override 131 public double getTimestamp() { 132 return t; 133 } 134 135 @Override 136 public Pose2d getPose() { 137 return new Pose2d(x, y, Rotation2d.fromRadians(heading)); 138 } 139 140 @Override 141 public ChassisSpeeds getChassisSpeeds() { 142 return new ChassisSpeeds(vx, vy, omega); 143 } 144 145 @Override 146 public SwerveSample interpolate(SwerveSample endValue, double timestamp) { 147 double scale = (timestamp - this.t) / (endValue.t - this.t); 148 149 double[] interp_fx = new double[4]; 150 double[] interp_fy = new double[4]; 151 for (int i = 0; i < 4; ++i) { 152 interp_fx[i] = 153 MathUtil.interpolate(this.moduleForcesX()[i], endValue.moduleForcesX()[i], scale); 154 interp_fy[i] = 155 MathUtil.interpolate(this.moduleForcesY()[i], endValue.moduleForcesY()[i], scale); 156 } 157 158 // Integrate the acceleration to get the rest of the state, since linearly 159 // interpolating the state gives an inaccurate result if the accelerations are changing between 160 // states 161 // 162 // τ = timestamp − tₖ 163 // 164 // x(τ) = xₖ + vₖτ + 1/2 aₖτ² 165 // v(τ) = vₖ + aₖτ 166 double τ = timestamp - this.t; 167 double τ2 = τ * τ; 168 return new SwerveSample( 169 timestamp, 170 this.x + this.vx * τ + 0.5 * this.ax * τ2, 171 this.y + this.vy * τ + 0.5 * this.ay * τ2, 172 this.heading + this.omega * τ + 0.5 * this.alpha * τ2, 173 this.vx + this.ax * τ, 174 this.vy + this.ay * τ, 175 this.omega + this.alpha * τ, 176 this.ax, 177 this.ay, 178 this.alpha, 179 interp_fx, 180 interp_fy); 181 } 182 183 @Override 184 public SwerveSample offsetBy(double timestampOffset) { 185 return new SwerveSample( 186 this.t + timestampOffset, 187 this.x, 188 this.y, 189 this.heading, 190 this.vx, 191 this.vy, 192 this.omega, 193 this.ax, 194 this.ay, 195 this.alpha, 196 this.moduleForcesX(), 197 this.moduleForcesY()); 198 } 199 200 @Override 201 public SwerveSample flipped() { 202 return switch (ChoreoAllianceFlipUtil.getFlipper()) { 203 case MIRRORED -> 204 new SwerveSample( 205 this.t, 206 ChoreoAllianceFlipUtil.flipX(this.x), 207 ChoreoAllianceFlipUtil.flipY(this.y), 208 ChoreoAllianceFlipUtil.flipHeading(this.heading), 209 -this.vx, 210 this.vy, 211 -this.omega, 212 -this.ax, 213 this.ay, 214 -this.alpha, 215 // FL, FR, BL, BR 216 // Mirrored 217 // -FR, -FL, -BR, -BL 218 new double[] { 219 -this.moduleForcesX()[1], 220 -this.moduleForcesX()[0], 221 -this.moduleForcesX()[3], 222 -this.moduleForcesX()[2] 223 }, 224 // FL, FR, BL, BR 225 // Mirrored 226 // FR, FL, BR, BL 227 new double[] { 228 this.moduleForcesY()[1], 229 this.moduleForcesY()[0], 230 this.moduleForcesY()[3], 231 this.moduleForcesY()[2] 232 }); 233 case ROTATE_AROUND -> 234 new SwerveSample( 235 this.t, 236 ChoreoAllianceFlipUtil.flipX(this.x), 237 ChoreoAllianceFlipUtil.flipY(this.y), 238 ChoreoAllianceFlipUtil.flipHeading(this.heading), 239 -this.vx, 240 -this.vy, 241 this.omega, 242 -this.ax, 243 -this.ay, 244 this.alpha, 245 Arrays.stream(this.moduleForcesX()).map(x -> -x).toArray(), 246 Arrays.stream(this.moduleForcesY()).map(y -> -y).toArray()); 247 }; 248 } 249 250 /** The struct for the SwerveSample class. */ 251 public static final Struct<SwerveSample> struct = new SwerveSampleStruct(); 252 253 private static final class SwerveSampleStruct implements Struct<SwerveSample> { 254 @Override 255 public Class<SwerveSample> getTypeClass() { 256 return SwerveSample.class; 257 } 258 259 @Override 260 public String getTypeName() { 261 return "SwerveSample"; 262 } 263 264 @Override 265 public int getSize() { 266 return Struct.kSizeDouble * 18; 267 } 268 269 @Override 270 public String getSchema() { 271 return "double timestamp;" 272 + "Pose2d pose;" 273 + "double vx;" 274 + "double vy;" 275 + "double omega;" 276 + "double ax;" 277 + "double ay;" 278 + "double alpha;" 279 + "double moduleForcesX[4];" 280 + "double moduleForcesY[4];"; 281 } 282 283 @Override 284 public Struct<?>[] getNested() { 285 return new Struct<?>[] {Pose2d.struct}; 286 } 287 288 @Override 289 public SwerveSample unpack(ByteBuffer bb) { 290 return new SwerveSample( 291 bb.getDouble(), 292 bb.getDouble(), 293 bb.getDouble(), 294 bb.getDouble(), 295 bb.getDouble(), 296 bb.getDouble(), 297 bb.getDouble(), 298 bb.getDouble(), 299 bb.getDouble(), 300 bb.getDouble(), 301 new double[] {bb.getDouble(), bb.getDouble(), bb.getDouble(), bb.getDouble()}, 302 new double[] {bb.getDouble(), bb.getDouble(), bb.getDouble(), bb.getDouble()}); 303 } 304 305 @Override 306 public void pack(ByteBuffer bb, SwerveSample value) { 307 bb.putDouble(value.t); 308 bb.putDouble(value.x); 309 bb.putDouble(value.y); 310 bb.putDouble(value.heading); 311 bb.putDouble(value.vx); 312 bb.putDouble(value.vy); 313 bb.putDouble(value.omega); 314 bb.putDouble(value.ax); 315 bb.putDouble(value.ay); 316 bb.putDouble(value.alpha); 317 for (int i = 0; i < 4; ++i) { 318 bb.putDouble(value.moduleForcesX()[i]); 319 } 320 for (int i = 0; i < 4; ++i) { 321 bb.putDouble(value.moduleForcesY()[i]); 322 } 323 } 324 } 325 326 @Override 327 public boolean equals(Object obj) { 328 if (!(obj instanceof SwerveSample)) { 329 return false; 330 } 331 332 var other = (SwerveSample) obj; 333 return this.t == other.t 334 && this.x == other.x 335 && this.y == other.y 336 && this.heading == other.heading 337 && this.vx == other.vx 338 && this.vy == other.vy 339 && this.omega == other.omega 340 && this.ax == other.ax 341 && this.ay == other.ay 342 && this.alpha == other.alpha 343 && ChoreoArrayUtil.zipEquals( 344 this.fx, other.fx, (a, b) -> a.doubleValue() == b.doubleValue()) 345 && ChoreoArrayUtil.zipEquals( 346 this.fy, other.fy, (a, b) -> a.doubleValue() == b.doubleValue()); 347 } 348}