001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.math.system;
006
007import edu.wpi.first.math.Matrix;
008import edu.wpi.first.math.Nat;
009import edu.wpi.first.math.Num;
010import edu.wpi.first.math.numbers.N1;
011import java.util.function.BiFunction;
012import java.util.function.Function;
013
014public final class NumericalJacobian {
015  private NumericalJacobian() {
016    // Utility Class.
017  }
018
019  private static final double kEpsilon = 1e-5;
020
021  /**
022   * Computes the numerical Jacobian with respect to x for f(x).
023   *
024   * @param <Rows> Number of rows in the result of f(x).
025   * @param <States> Num representing the number of rows in the output of f.
026   * @param <Cols> Number of columns in the result of f(x).
027   * @param rows Number of rows in the result of f(x).
028   * @param cols Number of columns in the result of f(x).
029   * @param f Vector-valued function from which to compute the Jacobian.
030   * @param x Vector argument.
031   * @return The numerical Jacobian with respect to x for f(x, u, ...).
032   */
033  @SuppressWarnings({"ParameterName", "MethodTypeParameterName"})
034  public static <Rows extends Num, Cols extends Num, States extends Num>
035      Matrix<Rows, Cols> numericalJacobian(
036          Nat<Rows> rows,
037          Nat<Cols> cols,
038          Function<Matrix<Cols, N1>, Matrix<States, N1>> f,
039          Matrix<Cols, N1> x) {
040    var result = new Matrix<>(rows, cols);
041
042    for (int i = 0; i < cols.getNum(); i++) {
043      var dxPlus = x.copy();
044      var dxMinus = x.copy();
045      dxPlus.set(i, 0, dxPlus.get(i, 0) + kEpsilon);
046      dxMinus.set(i, 0, dxMinus.get(i, 0) - kEpsilon);
047      @SuppressWarnings("LocalVariableName")
048      var dF = f.apply(dxPlus).minus(f.apply(dxMinus)).div(2 * kEpsilon);
049
050      result.setColumn(i, Matrix.changeBoundsUnchecked(dF));
051    }
052
053    return result;
054  }
055
056  /**
057   * Returns numerical Jacobian with respect to x for f(x, u, ...).
058   *
059   * @param <Rows> Number of rows in the result of f(x, u).
060   * @param <States> Number of rows in x.
061   * @param <Inputs> Number of rows in the second input to f.
062   * @param <Outputs> Num representing the rows in the output of f.
063   * @param rows Number of rows in the result of f(x, u).
064   * @param states Number of rows in x.
065   * @param f Vector-valued function from which to compute Jacobian.
066   * @param x State vector.
067   * @param u Input vector.
068   * @return The numerical Jacobian with respect to x for f(x, u, ...).
069   */
070  @SuppressWarnings({"LambdaParameterName", "MethodTypeParameterName"})
071  public static <Rows extends Num, States extends Num, Inputs extends Num, Outputs extends Num>
072      Matrix<Rows, States> numericalJacobianX(
073          Nat<Rows> rows,
074          Nat<States> states,
075          BiFunction<Matrix<States, N1>, Matrix<Inputs, N1>, Matrix<Outputs, N1>> f,
076          Matrix<States, N1> x,
077          Matrix<Inputs, N1> u) {
078    return numericalJacobian(rows, states, _x -> f.apply(_x, u), x);
079  }
080
081  /**
082   * Returns the numerical Jacobian with respect to u for f(x, u).
083   *
084   * @param <States> The states of the system.
085   * @param <Inputs> The inputs to the system.
086   * @param <Rows> Number of rows in the result of f(x, u).
087   * @param rows Number of rows in the result of f(x, u).
088   * @param inputs Number of rows in u.
089   * @param f Vector-valued function from which to compute the Jacobian.
090   * @param x State vector.
091   * @param u Input vector.
092   * @return the numerical Jacobian with respect to u for f(x, u).
093   */
094  @SuppressWarnings({"LambdaParameterName", "MethodTypeParameterName"})
095  public static <Rows extends Num, States extends Num, Inputs extends Num>
096      Matrix<Rows, Inputs> numericalJacobianU(
097          Nat<Rows> rows,
098          Nat<Inputs> inputs,
099          BiFunction<Matrix<States, N1>, Matrix<Inputs, N1>, Matrix<States, N1>> f,
100          Matrix<States, N1> x,
101          Matrix<Inputs, N1> u) {
102    return numericalJacobian(rows, inputs, _u -> f.apply(x, _u), u);
103  }
104}