diff --git a/wpimath/src/main/native/include/frc/fmt/Eigen.h b/wpimath/src/main/native/include/frc/fmt/Eigen.h index f6cc7b6cc5..c6b2ee6afe 100644 --- a/wpimath/src/main/native/include/frc/fmt/Eigen.h +++ b/wpimath/src/main/native/include/frc/fmt/Eigen.h @@ -7,9 +7,10 @@ #include #include "Eigen/Core" +#include "Eigen/SparseCore" /** - * Formatter for Eigen::Matrix<>. + * Formatter for Eigen::Matrix. * * @tparam Rows Number of rows. * @tparam Cols Number of columns. @@ -51,12 +52,68 @@ struct fmt::formatter> { auto format(const Eigen::Matrix& mat, FormatContext& ctx) { auto out = ctx.out(); - for (int i = 0; i < Rows; ++i) { - for (int j = 0; j < Cols; ++j) { + for (int i = 0; i < mat.rows(); ++i) { + for (int j = 0; j < mat.cols(); ++j) { out = fmt::format_to(out, " {:f}", mat(i, j)); } - if (i < Rows - 1) { + if (i < mat.rows() - 1) { + out = fmt::format_to(out, "\n"); + } + } + + return out; + } +}; + +/** + * Formatter for Eigen::SparseMatrix. + * + * @tparam Options Union of bit flags controlling the storage scheme. + * @tparam StorageIndex The type of the indices. + */ +template +struct fmt::formatter> { + /** + * Storage for format specifier. + */ + char presentation = 'f'; + + /** + * Format string parser. + * + * @param ctx Format string context. + */ + constexpr auto parse(fmt::format_parse_context& ctx) { + auto it = ctx.begin(), end = ctx.end(); + if (it != end && (*it == 'f' || *it == 'e')) { + presentation = *it++; + } + + if (it != end && *it != '}') { + throw fmt::format_error("invalid format"); + } + + return it; + } + + /** + * Writes out a formatted matrix. + * + * @tparam FormatContext Format string context type. + * @param mat Matrix to format. + * @param ctx Format string context. + */ + template + auto format(const Eigen::SparseMatrix& mat, + FormatContext& ctx) { + auto out = ctx.out(); + for (int i = 0; i < mat.rows(); ++i) { + for (int j = 0; j < mat.cols(); ++j) { + out = fmt::format_to(out, " {:f}", mat.coeff(i, j)); + } + + if (i < mat.rows() - 1) { out = fmt::format_to(out, "\n"); } } diff --git a/wpimath/src/test/native/cpp/FormatterTest.cpp b/wpimath/src/test/native/cpp/FormatterTest.cpp index cd7ef5cbe4..0e2f77c25c 100644 --- a/wpimath/src/test/native/cpp/FormatterTest.cpp +++ b/wpimath/src/test/native/cpp/FormatterTest.cpp @@ -2,6 +2,8 @@ // 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 + #include #include "frc/fmt/Eigen.h" @@ -10,12 +12,33 @@ #include "units/velocity.h" TEST(FormatterTest, Eigen) { - Eigen::Matrix A{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}}; + Eigen::Matrix A{{0.0, 1.0}, {2.0, 3.0}, {4.0, 5.0}}; EXPECT_EQ( - " 1.000000 2.000000\n" - " 3.000000 4.000000\n" - " 5.000000 6.000000", + " 0.000000 1.000000\n" + " 2.000000 3.000000\n" + " 4.000000 5.000000", fmt::format("{}", A)); + + Eigen::MatrixXd B{{0.0, 1.0}, {2.0, 3.0}, {4.0, 5.0}}; + EXPECT_EQ( + " 0.000000 1.000000\n" + " 2.000000 3.000000\n" + " 4.000000 5.000000", + fmt::format("{}", B)); + + Eigen::SparseMatrix C{3, 2}; + std::vector> triplets; + triplets.emplace_back(0, 1, 1.0); + triplets.emplace_back(1, 0, 2.0); + triplets.emplace_back(1, 1, 3.0); + triplets.emplace_back(2, 0, 4.0); + triplets.emplace_back(2, 1, 5.0); + C.setFromTriplets(triplets.begin(), triplets.end()); + EXPECT_EQ( + " 0.000000 1.000000\n" + " 2.000000 3.000000\n" + " 4.000000 5.000000", + fmt::format("{}", C)); } TEST(FormatterTest, Units) {