Merge branch 'ml/optional-redox' into 'poet'

feat: Optional redox states of species in solution

See merge request naaice/iphreeqc!27

poet/include/PhreeqcMatrix.hpp

@@ -19,7 +19,8 @@ class IPhreeqc;
  * structures, minimizing the overhead of parsing and eliminate floating point
  * errors due to conversion.
  *
- * The class is also used to initialize the PhreeqcEngine class.
+ * The class is also used to initialize the PhreeqcEngine and PhreeqcRunner
+ * class.
  */
 class PhreeqcMatrix {
 public:
@@ -38,6 +39,20 @@ public:
    */
   ~PhreeqcMatrix() = default;
 
+  /**
+   * @brief Constructs a PhreeqcMatrix object with the specified database and
+   * input script.
+   *
+   * This constructor initializes a PhreeqcMatrix object using the provided
+   * database and input script. It sets the default values to exclude H(0) and
+   * O(0) in the output and to include redox states.
+   *
+   * @param database The path to the database file.
+   * @param input_script The input script to be used.
+   */
+  PhreeqcMatrix(const std::string &database, const std::string &input_script)
+      : PhreeqcMatrix(database, input_script, false, true) {}
+
   /**
    * @brief Construct a new Phreeqc Matrix object
    *
@@ -47,9 +62,10 @@ public:
    * @param database Phreeqc database as a string.
    * @param input_script Phreeqc input script as a string.
    * @param with_h0_o0 Whether to include H(0) and O(0) in the output or not.
+   * @param with_redox Whether to include redox states in the output or not.
    */
   PhreeqcMatrix(const std::string &database, const std::string &input_script,
-                bool with_h0_o0 = false);
+                bool with_h0_o0, bool with_redox);
 
   /**
    * @brief Construct a new Phreeqc Matrix object
@@ -287,6 +303,16 @@ public:
    */
   PhreeqcKnobs getKnobs() const { return *_m_knobs; }
 
+  /**
+   * @brief Checks if the redox states are included in solution.
+   *
+   * This function returns a boolean value indicating whether the redox states
+   * of a species are considered in the solution.
+   *
+   * @return true if redox states are included, false otherwise.
+   */
+  bool withRedox() const { return _m_with_redox; }
+
 private:
   std::map<int, std::vector<element>> _m_map;
   std::map<int, std::vector<base_names>> _m_internal_names;
@@ -303,4 +329,5 @@ private:
   std::string _m_database;
 
   bool _m_with_h0_o0;
+  bool _m_with_redox;
 };

poet/src/Engine.cpp

@@ -59,6 +59,7 @@ public:
   bool has_surface = false;
   struct InitCell {
     std::vector<std::string> solutions;
+    bool with_redox;
     std::vector<std::string> exchanger;
     std::vector<std::string> kinetics;
     std::vector<std::string> equilibrium;
@@ -96,6 +97,7 @@ PhreeqcEngine::Impl::Impl(const PhreeqcMatrix &pqc_mat, const int cell_id) {
   this->RunString(pqc_string.c_str());
 
   InitCell cell = {pqc_mat.getSolutionNames(),
+                   pqc_mat.withRedox(),
                    pqc_mat.getExchanger(cell_id),
                    pqc_mat.getKineticsNames(cell_id),
                    pqc_mat.getEquilibriumNames(cell_id),
@@ -133,8 +135,8 @@ void PhreeqcEngine::Impl::run(double time_step) {
 void PhreeqcEngine::Impl::init_wrappers(const InitCell &cell) {
 
   // Solutions
-  this->solutionWrapperPtr =
-      std::make_unique<SolutionWrapper>(this->Get_solution(1), cell.solutions);
+  this->solutionWrapperPtr = std::make_unique<SolutionWrapper>(
+      this->Get_solution(1), cell.solutions, cell.with_redox);
 
   if (this->Get_exchange(1) != nullptr) {
     this->exchangeWrapperPtr = std::make_unique<ExchangeWrapper>(

poet/src/PhreeqcMatrix/Ctor.cpp

@@ -9,8 +9,10 @@
 #include <string>
 
 PhreeqcMatrix::PhreeqcMatrix(const std::string &database,
-                             const std::string &input_script, bool with_h0_o0)
-    : _m_database(database), _m_with_h0_o0(with_h0_o0) {
+                             const std::string &input_script, bool with_h0_o0,
+                             bool with_redox)
+    : _m_database(database), _m_with_h0_o0(with_h0_o0),
+      _m_with_redox(with_redox) {
   this->_m_pqc = std::make_shared<IPhreeqc>();
 
   this->_m_pqc->LoadDatabaseString(database.c_str());

poet/src/PhreeqcMatrix/Init.cpp

@@ -19,13 +19,20 @@
 #include <vector>
 
 bool include_h0_o0 = false;
+bool with_redox = false;
 
 static std::vector<std::string> dump_solution_names(cxxSolution *solution,
                                                     Phreeqc *phreeqc) {
   std::vector<std::string> placeholder;
 
-  return phreeqc->find_all_valence_states(
-      SolutionWrapper::names(solution, include_h0_o0, placeholder));
+  std::vector<std::string> solnames =
+      SolutionWrapper::names(solution, include_h0_o0, placeholder, with_redox);
+
+  if (with_redox) {
+    solnames = phreeqc->find_all_valence_states(solnames);
+  }
+
+  return solnames;
 }
 
 template <enum PhreeqcMatrix::PhreeqcComponent comp, class T>
@@ -123,7 +130,8 @@ create_vector_from_phreeqc(Phreeqc *phreeqc, int id,
 
   // Solution
   SolutionWrapper sol_wrapper(
-      Utilities::Rxn_find(phreeqc->Get_Rxn_solution_map(), id), solution_names);
+      Utilities::Rxn_find(phreeqc->Get_Rxn_solution_map(), id), solution_names,
+      with_redox);
 
   base_add_to_element_vector<PhreeqcMatrix::PhreeqcComponent::SOLUTION>(
       sol_wrapper, solution_names, elements);
@@ -208,6 +216,7 @@ void PhreeqcMatrix::initialize() {
   }
 
   include_h0_o0 = this->_m_with_h0_o0;
+  with_redox = this->_m_with_redox;
 
   std::vector<std::string> solutions = find_all_solutions(phreeqc);
 

poet/src/Wrapper/SolutionWrapper.cpp

@@ -4,9 +4,11 @@
 #include <vector>
 
 SolutionWrapper::SolutionWrapper(
-    cxxSolution *soln, const std::vector<std::string> &_solution_order)
+    cxxSolution *soln, const std::vector<std::string> &_solution_order,
+    bool with_redox)
     : solution(soln), solution_order(_solution_order.begin() + NUM_ESSENTIALS,
-                                     _solution_order.end()) {
+                                     _solution_order.end()),
+      _with_redox(with_redox) {
   this->num_elements = _solution_order.size();
 
   auto &totals = solution->Get_totals();
@@ -17,10 +19,14 @@ void SolutionWrapper::get(std::span<LDBLE> &data) const {
   data[1] = solution->Get_total_o();
   data[2] = solution->Get_cb();
 
+  const cxxNameDouble &totals =
+      (_with_redox ? solution->Get_totals()
+                   : solution->Get_totals().Simplify_redox());
+
   std::size_t i = NUM_ESSENTIALS;
   for (const auto &tot_name : solution_order) {
-    auto it = solution->Get_totals().find(tot_name);
-    if (it == solution->Get_totals().end()) {
+    auto it = totals.find(tot_name);
+    if (it == totals.end()) {
       data[i++] = 0.0;
       continue;
     }
@@ -45,12 +51,15 @@ void SolutionWrapper::set(const std::span<LDBLE> &data) {
     new_totals[tot_name] = value;
   }
 
-  this->solution->Update(total_h, total_o, cb, new_totals);
+  this->solution->Update(total_h, total_o, cb,
+                         _with_redox ? new_totals
+                                     : new_totals.Simplify_redox());
 }
 
 std::vector<std::string>
 SolutionWrapper::names(cxxSolution *solution, bool include_h0_o0,
-                       std::vector<std::string> &solution_order) {
+                       std::vector<std::string> &solution_order,
+                       bool with_redox) {
   std::vector<std::string> names;
 
   names.insert(names.end(), ESSENTIALS.begin(), ESSENTIALS.end());
@@ -62,13 +71,16 @@ SolutionWrapper::names(cxxSolution *solution, bool include_h0_o0,
 
   std::set<std::string> names_set;
 
-  for (const auto &name : solution->Get_totals()) {
-    names_set.insert(name.first);
-  }
+  const cxxNameDouble &totals =
+      (with_redox ? solution->Get_totals()
+                  : solution->Get_totals().Simplify_redox());
 
-  // remove H(0) and O(0) from the set as they are already in the vector (if)
-  for (const auto &to_erase : {"H(0)", "O(0)"}) {
-    names_set.erase(to_erase);
+  for (const auto &[name, _] : totals) {
+    // Skip redox states of H and O
+    if (name == "H(0)" || name == "O(0)") {
+      continue;
+    }
+    names_set.insert(name);
   }
 
   names.insert(names.end(), names_set.begin(), names_set.end());

poet/src/Wrapper/SolutionWrapper.hpp

@@ -10,7 +10,8 @@
 class SolutionWrapper : public WrapperBase {
 public:
   SolutionWrapper(cxxSolution *soln,
-                  const std::vector<std::string> &solution_order);
+                  const std::vector<std::string> &solution_order,
+                  bool with_redox);
 
   void get(std::span<LDBLE> &data) const;
 
@@ -18,7 +19,7 @@ public:
 
   static std::vector<std::string>
   names(cxxSolution *solution, bool include_h0_o0,
-        std::vector<std::string> &solution_order);
+        std::vector<std::string> &solution_order, bool with_redox);
 
   std::vector<std::string> getEssentials() const;
 
@@ -30,4 +31,6 @@ private:
                                                              "Charge"};
 
   static constexpr std::size_t NUM_ESSENTIALS = ESSENTIALS.size();
+
+  const bool _with_redox;
 };

poet/test/testPhreeqcMatrix.cpp

@@ -1,5 +1,6 @@
 #include <cmath>
 #include <gtest/gtest.h>
+#include <string>
 #include <vector>
 
 #include <testInput.hpp>
@@ -208,3 +209,13 @@ POET_TEST(PhreeqcMatrixColumnMajorExport) {
   EXPECT_EQ(exported.values[0], 2);
   EXPECT_EQ(exported.values[1], 3);
 }
+
+POET_TEST(PhreeqcMatrixWithoutRedoxAndH0O0) {
+  PhreeqcMatrix pqc_mat(barite_db, barite_script, false, false);
+
+  const std::vector<std::string> expected_names_without_redox = {
+      "H", "O", "Charge", "Ba", "Cl", "S", "Sr",
+  };
+
+  EXPECT_EQ(expected_names_without_redox, pqc_mat.getSolutionNames());
+}

src/phreeqcpp/GFZ.cpp

@@ -4,7 +4,7 @@
 const std::set<std::string> to_ignore = {"H", "O", "Charge", "H(0)", "O(0)"};
 
 std::vector<std::string> Phreeqc::find_all_valence_states(
-    const std::vector<std::string> &&solution_names) {
+    const std::vector<std::string> &solution_names) {
   std::vector<std::string> solution_with_valences;
   solution_with_valences.reserve(solution_names.size());
 

src/phreeqcpp/Phreeqc.h

@@ -1913,7 +1913,7 @@ protected:
 
 public:
   std::vector<std::string>
-  find_all_valence_states(const std::vector<std::string> &&solution_names);
+  find_all_valence_states(const std::vector<std::string> &solution_names);
 
   static const class const_iso iso_defaults[];
   static const int count_iso_defaults;