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 ChoreoAllianceFlipUtil.getFlipper().flip(this);
175  }
176
177  @Override
178  public DifferentialSample mirrorX() {
179    return ChoreoAllianceFlipUtil.getMirrorX().flip(this);
180  }
181
182  @Override
183  public DifferentialSample mirrorY() {
184    return ChoreoAllianceFlipUtil.getMirrorY().flip(this);
185  }
186
187  @Override
188  public DifferentialSample rotateAround() {
189    return ChoreoAllianceFlipUtil.getRotateAround().flip(this);
190  }
191
192  public DifferentialSample offsetBy(double timestampOffset) {
193    return new DifferentialSample(
194        t + timestampOffset, x, y, heading, vl, vr, omega, al, ar, alpha, fl, fr);
195  }
196
197  /** The struct for the DifferentialSample class. */
198  public static final Struct<DifferentialSample> struct = new DifferentialSampleStruct();
199
200  private static final class DifferentialSampleStruct implements Struct<DifferentialSample> {
201    @Override
202    public Class<DifferentialSample> getTypeClass() {
203      return DifferentialSample.class;
204    }
205
206    @Override
207    public String getTypeName() {
208      return "DifferentialSample";
209    }
210
211    @Override
212    public int getSize() {
213      return Struct.kSizeDouble * 10;
214    }
215
216    @Override
217    public String getSchema() {
218      return "double timestamp;"
219          + "Pose2d pose;"
220          + "double vl;"
221          + "double vr;"
222          + "double omega;"
223          + "double al;"
224          + "double ar;"
225          + "double alpha;"
226          + "double fl;"
227          + "double fr;";
228    }
229
230    @Override
231    public Struct<?>[] getNested() {
232      return new Struct<?>[] {Pose2d.struct};
233    }
234
235    @Override
236    public DifferentialSample unpack(ByteBuffer bb) {
237      return new DifferentialSample(
238          bb.getDouble(),
239          bb.getDouble(),
240          bb.getDouble(),
241          bb.getDouble(),
242          bb.getDouble(),
243          bb.getDouble(),
244          bb.getDouble(),
245          bb.getDouble(),
246          bb.getDouble(),
247          bb.getDouble(),
248          bb.getDouble(),
249          bb.getDouble());
250    }
251
252    @Override
253    public void pack(ByteBuffer bb, DifferentialSample value) {
254      bb.putDouble(value.t);
255      bb.putDouble(value.x);
256      bb.putDouble(value.y);
257      bb.putDouble(value.heading);
258      bb.putDouble(value.vl);
259      bb.putDouble(value.vr);
260      bb.putDouble(value.omega);
261      bb.putDouble(value.al);
262      bb.putDouble(value.ar);
263      bb.putDouble(value.alpha);
264      bb.putDouble(value.fl);
265      bb.putDouble(value.fr);
266    }
267  }
268
269  @Override
270  public boolean equals(Object obj) {
271    if (!(obj instanceof DifferentialSample)) {
272      return false;
273    }
274
275    var other = (DifferentialSample) obj;
276    return MathUtil.isNear(this.t, other.t, 1E-6)
277        && MathUtil.isNear(this.x, other.x, 1E-6)
278        && MathUtil.isNear(this.y, other.y, 1E-6)
279        && MathUtil.isNear(this.heading, other.heading, 1E-6)
280        && MathUtil.isNear(this.vl, other.vl, 1E-6)
281        && MathUtil.isNear(this.vr, other.vr, 1E-6)
282        && MathUtil.isNear(this.omega, other.omega, 1E-6)
283        && MathUtil.isNear(this.al, other.al, 1E-6)
284        && MathUtil.isNear(this.ar, other.ar, 1E-6)
285        && MathUtil.isNear(this.alpha, other.alpha, 1E-6)
286        && MathUtil.isNear(this.fl, other.fl, 1E-6)
287        && MathUtil.isNear(this.fr, other.fr, 1E-6);
288  }
289}