Reuse CSR pattern in PhasorDynamics system Jacobian

CHANGELOG.md

@@ -44,7 +44,8 @@
 - Minor performance improvements to residual evaluation in PowerElectronics module.
 - Added full support for sparse Jacobians obtained with Enzyme in PhasorDynamics.
 - Added `Node` class to the PowerElectronics module to separate nodes from circuit components.
-- Refactor Jacobian assembly in `PowerElectronics` module to reuse the CSR pattern.
+- Refactored Jacobian assembly in `PowerElectronics` module to reuse the CSR pattern.
+- Refactored Jacobian assembly in `PhasorDyanmcics` module to reuse the CSR pattern.
 ## v0.1
 
 - Refactored code to support adding different model families.

GridKit/LinearAlgebra/SparseMatrix/CMakeLists.txt

@@ -1,4 +1,4 @@
-gridkit_add_library(SparseMatrix
+gridkit_add_library(sparse_matrix
   SOURCES
     CooMatrix.cpp
     CsrMatrix.cpp

GridKit/LinearAlgebra/SparseMatrix/CooMatrix.cpp

@@ -400,35 +400,42 @@ namespace GridKit
       map_to_dedup_      = new IdxT[static_cast<size_t>(nnz_)];
       IdxT* csr_row_data = new IdxT[static_cast<size_t>(n_ + 1)]();
 
-      map_to_dedup_[0] = 0;
-      csr_row_data[h_row_data_[0] + 1]++;
-
-      // Deduplicate sorted entries
-      IdxT w = 0;
-      for (IdxT i = 1; i < nnz_; i++)
+      if (nnz_ == 0)
+      {
+        csr_row_data[0] = 0;
+      }
+      else
       {
-        if (h_row_data_[i] == h_row_data_[w] && h_col_data_[i] == h_col_data_[w])
+        map_to_dedup_[0] = 0;
+        csr_row_data[h_row_data_[0] + 1]++;
+
+        // Deduplicate sorted entries
+        IdxT w = 0;
+        for (IdxT i = 1; i < nnz_; i++)
         {
-          h_val_data_[w]   += h_val_data_[i];
-          map_to_dedup_[i]  = w;
+          if (h_row_data_[i] == h_row_data_[w] && h_col_data_[i] == h_col_data_[w])
+          {
+            h_val_data_[w]   += h_val_data_[i];
+            map_to_dedup_[i]  = w;
+          }
+          else
+          {
+            ++w;
+            h_row_data_[w]   = h_row_data_[i];
+            h_col_data_[w]   = h_col_data_[i];
+            h_val_data_[w]   = h_val_data_[i];
+            map_to_dedup_[i] = w;
+            csr_row_data[h_row_data_[w] + 1]++;
+          }
         }
-        else
+        // nnz after deduplication
+        nnz_ = w + 1;
+
+        for (IdxT i = 0; i < n_; i++)
         {
-          ++w;
-          h_row_data_[w]   = h_row_data_[i];
-          h_col_data_[w]   = h_col_data_[i];
-          h_val_data_[w]   = h_val_data_[i];
-          map_to_dedup_[i] = w;
-          csr_row_data[h_row_data_[w] + 1]++;
+          csr_row_data[i + 1] += csr_row_data[i];
         }
       }
-      // nnz after deduplication
-      nnz_ = w + 1;
-
-      for (IdxT i = 0; i < n_; i++)
-      {
-        csr_row_data[i + 1] += csr_row_data[i];
-      }
 
       return csr_row_data;
     }

GridKit/Model/Evaluator.hpp

@@ -4,6 +4,7 @@
 
 #include <GridKit/Constants.hpp>
 #include <GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp>
+#include <GridKit/LinearAlgebra/SparseMatrix/CooMatrix.hpp>
 #include <GridKit/LinearAlgebra/SparseMatrix/CsrMatrix.hpp>
 #include <GridKit/Model/VariableMonitor.hpp>
 #include <GridKit/ScalarTraits.hpp>

GridKit/Model/PhasorDynamics/CMakeLists.txt

@@ -20,6 +20,7 @@ gridkit_add_library(phasor_dynamics_core
   LINK_LIBRARIES
     PUBLIC GridKit::definitions
     PUBLIC GridKit::utilities_logger
+    PUBLIC GridKit::sparse_matrix
   INCLUDE_DIRECTORIES
     PRIVATE ${GRIDKIT_THIRD_PARTY_DIR}/nlohmann-json/include
     PRIVATE ${GRIDKIT_THIRD_PARTY_DIR}/magic-enum/include)

GridKit/Model/PhasorDynamics/SystemModel.hpp

@@ -36,6 +36,7 @@ namespace GridKit
       using signal_type    = PhasorDynamics::SignalNode<ScalarT, IdxT>;
       using component_type = PhasorDynamics::Component<ScalarT, IdxT>;
       using RealT          = typename Model::Evaluator<ScalarT, IdxT>::RealT;
+      using CsrMatrix      = typename Model::Evaluator<ScalarT, IdxT>::CsrMatrix;
 
       using PhasorDynamics::Component<ScalarT, IdxT>::gridkit_component_id_;
       using PhasorDynamics::Component<ScalarT, IdxT>::size_;
@@ -270,6 +271,16 @@ namespace GridKit
             delete signal;
           }
         }
+        if (csr_jac_ != nullptr)
+        {
+          delete csr_jac_;
+          csr_jac_ = nullptr;
+        }
+        if (map_to_csr_ != nullptr)
+        {
+          delete[] map_to_csr_;
+          map_to_csr_ = nullptr;
+        }
       }
 
       /**
@@ -605,77 +616,157 @@ namespace GridKit
        */
       int evaluateJacobian() override
       {
-        bool sort = !jacobian_allocated_; // sort in axpy only if not allocated
-        if (!jacobian_allocated_)
-        {
-          J_.zeroMatrix();
-        }
-        else
-        {
-          J_.zeroValuedMatrix();
-        }
-
-        std::vector<IdxT>  ctemp{};
-        std::vector<IdxT>  rtemp{};
-        std::vector<RealT> valtemp{};
-
         // Initialize bus Jacobians
         for (const auto& bus : buses_)
         {
           bus->evaluateJacobian();
         }
 
-        // Jacobian blocks owed by components
-        // Also updates bus Jacobians
+        // Evaluate component Jacobians and update bus Jacobians
         for (const auto& component : components_)
         {
           component->evaluateJacobian();
-          auto component_jacobian = component->getJacobian();
+        }
 
-          std::tuple<std::vector<IdxT>&, std::vector<IdxT>&, std::vector<RealT>&> component_jacobian_entries = component_jacobian.getEntries();
-          const auto [rows, columns, values]                                                                 = component_jacobian_entries;
-          for (size_t i = 0; i < rows.size(); ++i)
+        // Build or update system CSR Jacobian
+        if (csr_jac_ == nullptr)
+        {
+          // Count the number of non-zeros
+          IdxT nnz_dup = 0;
+          for (const auto& component : components_)
           {
-            rtemp.push_back(rows[i]);
-            ctemp.push_back(columns[i]);
-            valtemp.push_back(values[i]);
+            auto component_jacobian = component->getJacobian();
+
+            std::tuple<std::vector<IdxT>&, std::vector<IdxT>&, std::vector<RealT>&> component_jacobian_entries = component_jacobian.getEntries();
+            const auto [rows, columns, values]                                                                 = component_jacobian_entries;
+            for (size_t i = 0; i < rows.size(); ++i)
+            {
+              ++nnz_dup;
+            }
           }
-          // Need axpy because some branches may connect to the same buses, and own the same
-          // off-diagonal locations. @todo implement a more efficient approach.
-          J_.axpy(1.0, rtemp, ctemp, valtemp, sort);
-          rtemp.clear();
-          ctemp.clear();
-          valtemp.clear();
-        }
+          for (const auto& bus : buses_)
+          {
+            auto bus_jacobian = bus->getJacobian();
 
-        // Bus Jacobians
-        for (const auto& bus : buses_)
-        {
-          auto bus_jacobian = bus->getJacobian();
+            std::tuple<std::vector<IdxT>&, std::vector<IdxT>&, std::vector<RealT>&> bus_jacobian_entries = bus_jacobian.getEntries();
+            const auto [rows, columns, values]                                                           = bus_jacobian_entries;
+            for (size_t i = 0; i < rows.size(); ++i)
+            {
+              ++nnz_dup;
+            }
+          }
+
+          // Allocate COO triplet arrays (we own these until we hand off to CsrMatrix)
+          IdxT*  rows_dup = new IdxT[nnz_dup];
+          IdxT*  cols_dup = new IdxT[nnz_dup];
+          RealT* vals_dup = new RealT[nnz_dup];
 
-          std::tuple<std::vector<IdxT>&, std::vector<IdxT>&, std::vector<RealT>&> bus_jacobian_entries = bus_jacobian.getEntries();
-          const auto [rows, columns, values]                                                           = bus_jacobian_entries;
-          for (size_t i = 0; i < rows.size(); ++i)
+          IdxT counter = 0;
+          for (const auto& component : components_)
           {
-            rtemp.push_back(rows[i]);
-            ctemp.push_back(columns[i]);
-            valtemp.push_back(values[i]);
+            auto component_jacobian = component->getJacobian();
+
+            std::tuple<std::vector<IdxT>&, std::vector<IdxT>&, std::vector<RealT>&> component_jacobian_entries = component_jacobian.getEntries();
+            const auto [rows, columns, values]                                                                 = component_jacobian_entries;
+            for (size_t i = 0; i < rows.size(); ++i)
+            {
+              rows_dup[counter] = rows[i];
+              cols_dup[counter] = columns[i];
+              vals_dup[counter] = values[i];
+              counter++;
+            }
           }
-        }
+          for (const auto& bus : buses_)
+          {
+            auto bus_jacobian = bus->getJacobian();
+
+            std::tuple<std::vector<IdxT>&, std::vector<IdxT>&, std::vector<RealT>&> bus_jacobian_entries = bus_jacobian.getEntries();
+            const auto [rows, columns, values]                                                           = bus_jacobian_entries;
+            for (size_t i = 0; i < rows.size(); ++i)
+            {
+              rows_dup[counter] = rows[i];
+              cols_dup[counter] = columns[i];
+              vals_dup[counter] = values[i];
+              counter++;
+            }
+          }
+
+          // Build the system COO Jacobian
+          LinearAlgebra::CooMatrix<RealT, IdxT> jac(size_, size_, nnz_dup, &rows_dup, &cols_dup, &vals_dup);
+
+          // Populate CSR data with sort and deduplicate
+          IdxT* row_ptrs = jac.getCsrRowData();
+
+          // Deduplicated nnz
+          nnz_ = jac.getNnz();
+
+          // Allocate cols/vals with deduplicated nnz
+          IdxT*  cols = new IdxT[nnz_];
+          RealT* vals = new RealT[nnz_];
+
+          std::copy(jac.getColData(), jac.getColData() + nnz_, cols);
+          std::copy(jac.getValues(), jac.getValues() + nnz_, vals);
 
-        J_.axpy(1.0, rtemp, ctemp, valtemp, sort);
+          // Create the CSR Jacobian
+          csr_jac_ = new CsrMatrix(size_, size_, nnz_, &row_ptrs, &cols, &vals);
 
-        // Flag Jacobian as allocated
-        if (!jacobian_allocated_)
+          const IdxT* map_to_sorted = jac.getMapToSorted();
+          const IdxT* map_to_dedup  = jac.getMapToDeduplicated();
+
+          // Build a mappping from original COO index to CSR index
+          map_to_csr_ = new IdxT[nnz_dup];
+          for (IdxT i = 0; i < nnz_dup; ++i)
+          {
+            map_to_csr_[map_to_sorted[i]] = map_to_dedup[i];
+          }
+        }
+        else
         {
-          jacobian_allocated_ = true;
+          // Zero out values
+          RealT* vals = csr_jac_->getValues();
+          for (IdxT i = 0; i < csr_jac_->getNnz(); ++i)
+          {
+            vals[i] = 0.0;
+          }
+
+          // Update CSR values from component and bus Jacobians
+          IdxT counter = 0;
+          for (const auto& component : components_)
+          {
+            auto component_jacobian = component->getJacobian();
+
+            std::tuple<std::vector<IdxT>&, std::vector<IdxT>&, std::vector<RealT>&> component_jacobian_entries = component_jacobian.getEntries();
+            const auto [rows, columns, values]                                                                 = component_jacobian_entries;
+            for (size_t i = 0; i < rows.size(); ++i)
+            {
+              vals[map_to_csr_[counter]] += values[i];
+              ++counter;
+            }
+          }
+          for (const auto& bus : buses_)
+          {
+            auto bus_jacobian = bus->getJacobian();
+
+            std::tuple<std::vector<IdxT>&, std::vector<IdxT>&, std::vector<RealT>&> bus_jacobian_entries = bus_jacobian.getEntries();
+            const auto [rows, columns, values]                                                           = bus_jacobian_entries;
+            for (size_t i = 0; i < rows.size(); ++i)
+            {
+              vals[map_to_csr_[counter]] += values[i];
+              ++counter;
+            }
+          }
         }
 
         // J_.printMatrix("System Jacobian");
 
         return 0;
       }
 
+      CsrMatrix* getCsrJacobian() const override
+      {
+        return csr_jac_;
+      }
+
       bool monitoring() const override
       {
         return !monitor_.empty();
@@ -839,7 +930,9 @@ namespace GridKit
       std::map<IdxT, IdxT> gridkit_fault_indices_;  ///< Map between fault_id and component_id
 
       bool owns_components_{false};
-      bool jacobian_allocated_{false};
+
+      IdxT*      map_to_csr_{nullptr};
+      CsrMatrix* csr_jac_{nullptr};
 
       /// Variable monitor
       Model::VariableMonitorController<ScalarT> monitor_;

GridKit/Model/PowerElectronics/SystemModelPowerElectronics.hpp

@@ -9,9 +9,6 @@
 #include <vector>
 
 #include <GridKit/Constants.hpp>
-#include <GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp>
-#include <GridKit/LinearAlgebra/SparseMatrix/CooMatrix.hpp>
-#include <GridKit/LinearAlgebra/SparseMatrix/CsrMatrix.hpp>
 #include <GridKit/Model/PowerElectronics/CircuitComponent.hpp>
 #include <GridKit/Model/PowerElectronics/CircuitNode.hpp>
 #include <GridKit/ScalarTraits.hpp>

GridKit/Solver/Dynamic/CMakeLists.txt

@@ -21,7 +21,7 @@ if (GRIDKIT_ENABLE_SUNDIALS_SPARSE)
       PUBLIC SUNDIALS::sunlinsolklu
       PUBLIC GridKit::definitions
       PUBLIC GridKit::utilities_logger
-      PUBLIC GridKit::SparseMatrix)
+      PUBLIC GridKit::sparse_matrix)
 else()
   gridkit_add_library(solvers_dyn
     SOURCES
@@ -33,5 +33,5 @@ else()
       PUBLIC SUNDIALS::idas
       PUBLIC GridKit::definitions
       PUBLIC GridKit::utilities_logger
-      PUBLIC GridKit::SparseMatrix)
+      PUBLIC GridKit::sparse_matrix)
 endif()