From 8c2cbe77e1ff319c825b9c908e4a31272c561b61 Mon Sep 17 00:00:00 2001
From: Tyler Veness <calcmogul@gmail.com>
Date: Thu, 9 Jan 2025 22:09:18 -0800
Subject: [PATCH] Make .spy file format more efficient (#678)

---
 include/sleipnir/util/Spy.hpp             | 142 +++++++++++++---------
 src/optimization/solver/InteriorPoint.cpp |  32 +++--
 src/optimization/solver/Newton.cpp        |  14 +--
 src/optimization/solver/SQP.cpp           |  23 ++--
 tools/spy.py                              | 127 +++++++------------
 5 files changed, 164 insertions(+), 174 deletions(-)

diff --git a/include/sleipnir/util/Spy.hpp b/include/sleipnir/util/Spy.hpp
index cb9b4e19..1abeb469 100644
--- a/include/sleipnir/util/Spy.hpp
+++ b/include/sleipnir/util/Spy.hpp
@@ -2,6 +2,9 @@
 
 #pragma once
 
+#include <stdint.h>
+
+#include <bit>
 #include <fstream>
 #include <string>
 #include <string_view>
@@ -9,81 +12,110 @@
 #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
diff --git a/src/optimization/solver/InteriorPoint.cpp b/src/optimization/solver/InteriorPoint.cpp
index 46b5c9b1..3f27847a 100644
--- a/src/optimization/solver/InteriorPoint.cpp
+++ b/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
diff --git a/src/optimization/solver/Newton.cpp b/src/optimization/solver/Newton.cpp
index 2f8c3516..2d27d0d0 100644
--- a/src/optimization/solver/Newton.cpp
+++ b/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
diff --git a/src/optimization/solver/SQP.cpp b/src/optimization/solver/SQP.cpp
index 38aa75fb..d2170aed 100644
--- a/src/optimization/solver/SQP.cpp
+++ b/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
diff --git a/tools/spy.py b/tools/spy.py
index 89889976..a950230b 100755
--- a/tools/spy.py
+++ b/tools/spy.py
@@ -1,8 +1,7 @@
 #!/usr/bin/env python3
 
-"""Loads and displays the sparsity patterns from A_e.spy, A_i.spy, and H.spy."""
-
-import re
+import argparse
+import struct
 import sys
 
 import matplotlib.pyplot as plt
@@ -12,18 +11,10 @@
 from matplotlib.colors import ListedColormap
 
 
-def plot_csv(filename, title, xlabel, ylabel):
-    print(f"Loading sparsity pattern for {title}...", end="")
+def plot_csv(filename: str) -> animation.FuncAnimation:
+    print(f"Loading sparsity pattern {filename}...", end="")
     sys.stdout.flush()
 
-    with open(filename) as f:
-        contents = f.read()
-
-    max_row_idx = 0
-    max_col_idx = 0
-    row = 0
-    col = 0
-
     xs = []
     ys = []
     vs = []
@@ -32,51 +23,40 @@ def plot_csv(filename, title, xlabel, ylabel):
     y = []
     v = []
 
-    for m in re.finditer(r"\.|\+|\-|\n\n|\n", contents):
-        token = m.group()
-
-        if token == ".":
-            # Do nothing since zero entries aren't logged
-            col += 1
-        elif token == "+":
-            # Log positive entry
-            x.append(col)
-            y.append(row)
-            v.append(1.0)
-            col += 1
-        elif token == "-":
-            # Log negative entry
-            x.append(col)
-            y.append(row)
-            v.append(-1.0)
-            col += 1
-        elif token == "\n\n":
-            # Prep for new matrix. Log old one if it's not truncated.
-            if row >= max_row_idx and col >= max_col_idx:
-                max_row_idx = row
-                max_col_idx = col
-
+    with open(filename, mode="rb") as f:
+        size = struct.unpack("<i", f.read(4))[0]
+        title = struct.unpack(f"{size}s", f.read(size))[0].decode("utf-8")
+
+        size = struct.unpack("<i", f.read(4))[0]
+        row_label = struct.unpack(f"{size}s", f.read(size))[0].decode("utf-8")
+
+        size = struct.unpack("<i", f.read(4))[0]
+        col_label = struct.unpack(f"{size}s", f.read(size))[0].decode("utf-8")
+
+        rows = struct.unpack("<i", f.read(4))[0]
+        cols = struct.unpack("<i", f.read(4))[0]
+
+        try:
+            while True:
+                num_coords = struct.unpack("<i", f.read(4))[0]
+                for _ in range(num_coords):
+                    y.append(struct.unpack("<i", f.read(4))[0])
+                    x.append(struct.unpack("<i", f.read(4))[0])
+                    sign = struct.unpack("c", f.read(1))[0]
+                    if sign == b"+":
+                        v.append(1.0)
+                    elif sign == b"-":
+                        v.append(-1.0)
+                    else:
+                        v.append(0.0)
                 xs.append(x)
                 ys.append(y)
                 vs.append(v)
-            x = []
-            y = []
-            v = []
-            row = 0
-            col = 0
-        elif token == "\n" and m.span()[0] < len(contents) - 1:
-            # Prep for new row if it isn't the last
-            row += 1
-            col = 0
-
-    # Log leftover matrix if it's not truncated
-    if row >= max_row_idx and col >= max_col_idx:
-        max_row_idx = row
-        max_col_idx = col
-
-        xs.append(x)
-        ys.append(y)
-        vs.append(v)
+                x = []
+                y = []
+                v = []
+        except:
+            pass
 
     print(" done.")
     sys.stdout.flush()
@@ -85,7 +65,7 @@ def plot_csv(filename, title, xlabel, ylabel):
     ax = fig.add_subplot()
 
     # Display scatter plot
-    cmap = ListedColormap(["red", "green"])
+    cmap = ListedColormap(["red", "black", "green"])
     sc = ax.scatter(xs[0], ys[0], s=1, c=vs[0], marker=".", cmap=cmap, vmin=-1, vmax=1)
     iteration = ax.text(0.02, 0.95, "", transform=ax.transAxes)
 
@@ -97,15 +77,15 @@ def plot_csv(filename, title, xlabel, ylabel):
     colorbar.ax.set_yticklabels(ticklabels)
 
     ax.set_title(title)
-    ax.set_xlabel(xlabel)
-    ax.set_ylabel(ylabel)
+    ax.set_xlabel(col_label)
+    ax.set_ylabel(row_label)
 
-    ax.set_xlim([0, max_col_idx])
-    ax.set_ylim([0, max_row_idx])
+    ax.set_xlim([0, cols])
+    ax.set_ylim([0, rows])
     ax.invert_yaxis()
     ax.set_aspect(1.0)
 
-    def animate(i):
+    def animate(i: int):
         sc.set_offsets(np.c_[xs[i], ys[i]])
         sc.set_array(vs[i])
         iteration.set_text(f"iter {i}/{max(0, len(vs) - 1)}")
@@ -123,27 +103,14 @@ def animate(i):
 
 
 def main():
-    # pragma pylint: disable=unused-variable
-    anim1 = plot_csv(
-        "A_e.spy",
-        title="Equality constraint Jacobian",
-        xlabel="Decision variables",
-        ylabel="Constraints",
-    )
-    # pragma pylint: disable=unused-variable
-    anim2 = plot_csv(
-        "A_i.spy",
-        title="Inequality constraint Jacobian",
-        xlabel="Decision variables",
-        ylabel="Constraints",
+    parser = argparse.ArgumentParser(
+        description="Displays sparsity pattern (.spy) files."
     )
+    parser.add_argument("filename", nargs="+")
+    args = parser.parse_args()
+
     # pragma pylint: disable=unused-variable
-    anim3 = plot_csv(
-        "H.spy",
-        title="Hessian",
-        xlabel="Decision variables",
-        ylabel="Decision variables",
-    )
+    animations = [plot_csv(filename) for filename in args.filename]  # noqa
     plt.show()