[wpimath] Fix LQR matrix constructor overload for Q, R, and N (#3884)

It was using the continuous B matrix to compute the feedback gain
instead of the discrete B matrix.

Tests were added for the matrix constructor overloads.

wpimath/src/test/java/edu/wpi/first/math/controller/LinearQuadraticRegulatorTest.java

@@ -6,7 +6,11 @@ package edu.wpi.first.math.controller;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
+import edu.wpi.first.math.Matrix;
+import edu.wpi.first.math.Nat;
+import edu.wpi.first.math.Num;
 import edu.wpi.first.math.VecBuilder;
+import edu.wpi.first.math.system.Discretization;
 import edu.wpi.first.math.system.plant.DCMotor;
 import edu.wpi.first.math.system.plant.LinearSystemId;
 import org.junit.jupiter.api.Test;
@@ -27,15 +31,14 @@ class LinearQuadraticRegulatorTest {
     var rElms = VecBuilder.fill(12.0);
     var dt = 0.00505;
 
-    var controller = new LinearQuadraticRegulator<>(plant, qElms, rElms, dt);
+    var K = new LinearQuadraticRegulator<>(plant, qElms, rElms, dt).getK();
 
-    var k = controller.getK();
-
-    assertEquals(522.153, k.get(0, 0), 0.1);
-    assertEquals(38.2, k.get(0, 1), 0.1);
+    assertEquals(522.153, K.get(0, 0), 0.1);
+    assertEquals(38.2, K.get(0, 1), 0.1);
   }
 
   @Test
+  @SuppressWarnings("LocalVariableName")
   void testFourMotorElevator() {
     var dt = 0.020;
 
@@ -43,11 +46,12 @@ class LinearQuadraticRegulatorTest {
         LinearSystemId.createElevatorSystem(
             DCMotor.getVex775Pro(4), 8.0, 0.75 * 25.4 / 1000.0, 14.67);
 
-    var regulator =
-        new LinearQuadraticRegulator<>(plant, VecBuilder.fill(0.1, 0.2), VecBuilder.fill(12.0), dt);
+    var K =
+        new LinearQuadraticRegulator<>(plant, VecBuilder.fill(0.1, 0.2), VecBuilder.fill(12.0), dt)
+            .getK();
 
-    assertEquals(10.381, regulator.getK().get(0, 0), 1e-2);
-    assertEquals(0.6929, regulator.getK().get(0, 1), 1e-2);
+    assertEquals(10.381, K.get(0, 0), 1e-2);
+    assertEquals(0.6929, K.get(0, 1), 1e-2);
   }
 
   @Test
@@ -65,12 +69,95 @@ class LinearQuadraticRegulatorTest {
     var rElms = VecBuilder.fill(12.0);
     var dt = 0.00505;
 
-    var controller = new LinearQuadraticRegulator<>(plant, qElms, rElms, dt);
+    var K = new LinearQuadraticRegulator<>(plant, qElms, rElms, dt).getK();
 
-    var k = controller.getK();
+    assertEquals(19.16, K.get(0, 0), 0.1);
+    assertEquals(3.32, K.get(0, 1), 0.1);
+  }
 
-    assertEquals(19.16, k.get(0, 0), 0.1);
-    assertEquals(3.32, k.get(0, 1), 0.1);
+  /**
+   * Returns feedback control gain for implicit model following.
+   *
+   * <p>This is used to test the QRN overload of LQR.
+   *
+   * @param States Number of states.
+   * @param Inputs Number of inputs.
+   * @param A State matrix.
+   * @param B Input matrix.
+   * @param Q State cost matrix.
+   * @param R Input cost matrix.
+   * @param Aref Desired state matrix.
+   * @param dtSeconds Discretization timestep in seconds.
+   */
+  @SuppressWarnings({"LocalVariableName", "MethodTypeParameterName", "ParameterName"})
+  <States extends Num, Inputs extends Num> Matrix<Inputs, States> getImplicitModelFollowingK(
+      Matrix<States, States> A,
+      Matrix<States, Inputs> B,
+      Matrix<States, States> Q,
+      Matrix<Inputs, Inputs> R,
+      Matrix<States, States> Aref,
+      double dtSeconds) {
+    // Discretize real dynamics
+    var discABPair = Discretization.discretizeAB(A, B, dtSeconds);
+    var discA = discABPair.getFirst();
+    var discB = discABPair.getSecond();
+
+    // Discretize desired dynamics
+    var discAref = Discretization.discretizeA(Aref, dtSeconds);
+
+    Matrix<States, States> Qimf =
+        (discA.minus(discAref)).transpose().times(Q).times(discA.minus(discAref));
+    Matrix<Inputs, Inputs> Rimf = discB.transpose().times(Q).times(discB).plus(R);
+    Matrix<States, Inputs> Nimf = (discA.minus(discAref)).transpose().times(Q).times(discB);
+
+    return new LinearQuadraticRegulator<>(A, B, Qimf, Rimf, Nimf, dtSeconds).getK();
+  }
+
+  @Test
+  @SuppressWarnings("LocalVariableName")
+  void testMatrixOverloadsWithSingleIntegrator() {
+    var A = Matrix.mat(Nat.N2(), Nat.N2()).fill(0, 0, 0, 0);
+    var B = Matrix.mat(Nat.N2(), Nat.N2()).fill(1, 0, 0, 1);
+    var Q = Matrix.mat(Nat.N2(), Nat.N2()).fill(1, 0, 0, 1);
+    var R = Matrix.mat(Nat.N2(), Nat.N2()).fill(1, 0, 0, 1);
+
+    // QR overload
+    var K = new LinearQuadraticRegulator<>(A, B, Q, R, 0.005).getK();
+    assertEquals(0.99750312499512261, K.get(0, 0), 1e-10);
+    assertEquals(0.0, K.get(0, 1), 1e-10);
+    assertEquals(0.0, K.get(1, 0), 1e-10);
+    assertEquals(0.99750312499512261, K.get(1, 1), 1e-10);
+
+    // QRN overload
+    var N = Matrix.mat(Nat.N2(), Nat.N2()).fill(1, 0, 0, 1);
+    var Kimf = new LinearQuadraticRegulator<>(A, B, Q, R, N, 0.005).getK();
+    assertEquals(1.0, Kimf.get(0, 0), 1e-10);
+    assertEquals(0.0, Kimf.get(0, 1), 1e-10);
+    assertEquals(0.0, Kimf.get(1, 0), 1e-10);
+    assertEquals(1.0, Kimf.get(1, 1), 1e-10);
+  }
+
+  @Test
+  @SuppressWarnings("LocalVariableName")
+  void testMatrixOverloadsWithDoubleIntegrator() {
+    double Kv = 3.02;
+    double Ka = 0.642;
+
+    var A = Matrix.mat(Nat.N2(), Nat.N2()).fill(0, 1, 0, -Kv / Ka);
+    var B = Matrix.mat(Nat.N2(), Nat.N1()).fill(0, 1.0 / Ka);
+    var Q = Matrix.mat(Nat.N2(), Nat.N2()).fill(1, 0, 0, 0.2);
+    var R = Matrix.mat(Nat.N1(), Nat.N1()).fill(0.25);
+
+    // QR overload
+    var K = new LinearQuadraticRegulator<>(A, B, Q, R, 0.005).getK();
+    assertEquals(1.9960017786537287, K.get(0, 0), 1e-10);
+    assertEquals(0.51182128351092726, K.get(0, 1), 1e-10);
+
+    // QRN overload
+    var Aref = Matrix.mat(Nat.N2(), Nat.N2()).fill(0, 1, 0, -Kv / (Ka * 2.0));
+    var Kimf = getImplicitModelFollowingK(A, B, Q, R, Aref, 0.005);
+    assertEquals(0.0, Kimf.get(0, 0), 1e-10);
+    assertEquals(-5.367540084534802e-05, Kimf.get(0, 1), 1e-10);
   }
 
   @Test