Make .spy file format more efficient (#678)

include/sleipnir/util/Spy.hpp

@@ -2,88 +2,120 @@
 
 #pragma once
 
+#include <stdint.h>
+
+#include <bit>
 #include <fstream>
 #include <string>
 #include <string_view>
 
 #include <Eigen/SparseCore>
 
 #include "sleipnir/util/SymbolExports.hpp"
-#include "sleipnir/util/small_vector.hpp"
 
 namespace sleipnir {
 
 /**
- * Write the sparsity pattern of a sparse matrix to a file.
+ * Writes the sparsity pattern of a sparse matrix to a file.
+ *
+ * Each file represents the sparsity pattern of one matrix over time. <a
+ * href="https://github.com/SleipnirGroup/Sleipnir/blob/main/tools/spy.py">spy.py</a>
+ * can display it as an animation.
  *
- * Each character represents an element with '.' representing zero, '+'
- * representing positive, and '-' representing negative. Here's an example for a
- * 3x3 identity matrix.
+ * The file starts with the following header:
+ * <ol>
+ *   <li>Plot title (length as a little-endian int32, then characters)</li>
+ *   <li>Row label (length as a little-endian int32, then characters)</li>
+ *   <li>Column label (length as a little-endian int32, then characters)</li>
+ * </ol>
  *
- * "+.."
- * ".+."
- * "..+"
+ * Then, each sparsity pattern starts with:
+ * <ol>
+ *   <li>Number of coordinates as a little-endian int32</li>
+ * </ol>
+ *
+ * followed by that many coordinates in the following format:
+ * <ol>
+ *   <li>Row index as a little-endian int32</li>
+ *   <li>Column index as a little-endian int32</li>
+ *   <li>Sign as a character ('+' for positive, '-' for negative, or '0' for
+ *       zero)</li>
+ * </ol>
  *
  * @param[out] file A file stream.
  * @param[in] mat The sparse matrix.
  */
-SLEIPNIR_DLLEXPORT inline void Spy(std::ostream& file,
-                                   const Eigen::SparseMatrix<double>& mat) {
-  const int cells_width = mat.cols() + 1;
-  const int cells_height = mat.rows();
+class SLEIPNIR_DLLEXPORT Spy {
+ public:
+  /**
+   * Constructs a Spy instance.
+   *
+   * @param filename The filename.
+   * @param title Plot title.
+   * @param rowLabel Row label.
+   * @param colLabel Column label.
+   * @param rows The sparse matrix's number of rows.
+   * @param cols The sparse matrix's number of columns.
+   */
+  Spy(std::string_view filename, std::string_view title,
+      std::string_view rowLabel, std::string_view colLabel, int rows, int cols)
+      : m_file{std::string{filename}, std::ios::binary} {
+    // Write title
+    Write32le(title.size());
+    m_file.write(title.data(), title.size());
 
-  small_vector<uint8_t> cells;
+    // Write row label
+    Write32le(rowLabel.size());
+    m_file.write(rowLabel.data(), rowLabel.size());
 
-  // Allocate space for matrix of characters plus trailing newlines
-  cells.reserve(cells_width * cells_height);
+    // Write column label
+    Write32le(colLabel.size());
+    m_file.write(colLabel.data(), colLabel.size());
 
-  // Initialize cell array
-  for (int row = 0; row < mat.rows(); ++row) {
-    for (int col = 0; col < mat.cols(); ++col) {
-      cells.emplace_back('.');
-    }
-    cells.emplace_back('\n');
+    // Write row and column counts
+    Write32le(rows);
+    Write32le(cols);
   }
 
-  // Fill in non-sparse entries
-  for (int k = 0; k < mat.outerSize(); ++k) {
-    for (Eigen::SparseMatrix<double>::InnerIterator it{mat, k}; it; ++it) {
-      if (it.value() < 0.0) {
-        cells[it.row() * cells_width + it.col()] = '-';
-      } else if (it.value() > 0.0) {
-        cells[it.row() * cells_width + it.col()] = '+';
+  /**
+   * Adds a matrix to the file.
+   *
+   * @param mat The matrix.
+   */
+  void Add(const Eigen::SparseMatrix<double>& mat) {
+    // Write number of coordinates
+    Write32le(mat.nonZeros());
+
+    // Write coordinates
+    for (int k = 0; k < mat.outerSize(); ++k) {
+      for (Eigen::SparseMatrix<double>::InnerIterator it{mat, k}; it; ++it) {
+        Write32le(it.row());
+        Write32le(it.col());
+        if (it.value() > 0.0) {
+          m_file << '+';
+        } else if (it.value() < 0.0) {
+          m_file << '-';
+        } else {
+          m_file << '0';
+        }
       }
     }
   }
 
-  // Write cell array to file
-  for (const auto& c : cells) {
-    file << c;
-  }
-}
+ private:
+  std::ofstream m_file;
 
-/**
- * Write the sparsity pattern of a sparse matrix to a file.
- *
- * Each character represents an element with "." representing zero, "+"
- * representing positive, and "-" representing negative. Here's an example for a
- * 3x3 identity matrix.
- *
- * "+.."
- * ".+."
- * "..+"
- *
- * @param[in] filename The filename.
- * @param[in] mat The sparse matrix.
- */
-SLEIPNIR_DLLEXPORT inline void Spy(std::string_view filename,
-                                   const Eigen::SparseMatrix<double>& mat) {
-  std::ofstream file{std::string{filename}};
-  if (!file.is_open()) {
-    return;
+  /**
+   * Writes a 32-bit signed integer to the file as little-endian.
+   *
+   * @param num A 32-bit signed integer.
+   */
+  void Write32le(int32_t num) {
+    if constexpr (std::endian::native != std::endian::little) {
+      num = std::byteswap(num);
+    }
+    m_file.write(reinterpret_cast<char*>(&num), sizeof(num));
   }
-
-  Spy(file, mat);
-}
+};
 
 }  // namespace sleipnir

src/optimization/solver/InteriorPoint.cpp

@@ -5,8 +5,8 @@
 #include <algorithm>
 #include <chrono>
 #include <cmath>
-#include <fstream>
 #include <limits>
+#include <memory>
 
 #include <Eigen/SparseCholesky>
 
@@ -125,13 +125,18 @@ void InteriorPoint(std::span<Variable> decisionVariables,
   }
 
   // Sparsity pattern files written when spy flag is set in SolverConfig
-  std::ofstream H_spy;
-  std::ofstream A_e_spy;
-  std::ofstream A_i_spy;
+  std::unique_ptr<Spy> H_spy;
+  std::unique_ptr<Spy> A_e_spy;
+  std::unique_ptr<Spy> A_i_spy;
   if (config.spy) {
-    A_e_spy.open("A_e.spy");
-    A_i_spy.open("A_i.spy");
-    H_spy.open("H.spy");
+    H_spy = std::make_unique<Spy>("H.spy", "Hessian", "Decision variables",
+                                  "Decision variables", H.rows(), H.cols());
+    A_e_spy = std::make_unique<Spy>("A_e.spy", "Equality constraint Jacobian",
+                                    "Constraints", "Decision variables",
+                                    A_e.rows(), A_e.cols());
+    A_i_spy = std::make_unique<Spy>("A_i.spy", "Inequality constraint Jacobian",
+                                    "Constraints", "Decision variables",
+                                    A_i.rows(), A_i.cols());
   }
 
   if (config.diagnostics && !feasibilityRestoration) {
@@ -299,16 +304,9 @@ void InteriorPoint(std::span<Variable> decisionVariables,
 
     // Write out spy file contents if that's enabled
     if (config.spy) {
-      // Gap between sparsity patterns
-      if (iterations > 0) {
-        A_e_spy << "\n";
-        A_i_spy << "\n";
-        H_spy << "\n";
-      }
-
-      Spy(H_spy, H);
-      Spy(A_e_spy, A_e);
-      Spy(A_i_spy, A_i);
+      H_spy->Add(H);
+      A_e_spy->Add(A_e);
+      A_i_spy->Add(A_i);
     }
 
     // Call user callback

src/optimization/solver/Newton.cpp

@@ -5,8 +5,8 @@
 #include <algorithm>
 #include <chrono>
 #include <cmath>
-#include <fstream>
 #include <limits>
+#include <memory>
 
 #include <Eigen/SparseCholesky>
 
@@ -60,9 +60,10 @@ void Newton(std::span<Variable> decisionVariables, Variable& f,
   }
 
   // Sparsity pattern files written when spy flag is set in SolverConfig
-  std::ofstream H_spy;
+  std::unique_ptr<Spy> H_spy;
   if (config.spy) {
-    H_spy.open("H.spy");
+    H_spy = std::make_unique<Spy>("H.spy", "Hessian", "Decision variables",
+                                  "Decision variables", H.rows(), H.cols());
   }
 
   if (config.diagnostics) {
@@ -137,12 +138,7 @@ void Newton(std::span<Variable> decisionVariables, Variable& f,
 
     // Write out spy file contents if that's enabled
     if (config.spy) {
-      // Gap between sparsity patterns
-      if (iterations > 0) {
-        H_spy << "\n";
-      }
-
-      Spy(H_spy, H);
+      H_spy->Add(H);
     }
 
     // Call user callback

src/optimization/solver/SQP.cpp

@@ -5,8 +5,8 @@
 #include <algorithm>
 #include <chrono>
 #include <cmath>
-#include <fstream>
 #include <limits>
+#include <memory>
 
 #include <Eigen/SparseCholesky>
 
@@ -100,11 +100,14 @@ void SQP(std::span<Variable> decisionVariables,
   }
 
   // Sparsity pattern files written when spy flag is set in SolverConfig
-  std::ofstream H_spy;
-  std::ofstream A_e_spy;
+  std::unique_ptr<Spy> H_spy;
+  std::unique_ptr<Spy> A_e_spy;
   if (config.spy) {
-    A_e_spy.open("A_e.spy");
-    H_spy.open("H.spy");
+    H_spy = std::make_unique<Spy>("H.spy", "Hessian", "Decision variables",
+                                  "Decision variables", H.rows(), H.cols());
+    A_e_spy = std::make_unique<Spy>("A_e.spy", "Equality constraint Jacobian",
+                                    "Constraints", "Decision variables",
+                                    A_e.rows(), A_e.cols());
   }
 
   if (config.diagnostics) {
@@ -207,14 +210,8 @@ void SQP(std::span<Variable> decisionVariables,
 
     // Write out spy file contents if that's enabled
     if (config.spy) {
-      // Gap between sparsity patterns
-      if (iterations > 0) {
-        A_e_spy << "\n";
-        H_spy << "\n";
-      }
-
-      Spy(H_spy, H);
-      Spy(A_e_spy, A_e);
+      H_spy->Add(H);
+      A_e_spy->Add(A_e);
     }
 
     // Call user callback