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;
006
007import org.ejml.simple.SimpleMatrix;
008
009public final class Drake {
010  private Drake() {}
011
012  /**
013   * Solves the discrete alegebraic Riccati equation.
014   *
015   * @param A System matrix.
016   * @param B Input matrix.
017   * @param Q State cost matrix.
018   * @param R Input cost matrix.
019   * @return Solution of DARE.
020   */
021  @SuppressWarnings({"LocalVariableName", "ParameterName"})
022  public static SimpleMatrix discreteAlgebraicRiccatiEquation(
023      SimpleMatrix A, SimpleMatrix B, SimpleMatrix Q, SimpleMatrix R) {
024    var S = new SimpleMatrix(A.numRows(), A.numCols());
025    WPIMathJNI.discreteAlgebraicRiccatiEquation(
026        A.getDDRM().getData(),
027        B.getDDRM().getData(),
028        Q.getDDRM().getData(),
029        R.getDDRM().getData(),
030        A.numCols(),
031        B.numCols(),
032        S.getDDRM().getData());
033    return S;
034  }
035
036  /**
037   * Solves the discrete alegebraic Riccati equation.
038   *
039   * @param <States> Number of states.
040   * @param <Inputs> Number of inputs.
041   * @param A System matrix.
042   * @param B Input matrix.
043   * @param Q State cost matrix.
044   * @param R Input cost matrix.
045   * @return Solution of DARE.
046   */
047  @SuppressWarnings({"ParameterName", "MethodTypeParameterName"})
048  public static <States extends Num, Inputs extends Num>
049      Matrix<States, States> discreteAlgebraicRiccatiEquation(
050          Matrix<States, States> A,
051          Matrix<States, Inputs> B,
052          Matrix<States, States> Q,
053          Matrix<Inputs, Inputs> R) {
054    return new Matrix<>(
055        discreteAlgebraicRiccatiEquation(
056            A.getStorage(), B.getStorage(), Q.getStorage(), R.getStorage()));
057  }
058
059  /**
060   * Solves the discrete alegebraic Riccati equation.
061   *
062   * @param A System matrix.
063   * @param B Input matrix.
064   * @param Q State cost matrix.
065   * @param R Input cost matrix.
066   * @param N State-input cross-term cost matrix.
067   * @return Solution of DARE.
068   */
069  @SuppressWarnings({"LocalVariableName", "ParameterName"})
070  public static SimpleMatrix discreteAlgebraicRiccatiEquation(
071      SimpleMatrix A, SimpleMatrix B, SimpleMatrix Q, SimpleMatrix R, SimpleMatrix N) {
072    // See
073    // https://en.wikipedia.org/wiki/Linear%E2%80%93quadratic_regulator#Infinite-horizon,_discrete-time_LQR
074    // for the change of variables used here.
075    var scrA = A.minus(B.mult(R.solve(N.transpose())));
076    var scrQ = Q.minus(N.mult(R.solve(N.transpose())));
077
078    var S = new SimpleMatrix(A.numRows(), A.numCols());
079    WPIMathJNI.discreteAlgebraicRiccatiEquation(
080        scrA.getDDRM().getData(),
081        B.getDDRM().getData(),
082        scrQ.getDDRM().getData(),
083        R.getDDRM().getData(),
084        A.numCols(),
085        B.numCols(),
086        S.getDDRM().getData());
087    return S;
088  }
089
090  /**
091   * Solves the discrete alegebraic Riccati equation.
092   *
093   * @param <States> Number of states.
094   * @param <Inputs> Number of inputs.
095   * @param A System matrix.
096   * @param B Input matrix.
097   * @param Q State cost matrix.
098   * @param R Input cost matrix.
099   * @param N State-input cross-term cost matrix.
100   * @return Solution of DARE.
101   */
102  @SuppressWarnings({"ParameterName", "MethodTypeParameterName"})
103  public static <States extends Num, Inputs extends Num>
104      Matrix<States, States> discreteAlgebraicRiccatiEquation(
105          Matrix<States, States> A,
106          Matrix<States, Inputs> B,
107          Matrix<States, States> Q,
108          Matrix<Inputs, Inputs> R,
109          Matrix<States, Inputs> N) {
110    // See
111    // https://en.wikipedia.org/wiki/Linear%E2%80%93quadratic_regulator#Infinite-horizon,_discrete-time_LQR
112    // for the change of variables used here.
113    var scrA = A.minus(B.times(R.solve(N.transpose())));
114    var scrQ = Q.minus(N.times(R.solve(N.transpose())));
115
116    return new Matrix<>(
117        discreteAlgebraicRiccatiEquation(
118            scrA.getStorage(), B.getStorage(), scrQ.getStorage(), R.getStorage()));
119  }
120}