Refactor scalability benchmark (#200)

benchmarks/scalability/CartPole/Main.cpp

@@ -1,12 +1,15 @@
 // Copyright (c) Sleipnir contributors
 
+#include <string_view>
 #include <vector>
 
 #include "CasADi.hpp"
 #include "Sleipnir.hpp"
 #include "Util.hpp"
 
-int main() {
+int main(int argc, char* argv[]) {
+  std::vector<std::string_view> args{argv + 1, argv + argc};
+
   constexpr bool diagnostics = false;
 
   constexpr auto T = 5_s;
@@ -19,7 +22,16 @@ int main() {
 
   fmt::print("Solving cart-pole problem from N = {} to N = {}.\n",
              sampleSizesToTest.front(), sampleSizesToTest.back());
-  return RunBenchmarksAndLog("cart-pole-scalability-results.csv", diagnostics,
-                             T, sampleSizesToTest, &CartPoleCasADi,
-                             &CartPoleSleipnir);
+  if (args.size() == 0 ||
+      std::find(args.begin(), args.end(), "--casadi") != args.end()) {
+    RunBenchmarksAndLog<casadi::Opti>(
+        "cart-pole-scalability-results-casadi.csv", diagnostics, T,
+        sampleSizesToTest, &CartPoleCasADi);
+  }
+  if (args.size() == 0 ||
+      std::find(args.begin(), args.end(), "--sleipnir") != args.end()) {
+    RunBenchmarksAndLog<sleipnir::OptimizationProblem>(
+        "cart-pole-scalability-results-sleipnir.csv", diagnostics, T,
+        sampleSizesToTest, &CartPoleSleipnir);
+  }
 }

benchmarks/scalability/Flywheel/Main.cpp

@@ -1,12 +1,16 @@
 // Copyright (c) Sleipnir contributors
 
+#include <algorithm>
+#include <string_view>
 #include <vector>
 
 #include "CasADi.hpp"
 #include "Sleipnir.hpp"
 #include "Util.hpp"
 
-int main() {
+int main(int argc, char* argv[]) {
+  std::vector<std::string_view> args{argv + 1, argv + argc};
+
   constexpr bool diagnostics = false;
 
   constexpr auto T = 5_s;
@@ -22,7 +26,16 @@ int main() {
 
   fmt::print("Solving flywheel problem from N = {} to N = {}.\n",
              sampleSizesToTest.front(), sampleSizesToTest.back());
-  return RunBenchmarksAndLog("flywheel-scalability-results.csv", diagnostics, T,
-                             sampleSizesToTest, &FlywheelCasADi,
-                             &FlywheelSleipnir);
+  if (args.size() == 0 ||
+      std::find(args.begin(), args.end(), "--casadi") != args.end()) {
+    RunBenchmarksAndLog<casadi::Opti>("flywheel-scalability-results-casadi.csv",
+                                      diagnostics, T, sampleSizesToTest,
+                                      &FlywheelCasADi);
+  }
+  if (args.size() == 0 ||
+      std::find(args.begin(), args.end(), "--sleipnir") != args.end()) {
+    RunBenchmarksAndLog<sleipnir::OptimizationProblem>(
+        "flywheel-scalability-results-sleipnir.csv", diagnostics, T,
+        sampleSizesToTest, &FlywheelSleipnir);
+  }
 }

benchmarks/scalability/Util.hpp

@@ -3,12 +3,15 @@
 #pragma once
 
 #include <chrono>
+#include <concepts>
 #include <fstream>
 #include <functional>
 #include <span>
+#include <string>
 #include <string_view>
 
 #include <casadi/casadi.hpp>
+#include <fmt/core.h>
 #include <sleipnir/optimization/OptimizationProblem.hpp>
 #include <units/time.h>
 
@@ -65,6 +68,8 @@ void RunBenchmark(std::ofstream& results, std::function<Problem()> setup,
  * The scale of the problem is iteratively increased by increasing the number of
  * timesteps within the time horizon.
  *
+ * @tparam Problem The optimization problem's type (casadi::Opti or
+ *   sleipnir::OptimizationProblem).
  * @param filename Results CSV filename.
  * @param diagnostics Whether to enable diagnostic prints.
  * @param T The time horizon of the optimization problem.
@@ -73,14 +78,49 @@ void RunBenchmark(std::ofstream& results, std::function<Problem()> setup,
  *   problem.
  * @param maxPower The maximum power of 10 for the number of samples in the
  *   problem.
- * @param casadiSetup A function that takes a time horizon and number of samples
- *   and returns a CasADi optimization problem instance.
  * @param sleipnirSetup A function that takes a time horizon and number of
- *   samples and returns a Sleipnir optimization problem instance.
+ *   samples and returns an optimization problem instance.
  */
-int RunBenchmarksAndLog(
-    std::string_view filename, bool diagnostics, units::second_t T,
-    std::span<int> sampleSizesToTest,
-    std::function<casadi::Opti(units::second_t, int)> casadiSetup,
-    std::function<sleipnir::OptimizationProblem(units::second_t, int)>
-        sleipnirSetup);
+template <typename Problem>
+int RunBenchmarksAndLog(std::string_view filename, bool diagnostics,
+                        units::second_t T, std::span<int> sampleSizesToTest,
+                        std::function<Problem(units::second_t, int)> setup) {
+  std::ofstream results{std::string{filename}};
+  if (!results.is_open()) {
+    return 1;
+  }
+
+  results << "Samples,"
+          << "Setup time (ms),Solve time (ms)\n";
+  std::flush(results);
+
+  for (int N : sampleSizesToTest) {
+    results << N << ",";
+    std::flush(results);
+
+    units::second_t dt = T / N;
+
+    fmt::print(stderr, "N = {}...", N);
+    RunBenchmark<Problem>(
+        results, [=] { return setup(dt, N); },
+        [=](Problem& problem) {
+          if constexpr (std::same_as<Problem, casadi::Opti>) {
+            if (diagnostics) {
+              problem.solver("ipopt");
+            } else {
+              problem.solver("ipopt", {{"print_time", 0}},
+                             {{"print_level", 0}, {"sb", "yes"}});
+            }
+            problem.solve();
+          } else {
+            problem.Solve({.diagnostics = diagnostics});
+          }
+        });
+    fmt::print(stderr, " done.\n");
+
+    results << "\n";
+    std::flush(results);
+  }
+
+  return 0;
+}

tools/generate-scalability-results-sleipnir.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+./build/FlywheelScalabilityBenchmark --sleipnir
+./tools/plot_scalability_results.py \
+  --filenames \
+    flywheel-scalability-results-sleipnir.csv \
+  --labels \
+    "Sleipnir" \
+  --title Flywheel \
+  --noninteractive
+
+./build/CartPoleScalabilityBenchmark --sleipnir
+./tools/plot_scalability_results.py \
+  --filenames \
+    cart-pole-scalability-results-sleipnir.csv \
+  --labels \
+    "Sleipnir" \
+  --title Cart-pole \
+  --noninteractive

tools/generate-scalability-results.sh

@@ -1,6 +1,24 @@
 #!/bin/bash
-./build/FlywheelScalabilityBenchmark
-./tools/plot_scalability_results.py --filename flywheel-scalability-results.csv --title Flywheel --noninteractive
+./build/FlywheelScalabilityBenchmark --casadi
+./build/FlywheelScalabilityBenchmark --sleipnir
+./tools/plot_scalability_results.py \
+  --filenames \
+    flywheel-scalability-results-casadi.csv \
+    flywheel-scalability-results-sleipnir.csv \
+  --labels \
+    "CasADi + ipopt" \
+    "Sleipnir" \
+  --title Flywheel \
+  --noninteractive
 
-./build/CartPoleScalabilityBenchmark
-./tools/plot_scalability_results.py --filename cart-pole-scalability-results.csv --title Cart-pole --noninteractive
+./build/CartPoleScalabilityBenchmark --casadi
+./build/CartPoleScalabilityBenchmark --sleipnir
+./tools/plot_scalability_results.py \
+  --filenames \
+    cart-pole-scalability-results-casadi.csv \
+    cart-pole-scalability-results-sleipnir.csv \
+  --labels \
+    "CasADi + ipopt" \
+    "Sleipnir" \
+  --title Cart-pole \
+  --noninteractive