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
014/** Numerical Jacobian utilities. */
015public final class NumericalJacobian {
016  private NumericalJacobian() {
017    // Utility Class.
018  }
019
020  private static final double kEpsilon = 1e-5;
021
022  /**
023   * Computes the numerical Jacobian with respect to x for f(x).
024   *
025   * @param <Rows> Number of rows in the result of f(x).
026   * @param <States> Num representing the number of rows in the output of f.
027   * @param <Cols> Number of columns in the result of f(x).
028   * @param rows Number of rows in the result of f(x).
029   * @param cols Number of columns in the result of f(x).
030   * @param f Vector-valued function from which to compute the Jacobian.
031   * @param x Vector argument.
032   * @return The numerical Jacobian with respect to x for f(x, u, ...).
033   */
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      var dF = f.apply(dxPlus).minus(f.apply(dxMinus)).div(2 * kEpsilon);
048
049      result.setColumn(i, Matrix.changeBoundsUnchecked(dF));
050    }
051
052    return result;
053  }
054
055  /**
056   * Returns numerical Jacobian with respect to x for f(x, u, ...).
057   *
058   * @param <Rows> Number of rows in the result of f(x, u).
059   * @param <States> Number of rows in x.
060   * @param <Inputs> Number of rows in the second input to f.
061   * @param <Outputs> Num representing the rows in the output of f.
062   * @param rows Number of rows in the result of f(x, u).
063   * @param states Number of rows in x.
064   * @param f Vector-valued function from which to compute Jacobian.
065   * @param x State vector.
066   * @param u Input vector.
067   * @return The numerical Jacobian with respect to x for f(x, u, ...).
068   */
069  public static <Rows extends Num, States extends Num, Inputs extends Num, Outputs extends Num>
070      Matrix<Rows, States> numericalJacobianX(
071          Nat<Rows> rows,
072          Nat<States> states,
073          BiFunction<Matrix<States, N1>, Matrix<Inputs, N1>, Matrix<Outputs, N1>> f,
074          Matrix<States, N1> x,
075          Matrix<Inputs, N1> u) {
076    return numericalJacobian(rows, states, _x -> f.apply(_x, u), x);
077  }
078
079  /**
080   * Returns the numerical Jacobian with respect to u for f(x, u).
081   *
082   * @param <States> The states of the system.
083   * @param <Inputs> The inputs to the system.
084   * @param <Rows> Number of rows in the result of f(x, u).
085   * @param rows Number of rows in the result of f(x, u).
086   * @param inputs Number of rows in u.
087   * @param f Vector-valued function from which to compute the Jacobian.
088   * @param x State vector.
089   * @param u Input vector.
090   * @return the numerical Jacobian with respect to u for f(x, u).
091   */
092  public static <Rows extends Num, States extends Num, Inputs extends Num>
093      Matrix<Rows, Inputs> numericalJacobianU(
094          Nat<Rows> rows,
095          Nat<Inputs> inputs,
096          BiFunction<Matrix<States, N1>, Matrix<Inputs, N1>, Matrix<States, N1>> f,
097          Matrix<States, N1> x,
098          Matrix<Inputs, N1> u) {
099    return numericalJacobian(rows, inputs, _u -> f.apply(x, _u), u);
100  }
101}