[wpimath] Make cost and covariance matrix functions constexpr (wpilib…

…suite#6444)
spacey-sooty · May 7, 2024 · dc00a13 · dc00a13
1 parent bdc7344
commit dc00a13
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 30 deletions.
diff --git a/wpimath/src/main/native/include/frc/StateSpaceUtil.h b/wpimath/src/main/native/include/frc/StateSpaceUtil.h
@@ -34,18 +34,28 @@ namespace frc {
  * @return State excursion or control effort cost matrix.
  */
 template <std::same_as<double>... Ts>
-Matrixd<sizeof...(Ts), sizeof...(Ts)> MakeCostMatrix(Ts... tolerances) {
-  Eigen::DiagonalMatrix<double, sizeof...(Ts)> result;
-  auto& diag = result.diagonal();
+constexpr Matrixd<sizeof...(Ts), sizeof...(Ts)> MakeCostMatrix(
+    Ts... tolerances) {
+  Matrixd<sizeof...(Ts), sizeof...(Ts)> result;
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      if (row != col) {
+        result.coeffRef(row, col) = 0.0;
+      }
+    }
+  }
+
   wpi::for_each(
       [&](int i, double tolerance) {
         if (tolerance == std::numeric_limits<double>::infinity()) {
-          diag(i) = 0.0;
+          result.coeffRef(i, i) = 0.0;
         } else {
-          diag(i) = 1.0 / std::pow(tolerance, 2);
+          result.coeffRef(i, i) = 1.0 / (tolerance * tolerance);
         }
       },
       tolerances...);
+
   return result;
 }
 
@@ -62,11 +72,21 @@ Matrixd<sizeof...(Ts), sizeof...(Ts)> MakeCostMatrix(Ts... tolerances) {
  * @return Process noise or measurement noise covariance matrix.
  */
 template <std::same_as<double>... Ts>
-Matrixd<sizeof...(Ts), sizeof...(Ts)> MakeCovMatrix(Ts... stdDevs) {
-  Eigen::DiagonalMatrix<double, sizeof...(Ts)> result;
-  auto& diag = result.diagonal();
-  wpi::for_each([&](int i, double stdDev) { diag(i) = std::pow(stdDev, 2); },
-                stdDevs...);
+constexpr Matrixd<sizeof...(Ts), sizeof...(Ts)> MakeCovMatrix(Ts... stdDevs) {
+  Matrixd<sizeof...(Ts), sizeof...(Ts)> result;
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      if (row != col) {
+        result.coeffRef(row, col) = 0.0;
+      }
+    }
+  }
+
+  wpi::for_each(
+      [&](int i, double stdDev) { result.coeffRef(i, i) = stdDev * stdDev; },
+      stdDevs...);
+
   return result;
 }
 
@@ -84,16 +104,23 @@ Matrixd<sizeof...(Ts), sizeof...(Ts)> MakeCovMatrix(Ts... stdDevs) {
  * @return State excursion or control effort cost matrix.
  */
 template <size_t N>
-Matrixd<N, N> MakeCostMatrix(const std::array<double, N>& costs) {
-  Eigen::DiagonalMatrix<double, N> result;
-  auto& diag = result.diagonal();
-  for (size_t i = 0; i < costs.size(); ++i) {
-    if (costs[i] == std::numeric_limits<double>::infinity()) {
-      diag(i) = 0.0;
-    } else {
-      diag(i) = 1.0 / std::pow(costs[i], 2);
+constexpr Matrixd<N, N> MakeCostMatrix(const std::array<double, N>& costs) {
+  Matrixd<N, N> result;
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      if (row == col) {
+        if (costs[row] == std::numeric_limits<double>::infinity()) {
+          result.coeffRef(row, col) = 0.0;
+        } else {
+          result.coeffRef(row, col) = 1.0 / (costs[row] * costs[row]);
+        }
+      } else {
+        result.coeffRef(row, col) = 0.0;
+      }
     }
   }
+
   return result;
 }
 
@@ -110,12 +137,19 @@ Matrixd<N, N> MakeCostMatrix(const std::array<double, N>& costs) {
  * @return Process noise or measurement noise covariance matrix.
  */
 template <size_t N>
-Matrixd<N, N> MakeCovMatrix(const std::array<double, N>& stdDevs) {
-  Eigen::DiagonalMatrix<double, N> result;
-  auto& diag = result.diagonal();
-  for (size_t i = 0; i < N; ++i) {
-    diag(i) = std::pow(stdDevs[i], 2);
+constexpr Matrixd<N, N> MakeCovMatrix(const std::array<double, N>& stdDevs) {
+  Matrixd<N, N> result;
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      if (row == col) {
+        result.coeffRef(row, col) = stdDevs[row] * stdDevs[row];
+      } else {
+        result.coeffRef(row, col) = 0.0;
+      }
+    }
   }
+
   return result;
 }
 

diff --git a/wpimath/src/test/native/cpp/StateSpaceUtilTest.cpp b/wpimath/src/test/native/cpp/StateSpaceUtilTest.cpp
@@ -2,16 +2,13 @@
 // Open Source Software; you can modify and/or share it under the terms of
 // the WPILib BSD license file in the root directory of this project.
 
-#include <array>
-
 #include <gtest/gtest.h>
 
 #include "frc/EigenCore.h"
 #include "frc/StateSpaceUtil.h"
-#include "frc/system/NumericalIntegration.h"
 
 TEST(StateSpaceUtilTest, CostParameterPack) {
-  frc::Matrixd<3, 3> mat = frc::MakeCostMatrix(1.0, 2.0, 3.0);
+  constexpr frc::Matrixd<3, 3> mat = frc::MakeCostMatrix(1.0, 2.0, 3.0);
   EXPECT_NEAR(mat(0, 0), 1.0, 1e-3);
   EXPECT_NEAR(mat(0, 1), 0.0, 1e-3);
   EXPECT_NEAR(mat(0, 2), 0.0, 1e-3);
@@ -24,7 +21,7 @@ TEST(StateSpaceUtilTest, CostParameterPack) {
 }
 
 TEST(StateSpaceUtilTest, CostArray) {
-  frc::Matrixd<3, 3> mat = frc::MakeCostMatrix<3>({1.0, 2.0, 3.0});
+  constexpr frc::Matrixd<3, 3> mat = frc::MakeCostMatrix<3>({1.0, 2.0, 3.0});
   EXPECT_NEAR(mat(0, 0), 1.0, 1e-3);
   EXPECT_NEAR(mat(0, 1), 0.0, 1e-3);
   EXPECT_NEAR(mat(0, 2), 0.0, 1e-3);
@@ -37,7 +34,7 @@ TEST(StateSpaceUtilTest, CostArray) {
 }
 
 TEST(StateSpaceUtilTest, CovParameterPack) {
-  frc::Matrixd<3, 3> mat = frc::MakeCovMatrix(1.0, 2.0, 3.0);
+  constexpr frc::Matrixd<3, 3> mat = frc::MakeCovMatrix(1.0, 2.0, 3.0);
   EXPECT_NEAR(mat(0, 0), 1.0, 1e-3);
   EXPECT_NEAR(mat(0, 1), 0.0, 1e-3);
   EXPECT_NEAR(mat(0, 2), 0.0, 1e-3);
@@ -50,7 +47,7 @@ TEST(StateSpaceUtilTest, CovParameterPack) {
 }
 
 TEST(StateSpaceUtilTest, CovArray) {
-  frc::Matrixd<3, 3> mat = frc::MakeCovMatrix<3>({1.0, 2.0, 3.0});
+  constexpr frc::Matrixd<3, 3> mat = frc::MakeCovMatrix<3>({1.0, 2.0, 3.0});
   EXPECT_NEAR(mat(0, 0), 1.0, 1e-3);
   EXPECT_NEAR(mat(0, 1), 0.0, 1e-3);
   EXPECT_NEAR(mat(0, 2), 0.0, 1e-3);