From 38268b4aad03e6ce4755315f4cd690f007fa2720 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Max=20L=C3=BCbke?= <mluebke@uni-potsdam.de>
Date: Wed, 4 Dec 2024 09:06:40 +0000
Subject: [PATCH] feat: add support for ignoring specified cells in simulation
 run

---
 poet/include/PhreeqcRunner.hpp  | 12 ++++++++++++
 poet/src/Runner.cpp             | 28 ++++++++++++++++++++++++++++
 poet/test/testPhreeqcRunner.cpp | 22 ++++++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/poet/include/PhreeqcRunner.hpp b/poet/include/PhreeqcRunner.hpp
index 5c5230a3..ef95a5c3 100644
--- a/poet/include/PhreeqcRunner.hpp
+++ b/poet/include/PhreeqcRunner.hpp
@@ -50,6 +50,18 @@ public:
   void run(std::vector<std::vector<double>> &simulationInOut,
            const double time_step);
 
+  /**
+   * @brief Runs the simulation with the given input and output data.
+   *
+   * @param simulationInOut A reference to a 2D vector containing the simulation
+   * input and output data.
+   * @param time_step The time step for the simulation.
+   * @param to_ignore A vector of indices specifying which elements to ignore
+   * during the simulation.
+   */
+  void run(std::vector<std::vector<double>> &simulationInOut,
+           const double time_step, const std::vector<std::size_t> &to_ignore);
+
   /**
    * @brief Returns the number of engines currently stored.
    *
diff --git a/poet/src/Runner.cpp b/poet/src/Runner.cpp
index 44913b20..3d9e7cf6 100644
--- a/poet/src/Runner.cpp
+++ b/poet/src/Runner.cpp
@@ -1,9 +1,11 @@
 #include "PhreeqcEngine.hpp"
 #include "PhreeqcMatrix.hpp"
 #include "PhreeqcRunner.hpp"
+#include <algorithm>
 #include <cmath>
 #include <cstddef>
 #include <memory>
+#include <set>
 #include <vector>
 
 PhreeqcRunner::PhreeqcRunner(const PhreeqcMatrix &matrix) {
@@ -50,6 +52,32 @@ void PhreeqcRunner::run(std::vector<std::vector<double>> &simulationInOut,
 
     this->_engineStorage.at(pqc_id)->runCell(this->_buffer, time_step);
 
+    // Copy the buffer back to the output while ignoring the first element and
+    // NaNs
+    copy_from_buffer(this->_buffer, simulationInOut[i]);
+  }
+}
+
+void PhreeqcRunner::run(std::vector<std::vector<double>> &simulationInOut,
+                        const double time_step,
+                        const std::vector<std::size_t> &to_ignore) {
+  const std::set<std::size_t> to_ignore_set(to_ignore.begin(), to_ignore.end());
+
+  for (std::size_t i = 0; i < simulationInOut.size(); i++) {
+    if (to_ignore_set.find(i) != to_ignore_set.end()) {
+      continue;
+    }
+
+    const auto &current_inout = simulationInOut[i];
+    const auto pqc_id = static_cast<int>(current_inout[0]);
+
+    this->_buffer.clear();
+
+    // Copy the input to the buffer while ignoring the first element and NaNs
+    copy_to_buffer(this->_buffer, current_inout);
+
+    this->_engineStorage.at(pqc_id)->runCell(this->_buffer, time_step);
+
     // Copy the buffer back to the output while ignoring the first element and
     // NaNs
     copy_from_buffer(this->_buffer, simulationInOut[i]);
diff --git a/poet/test/testPhreeqcRunner.cpp b/poet/test/testPhreeqcRunner.cpp
index 41f0d1d5..08a6e995 100644
--- a/poet/test/testPhreeqcRunner.cpp
+++ b/poet/test/testPhreeqcRunner.cpp
@@ -72,4 +72,26 @@ POET_TEST(PhreeqcRunnerUnknownID) {
   simulationInOut[0][0] = 1000;
 
   EXPECT_THROW(runner.run(simulationInOut, 100), std::out_of_range);
+}
+
+POET_TEST(PhreeqcRunnerSimulationWithIgnoredCells) {
+  PhreeqcMatrix pqc_mat(test_database, test_script);
+  const auto subsetted_pqc_mat = pqc_mat.subset({2, 3});
+  PhreeqcRunner runner(subsetted_pqc_mat);
+
+  const auto stl_mat = subsetted_pqc_mat.get();
+  const auto matrix_values = stl_mat.values;
+  const auto num_columns = stl_mat.names.size();
+
+  std::vector<double> second_line(matrix_values.begin() + num_columns,
+                                  matrix_values.end());
+  std::vector<std::vector<double>> simulationInOut;
+
+  simulationInOut.push_back(second_line);
+
+  EXPECT_NO_THROW(runner.run(simulationInOut, 10000, {0}));
+
+  for (std::size_t i = 0; i < num_columns; ++i) {
+    EXPECT_DOUBLE_EQ(simulationInOut[0][i], second_line[i]);
+  }
 }
\ No newline at end of file