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