001// Copyright (c) Choreo contributors
002
003package choreo.trajectory;
004
005import choreo.util.ChoreoAllianceFlipUtil;
006import edu.wpi.first.math.MathUtil;
007import edu.wpi.first.math.Matrix;
008import edu.wpi.first.math.VecBuilder;
009import edu.wpi.first.math.geometry.Pose2d;
010import edu.wpi.first.math.geometry.Rotation2d;
011import edu.wpi.first.math.kinematics.ChassisSpeeds;
012import edu.wpi.first.math.numbers.N1;
013import edu.wpi.first.math.numbers.N3;
014import edu.wpi.first.math.numbers.N6;
015import edu.wpi.first.math.system.NumericalIntegration;
016import edu.wpi.first.util.struct.Struct;
017import java.nio.ByteBuffer;
018import java.util.function.BiFunction;
019
020/** A single differential drive robot sample in a Trajectory. */
021public class DifferentialSample implements TrajectorySample<DifferentialSample> {
022  /** The timestamp of this sample relative to the beginning of the trajectory. */
023  public final double t;
024
025  /** The X position of the sample relative to the blue alliance wall origin in meters. */
026  public final double x;
027
028  /** The Y position of the sample relative to the blue alliance wall origin in meters. */
029  public final double y;
030
031  /** The heading of the sample in radians, with 0 being in the +X direction. */
032  public final double heading;
033
034  /** The velocity of the left side in m/s. */
035  public final double vl;
036
037  /** The velocity of the right side in m/s. */
038  public final double vr;
039
040  /** The chassis angular velocity in rad/s. */
041  public final double omega;
042
043  /** The acceleration of the left side in m/s². */
044  public final double al;
045
046  /** The acceleration of the right side in m/s². */
047  public final double ar;
048
049  /** The chassis angular acceleration in rad/s². */
050  public final double alpha;
051
052  /** The force of the left side in Newtons. */
053  public final double fl;
054
055  /** The force of the right side in Newtons. */
056  public final double fr;
057
058  /**
059   * Constructs a DifferentialSample with the specified parameters.
060   *
061   * @param timestamp The timestamp of this sample.
062   * @param x The X position of the sample in meters.
063   * @param y The Y position of the sample in meters.
064   * @param heading The heading of the sample in radians, with 0 being in the +X direction.
065   * @param vl The velocity of the left side in m/s.
066   * @param vr The velocity of the right side in m/s.
067   * @param omega The chassis angular velocity in rad/s.
068   * @param al The acceleration of the left side in m/s².
069   * @param ar The acceleration of the right side in m/s².
070   * @param alpha The chassis angular acceleration in rad/s².
071   * @param fl The force of the left side in Newtons.
072   * @param fr The force of the right side in Newtons.
073   */
074  public DifferentialSample(
075      double timestamp,
076      double x,
077      double y,
078      double heading,
079      double vl,
080      double vr,
081      double omega,
082      double al,
083      double ar,
084      double alpha,
085      double fl,
086      double fr) {
087    this.t = timestamp;
088    this.x = x;
089    this.y = y;
090    this.heading = heading;
091    this.vl = vl;
092    this.vr = vr;
093    this.omega = omega;
094    this.al = al;
095    this.ar = ar;
096    this.alpha = alpha;
097    this.fl = fl;
098    this.fr = fr;
099  }
100
101  @Override
102  public double getTimestamp() {
103    return t;
104  }
105
106  @Override
107  public Pose2d getPose() {
108    return new Pose2d(x, y, Rotation2d.fromRadians(heading));
109  }
110
111  /**
112   * Returns the field-relative chassis speeds of this sample.
113   *
114   * @return the field-relative chassis speeds of this sample.
115   * @see edu.wpi.first.math.kinematics.DifferentialDriveKinematics#toChassisSpeeds
116   */
117  @Override
118  public ChassisSpeeds getChassisSpeeds() {
119    return new ChassisSpeeds((vl + vr) / 2, 0, omega);
120  }
121
122  @Override
123  public DifferentialSample interpolate(DifferentialSample endValue, double timestamp) {
124    double scale = (timestamp - this.t) / (endValue.t - this.t);
125
126    // Integrate the acceleration to get the rest of the state, since linearly
127    // interpolating the state gives an inaccurate result if the accelerations are changing between
128    // states
129    Matrix<N6, N1> initialState = VecBuilder.fill(x, y, heading, vl, vr, omega);
130
131    BiFunction<Matrix<N6, N1>, Matrix<N3, N1>, Matrix<N6, N1>> f =
132        (state, input) -> {
133          //  state =  [x, y, θ, vₗ, vᵣ, ω]
134          //  input =  [aₗ, aᵣ, α]
135          //
136          //  v = (vₗ + vᵣ)/2
137          //
138          //  ẋ = v cosθ
139          //  ẏ = v sinθ
140          //  θ̇ = ω
141          //  v̇ₗ = aₗ
142          //  v̇ᵣ = aᵣ
143          //  ω̇ = α
144          var θ = state.get(2, 0);
145          var vl = state.get(3, 0);
146          var vr = state.get(4, 0);
147          var ω = state.get(5, 0);
148          var al = input.get(0, 0);
149          var ar = input.get(1, 0);
150          var α = input.get(2, 0);
151          var v = (vl + vr) / 2;
152          return VecBuilder.fill(v * Math.cos(θ), v * Math.sin(θ), ω, al, ar, α);
153        };
154
155    double τ = timestamp - this.t;
156    var sample = NumericalIntegration.rkdp(f, initialState, VecBuilder.fill(al, ar, alpha), τ);
157
158    return new DifferentialSample(
159        MathUtil.interpolate(this.t, endValue.t, scale),
160        sample.get(0, 0),
161        sample.get(1, 0),
162        sample.get(2, 0),
163        sample.get(3, 0),
164        sample.get(4, 0),
165        sample.get(5, 0),
166        this.al,
167        this.ar,
168        this.alpha,
169        MathUtil.interpolate(this.fl, endValue.fl, scale),
170        MathUtil.interpolate(this.fr, endValue.fr, scale));
171  }
172
173  public DifferentialSample flipped() {
174    return switch (ChoreoAllianceFlipUtil.getFlipper()) {
175      case MIRRORED ->
176          new DifferentialSample(
177              t,
178              ChoreoAllianceFlipUtil.flipX(x),
179              ChoreoAllianceFlipUtil.flipY(y), // No-op for mirroring
180              ChoreoAllianceFlipUtil.flipHeading(heading),
181              vr,
182              vl,
183              -omega,
184              ar,
185              al,
186              -alpha,
187              fr,
188              fl);
189      case ROTATE_AROUND ->
190          new DifferentialSample(
191              t,
192              ChoreoAllianceFlipUtil.flipX(x),
193              ChoreoAllianceFlipUtil.flipY(y),
194              ChoreoAllianceFlipUtil.flipHeading(heading),
195              vl,
196              vr,
197              omega,
198              al,
199              ar,
200              alpha,
201              fl,
202              fr);
203    };
204  }
205
206  public DifferentialSample offsetBy(double timestampOffset) {
207    return new DifferentialSample(
208        t + timestampOffset, x, y, heading, vl, vr, omega, al, ar, alpha, fl, fr);
209  }
210
211  /** The struct for the DifferentialSample class. */
212  public static final Struct<DifferentialSample> struct = new DifferentialSampleStruct();
213
214  private static final class DifferentialSampleStruct implements Struct<DifferentialSample> {
215    @Override
216    public Class<DifferentialSample> getTypeClass() {
217      return DifferentialSample.class;
218    }
219
220    @Override
221    public String getTypeName() {
222      return "DifferentialSample";
223    }
224
225    @Override
226    public int getSize() {
227      return Struct.kSizeDouble * 10;
228    }
229
230    @Override
231    public String getSchema() {
232      return "double timestamp;"
233          + "Pose2d pose;"
234          + "double vl;"
235          + "double vr;"
236          + "double omega;"
237          + "double al;"
238          + "double ar;"
239          + "double alpha;"
240          + "double fl;"
241          + "double fr;";
242    }
243
244    @Override
245    public Struct<?>[] getNested() {
246      return new Struct<?>[] {Pose2d.struct};
247    }
248
249    @Override
250    public DifferentialSample unpack(ByteBuffer bb) {
251      return new DifferentialSample(
252          bb.getDouble(),
253          bb.getDouble(),
254          bb.getDouble(),
255          bb.getDouble(),
256          bb.getDouble(),
257          bb.getDouble(),
258          bb.getDouble(),
259          bb.getDouble(),
260          bb.getDouble(),
261          bb.getDouble(),
262          bb.getDouble(),
263          bb.getDouble());
264    }
265
266    @Override
267    public void pack(ByteBuffer bb, DifferentialSample value) {
268      bb.putDouble(value.t);
269      bb.putDouble(value.x);
270      bb.putDouble(value.y);
271      bb.putDouble(value.heading);
272      bb.putDouble(value.vl);
273      bb.putDouble(value.vr);
274      bb.putDouble(value.omega);
275      bb.putDouble(value.al);
276      bb.putDouble(value.ar);
277      bb.putDouble(value.alpha);
278      bb.putDouble(value.fl);
279      bb.putDouble(value.fr);
280    }
281  }
282
283  @Override
284  public boolean equals(Object obj) {
285    if (!(obj instanceof DifferentialSample)) {
286      return false;
287    }
288
289    var other = (DifferentialSample) obj;
290    return this.t == other.t
291        && this.x == other.x
292        && this.y == other.y
293        && this.heading == other.heading
294        && this.vl == other.vl
295        && this.vr == other.vr
296        && this.omega == other.omega
297        && this.al == other.al
298        && this.ar == other.ar
299        && this.alpha == other.alpha
300        && this.fl == other.fl
301        && this.fr == other.fr;
302  }
303}