Refactor Jacobian assembly by KakeruUeda · Pull Request #339 · ORNL/GridKit

KakeruUeda · 2026-02-13T22:26:21Z

Description

Jacobian assembly currently uses a COO matrix, and we sort and rebuild CSR row pointers each step for IDA. The overhead from this process can become a bottleneck as nnz grows. This change keeps the system Jacobian in CSR and preserves a COO→CSR mapping for direct CSR updates.

Closes #301

cc @nkoukpaizan

Proposed changes

Store the system Jacobian in CSR format in PowerElectronicsModel.
Update allocate() to populate CSR entries and preserve the mapping.
Add CsrJac for the solver configuration (PhasorDynamics still uses the Jac function).

Currently PowerElectronics only.

Checklist

All tests pass.
Code compiles cleanly with flags -Wall -Wpedantic -Wconversion -Wextra.
The new code follows GridKit™ style guidelines.
[N/A] There are unit tests for the new code.
The new code is documented.
[N/A] The feature branch is rebased with respect to the target branch.
I have updated CHANGELOG.md to reflect the changes in this PR. If this is a minor PR that is part of a larger fix already included in the file, state so.

KakeruUeda · 2026-02-16T13:47:21Z

Timed runSimulation(t_final) in the microgrid app. CSR assembly improved runtime, while the error remained unchanged.

Before

After

shakedregev

Still need to test this. Looks good except for naming inconsistencies.

shakedregev · 2026-02-16T15:02:13Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

      IdxT               rows_size_;
      IdxT               columns_size_;
      bool               sorted_;
+      std::vector<IdxT>  map2csr_;


Change this to follow naming conventions globally.
Use _ and to instead of 2

shakedregev · 2026-02-16T15:03:46Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+      const IdxT nnz_dup = static_cast<IdxT>(row_indices_.size());
+
+      // Sort the entries while preserving the mapping
+      std::vector<IdxT> map2sorted(nnz_dup);


shakedregev · 2026-02-16T15:03:55Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+
+      // Deduplicate entries by summing values
+      std::vector<IdxT> row_ptrs(rows_size_ + 1, 0);
+      std::vector<IdxT> map2dedup(nnz_dup);


shakedregev · 2026-02-16T15:04:37Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+      for (size_t i = 0; i < ordervec.size(); i++)
+      {
+        // permutation swap
+        while (ordervec[i] != ordervec[ordervec[i]])


underscores

shakedregev · 2026-02-16T15:07:04Z

GridKit/Solver/Dynamic/Ida.cpp

+    int Ida<ScalarT, IdxT>::CsrJac(RealT t, RealT cj, N_Vector yy, N_Vector yp, N_Vector, SUNMatrix J, void* user_data, N_Vector, N_Vector, N_Vector)
+    {
+      GridKit::Model::Evaluator<ScalarT, IdxT>* model = static_cast<GridKit::Model::Evaluator<ScalarT, IdxT>*>(user_data);
+
+      model->updateTime(t, cj);
+      copyVec(yy, model->y());
+      copyVec(yp, model->yp());
+
+      model->evaluateJacobian();
+
+      using CsrMatrix      = GridKit::LinearAlgebra::CsrMatrix<RealT, IdxT>;
+      const CsrMatrix* Jac = model->getCsrJacobian();
+
+      SUNMatZero(J);
+
+      sunindextype* sun_rptr = SUNSparseMatrix_IndexPointers(J);
+      sunindextype* sun_cind = SUNSparseMatrix_IndexValues(J);
+      RealT*        sun_vals = SUNSparseMatrix_Data(J);
+
+      IdxT n   = Jac->getRows();
+      IdxT nnz = Jac->getNnz();
+
+      // Get reference to the jacobian entries
+      const IdxT*  rptr = Jac->getRowPtr();
+      const IdxT*  cind = Jac->getColInd();
+      const RealT* vals = Jac->getValues();
+
+      // Copy data from model jac to sundials
+      std::copy(rptr, rptr + n + 1, sun_rptr);
+      std::copy(cind, cind + nnz, sun_cind);
+      std::copy(vals, vals + nnz, sun_vals);


I assume these are copied over, but this file violates the naming conventions in almost every variable and function.

Some of these are SUNDIALS functions, so they don't follow our style guidelines.

shakedregev

Approved pending fixes to variable names

shakedregev · 2026-02-16T21:53:25Z

Timed runSimulation(t_final) in the microgrid app. CSR assembly improved runtime, while the error remained unchanged.

Before

After

shakedregev · 2026-02-16T21:54:34Z

It's impressive that there are such substantial speed gains on even a small system. This means that for real systems it's probably astronomical. I wonder if we can test that.

shakedregev · 2026-02-16T22:04:03Z

examples/LinearAlgebra/SparseTest/SparseTest.cpp

+  // Basic Verification test
+  std::vector<size_t> csr_rtest   = {0, 2, 4, 6};
+  std::vector<size_t> csr_ctest   = {0, 1, 1, 2, 0, 2};
+  std::vector<double> csr_valtest = {2.6, 1.1, 2.9, 1.3, 1.4, 3.3};
+  std::vector<size_t> maptest     = {2, 5, 0, 0, 4, 2, 1, 5, 3};
+
+  std::vector<size_t>& map_to_csr = C.getMapToCsr();
+
+  assert(csr_rtest.size() == csr_r.size());
+  assert(csr_ctest.size() == csr_c.size());
+  assert(csr_valtest.size() == csr_v.size());
+  assert(maptest.size() == map_to_csr.size());
+
+  for (size_t i = 0; i < csr_rtest.size(); i++)
+  {
+    if (csr_r[i] != csr_rtest[i])


_ before test

shakedregev · 2026-02-16T22:06:02Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+      std::vector<IdxT>&     getMapToCsr();

      void printMatrix(std::string name = "");

      void printMatrixMarket(const std::string& filename, const std::string& comment);

      static void sortSparseCOO(std::vector<IdxT>& rows, std::vector<IdxT>& columns, std::vector<RealT>& values);
+      static void sortSparseCOO(std::vector<IdxT>& rows, std::vector<IdxT>& columns, std::vector<RealT>& values, std::vector<IdxT>& map);


We should consistently capitalize entire acronyms or just the first letter. I tend to favor the Csr vs COO style better. We are essentially using the acronym as a word. In comments, we can have it all capitalized.

PS: Looks like most of the code actually capitalizes the entire acronym for COO, but not for CSR. @pelesh - thoughts?

I agree on this

Please follow style guidelines. Some of the non-compliant code was grandfathered in. That is why we still have some inconsistencies.

shakedregev · 2026-02-16T22:09:42Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+      std::vector<size_t> order_vec(rows.size());
+      std::size_t         n(0);
+      std::generate(std::begin(order_vec), std::end(order_vec), [&]


I missed this (perhaps multiple times), but why are we using std::vector here? Is this to keep consistency with the rest of the code?
Edit: is this just for comparison?

Yes, to keep consistency with the rest of the code for now.

Please remove std::vector here.

pelesh

CSR matrix implementation should remain in the *.cpp file.
I would avoid using STL containers here because we want to prototype parallel implementations.
Do not use existing COO matrix implementation in GridKit. If needed, port COO matrix from Re::Solve.

I made some other (mostly nitpicking) comments about the code.

pelesh · 2026-02-16T23:03:47Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+    template <typename RealT, typename IdxT>
+    std::tuple<std::vector<IdxT>, std::vector<IdxT>, std::vector<RealT>> COO_Matrix<RealT, IdxT>::getCsrData()


Do not use std::tuple.

pelesh · 2026-02-16T23:04:52Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+      std::vector<IdxT>&     getMapToCsr();

      void printMatrix(std::string name = "");

      void printMatrixMarket(const std::string& filename, const std::string& comment);

      static void sortSparseCOO(std::vector<IdxT>& rows, std::vector<IdxT>& columns, std::vector<RealT>& values);
+      static void sortSparseCOO(std::vector<IdxT>& rows, std::vector<IdxT>& columns, std::vector<RealT>& values, std::vector<IdxT>& map);


Please follow style guidelines. Some of the non-compliant code was grandfathered in. That is why we still have some inconsistencies.

pelesh · 2026-02-16T23:05:31Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+      sortSparseCOO(row_indices_, column_indices_, values_, map_to_sorted);
+


Per guidelines:

Suggested change

sortSparseCOO(row_indices_, column_indices_, values_, map_to_sorted);

sortSparseCoo(row_indices_, column_indices_, values_, map_to_sorted);

pelesh · 2026-02-16T23:06:06Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+     * @param[in] n number of columns
+     */
+    template <typename RealT, typename IdxT>
+    void COO_Matrix<RealT, IdxT>::resetEntries(std::vector<IdxT> r, std::vector<IdxT> c, std::vector<RealT> v, IdxT m, IdxT n)


I would avoid using COO_Matrix altogether. We are going to remove this class.

pelesh · 2026-02-16T23:09:13Z

GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp

+      std::vector<size_t> order_vec(rows.size());
+      std::size_t         n(0);
+      std::generate(std::begin(order_vec), std::end(order_vec), [&]


Please remove std::vector here.

KakeruUeda added 4 commits February 13, 2026 16:45

Refactor Jacobian assembly in PowerElectronics

19a5a84

Update CHANGELOG.md

e822f5b

Remove unnecessary include

cfd55e8

Fix install error

96c2abc

KakeruUeda self-assigned this Feb 14, 2026

KakeruUeda added 2 commits February 13, 2026 20:52

Small fix

c103d52

Fix comments

de0de19

KakeruUeda requested review from pelesh and shakedregev February 16, 2026 13:48

shakedregev reviewed Feb 16, 2026

View reviewed changes

shakedregev approved these changes Feb 16, 2026

View reviewed changes

shakedregev closed this Feb 16, 2026

shakedregev reopened this Feb 16, 2026

KakeruUeda and others added 2 commits February 16, 2026 16:56

Rename some variables

ba794ab

Apply pre-commmit fixes

0555728

shakedregev reviewed Feb 16, 2026

View reviewed changes

pelesh requested changes Feb 16, 2026

View reviewed changes

pelesh marked this pull request as draft February 16, 2026 23:13

		template <typename RealT, typename IdxT>
		std::tuple<std::vector<IdxT>, std::vector<IdxT>, std::vector<RealT>> COO_Matrix<RealT, IdxT>::getCsrData()

		sortSparseCOO(row_indices_, column_indices_, values_, map_to_sorted);

Conversation

KakeruUeda commented Feb 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Proposed changes

Checklist

Uh oh!

KakeruUeda commented Feb 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shakedregev left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shakedregev left a comment

Choose a reason for hiding this comment

Uh oh!

shakedregev commented Feb 16, 2026

Uh oh!

shakedregev commented Feb 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pelesh Feb 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shakedregev Feb 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pelesh left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pelesh Feb 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pelesh Feb 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

KakeruUeda commented Feb 13, 2026 •

edited

Loading

KakeruUeda commented Feb 16, 2026 •

edited

Loading

pelesh Feb 16, 2026 •

edited

Loading

shakedregev Feb 16, 2026 •

edited

Loading

pelesh left a comment •

edited

Loading

pelesh Feb 16, 2026 •

edited

Loading

pelesh Feb 16, 2026 •

edited

Loading