Skip to content

Commit 25f59d8

Browse files
nkoukpaizannkoukpaizan
authored
PhasorDynamics simulations with sparse Jacobians (#312)
* Functional PhasorDynamics simulations with sparse Jacobians. * Cleaner Enzyme sparse wrapper implementation. * Replace int-int maps with vector of integers for residual/variable indices. * Make use of INVALID_INDEX constant. * More intuitive PhasorDynamics example test names. * Add sstream include to VariableMonitor. * Missing LINK_LIBRARIES for Logger. * CMake guard for Enzyme -O0 -enzyme-auto-sparsity=1 bug. * Misc documentation of Enzyme utilities.
1 parent fed7fa3 commit 25f59d8

92 files changed

Lines changed: 2178 additions & 1286 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
- Added capability to print monitored variables in multiple formats, triggered from `Ida::runSimulation`.
4343
- Added IDA statistics object which can be accumulated over multiple simulations.
4444
- Minor performance improvements to residual evaluation in PowerElectronics module.
45+
- Added full support for sparse Jacobians obtained with Enzyme in PhasorDynamics.
4546

4647
## v0.1
4748

‎GridKit/AutomaticDifferentiation/DependencyTracking/Variable.hpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#include <string>
1515
#include <vector>
1616

17+
#include <GridKit/Constants.hpp>
1718
#include <GridKit/ScalarTraits.hpp>
1819

1920
namespace GridKit
@@ -266,7 +267,7 @@ namespace GridKit
266267
bool is_fixed_; ///< Constant parameter flag.
267268

268269
mutable DependencyMap* dependencies_;
269-
static const size_t INVALID_VAR_NUMBER = static_cast<size_t>(-1);
270+
static const size_t INVALID_VAR_NUMBER = INVALID_INDEX<size_t>;
270271
};
271272

272273
//------------------------------------

‎GridKit/AutomaticDifferentiation/Enzyme/DfDwb.hpp‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
/**
2+
* @file DfDwb.hpp
3+
* @author Nicholson Koukpaizan (koukpaizannk@ornl.gov)
4+
*
5+
*/
6+
7+
#pragma once
8+
9+
#include <GridKit/AutomaticDifferentiation/Enzyme/EnzymeDefinitions.hpp>
10+
#include <GridKit/AutomaticDifferentiation/Enzyme/LowerSparseStorage.hpp>
11+
#include <GridKit/AutomaticDifferentiation/Enzyme/ModelWrappers.hpp>
12+
#include <GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp>
13+
#include <GridKit/ScalarTraits.hpp>
14+
15+
namespace GridKit
16+
{
17+
namespace Enzyme
18+
{
19+
namespace Sparse
20+
{
21+
/**
22+
* @brief Enzyme automatic differentiation Jacobian evaluator : df/dwb
23+
*
24+
* @tparam ModelT - model type
25+
* @tparam MemberFunctions - member function parameter key
26+
* @tparam ScalarT - scalar data type
27+
* @tparam IdxT - matrix index data type
28+
*/
29+
template <typename ModelT, MemberFunctions function, class ScalarT, typename IdxT>
30+
struct DfDwb
31+
{
32+
using RealT = typename GridKit::ScalarTraits<ScalarT>::RealT;
33+
using MatrixT = GridKit::LinearAlgebra::COO_Matrix<RealT, IdxT>;
34+
35+
/**
36+
* @param[in] model - Pointer to the model to be differentiated
37+
* @param[in] n_res - Number of residual functions
38+
* @param[in] n_var - Number of independent variables
39+
* @param[in] res_indices - Global residual indices
40+
* @param[in] var_indices - Global variable indices
41+
* @param[in] y - Internal variables
42+
* @param[in] yp - Internal variable derivatives
43+
* @param[in] wb - Bus variables
44+
* @param[in,out] jac - Jacobian
45+
*/
46+
static void eval(ModelT* model,
47+
size_t n_res,
48+
size_t n_var,
49+
const std::vector<IdxT>& res_indices,
50+
const std::vector<IdxT>& var_indices,
51+
ScalarT* y,
52+
ScalarT* yp,
53+
ScalarT* wb,
54+
MatrixT& jac)
55+
{
56+
if (n_res > 0 && n_var > 0)
57+
{
58+
std::vector<Triple<ScalarT>> triplets; //< @todo: Update once sparse storage format changes
59+
std::vector<ScalarT> elementary_v(n_var);
60+
for (size_t var_i = 0; var_i < n_var; ++var_i)
61+
{
62+
// Sparse storage. @see LowerSparseStorage.hpp
63+
ScalarT* output = __enzyme_todense<ScalarT*>((void*) ident_load<ScalarT>, (void*) ident_store<ScalarT>, var_i);
64+
ScalarT* d_output = __enzyme_todense<ScalarT*>((void*) sparse_load<ScalarT>, (void*) sparse_store<ScalarT>, var_i, &triplets);
65+
66+
// Elementary vector for Jacobian-vector product
67+
std::ranges::fill(elementary_v, 0.0);
68+
elementary_v[var_i] = 1.0;
69+
70+
// Core automatic differentiaation intrinsic that will be replaced by a derivative
71+
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function, ScalarT>::eval,
72+
enzyme_const,
73+
model,
74+
enzyme_const,
75+
y,
76+
enzyme_const,
77+
yp,
78+
enzyme_dup,
79+
wb,
80+
output,
81+
enzyme_dupnoneed,
82+
elementary_v.data(),
83+
d_output);
84+
}
85+
86+
// Store result
87+
std::vector<IdxT> ctemp{};
88+
std::vector<IdxT> rtemp{};
89+
std::vector<ScalarT> valtemp{};
90+
for (auto& tup : triplets)
91+
{
92+
rtemp.push_back(res_indices[static_cast<size_t>(tup.row)]);
93+
ctemp.push_back(var_indices[static_cast<size_t>(tup.col)]);
94+
valtemp.push_back(tup.val);
95+
}
96+
jac.setValues(rtemp, ctemp, valtemp); //< @todo: Update once sparse storage format changes
97+
}
98+
}
99+
100+
/**
101+
* @param[in] model - Pointer to the model to be differentiated
102+
* @param[in] n_res - Number of residual functions
103+
* @param[in] n_var - Number of independent variables
104+
* @param[in] res_indices - Global residual indices
105+
* @param[in] var_indices - Global variable indices
106+
* @param[in] y - Internal variables
107+
* @param[in] yp - Internal variable derivatives
108+
* @param[in] wb - Bus variables
109+
* @param[in] ws - Signal variables
110+
* @param[in,out] jac - Jacobian
111+
*/
112+
static void eval(ModelT* model,
113+
size_t n_res,
114+
size_t n_var,
115+
const std::vector<IdxT>& res_indices,
116+
const std::vector<IdxT>& var_indices,
117+
ScalarT* y,
118+
ScalarT* yp,
119+
ScalarT* wb,
120+
ScalarT* ws,
121+
MatrixT& jac)
122+
{
123+
if (n_res > 0 && n_var > 0)
124+
{
125+
std::vector<Triple<ScalarT>> triplets; //< @todo: Update once sparse storage format changes
126+
std::vector<ScalarT> elementary_v(n_var);
127+
for (size_t var_i = 0; var_i < n_var; ++var_i)
128+
{
129+
// Sparse storage. @see LowerSparseStorage.hpp
130+
ScalarT* output = __enzyme_todense<ScalarT*>((void*) ident_load<ScalarT>, (void*) ident_store<ScalarT>, var_i);
131+
ScalarT* d_output = __enzyme_todense<ScalarT*>((void*) sparse_load<ScalarT>, (void*) sparse_store<ScalarT>, var_i, &triplets);
132+
133+
// Elementary vector for Jacobian-vector product
134+
std::ranges::fill(elementary_v, 0.0);
135+
elementary_v[var_i] = 1.0;
136+
137+
// Core automatic differentiaation intrinsic that will be replaced by a derivative
138+
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function, ScalarT>::eval,
139+
enzyme_const,
140+
model,
141+
enzyme_const,
142+
y,
143+
enzyme_const,
144+
yp,
145+
enzyme_dup,
146+
wb,
147+
output,
148+
enzyme_const,
149+
ws,
150+
enzyme_dupnoneed,
151+
elementary_v.data(),
152+
d_output);
153+
}
154+
155+
// Store result
156+
std::vector<IdxT> ctemp{};
157+
std::vector<IdxT> rtemp{};
158+
std::vector<ScalarT> valtemp{};
159+
for (auto& tup : triplets)
160+
{
161+
rtemp.push_back(res_indices[static_cast<size_t>(tup.row)]);
162+
ctemp.push_back(var_indices[static_cast<size_t>(tup.col)]);
163+
valtemp.push_back(tup.val);
164+
}
165+
jac.setValues(rtemp, ctemp, valtemp); //< @todo: Update once sparse storage format changes
166+
}
167+
}
168+
};
169+
} // namespace Sparse
170+
} // namespace Enzyme
171+
} // namespace GridKit

‎GridKit/AutomaticDifferentiation/Enzyme/DfDws.hpp‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/**
2+
* @file DfDws.hpp
3+
* @author Nicholson Koukpaizan (koukpaizannk@ornl.gov)
4+
*
5+
*/
6+
7+
#pragma once
8+
9+
#include <GridKit/AutomaticDifferentiation/Enzyme/EnzymeDefinitions.hpp>
10+
#include <GridKit/AutomaticDifferentiation/Enzyme/LowerSparseStorage.hpp>
11+
#include <GridKit/AutomaticDifferentiation/Enzyme/ModelWrappers.hpp>
12+
#include <GridKit/Constants.hpp>
13+
#include <GridKit/LinearAlgebra/SparseMatrix/COO_Matrix.hpp>
14+
#include <GridKit/ScalarTraits.hpp>
15+
16+
namespace GridKit
17+
{
18+
namespace Enzyme
19+
{
20+
namespace Sparse
21+
{
22+
/**
23+
* @brief Enzyme automatic differentiation Jacobian evaluator: df/dws
24+
*
25+
* @tparam ModelT - model type
26+
* @tparam MemberFunctions - member function parameter key
27+
* @tparam ScalarT - scalar data type
28+
* @tparam IdxT - matrix index data type
29+
*/
30+
template <typename ModelT, MemberFunctions function, class ScalarT, typename IdxT>
31+
struct DfDws
32+
{
33+
using RealT = typename GridKit::ScalarTraits<ScalarT>::RealT;
34+
using MatrixT = GridKit::LinearAlgebra::COO_Matrix<RealT, IdxT>;
35+
36+
/**
37+
* @param[in] model - Pointer to the model to be differentiated
38+
* @param[in] n_res - Number of residual functions
39+
* @param[in] n_var - Number of independent variables
40+
* @param[in] res_indices - Global residual indices
41+
* @param[in] var_indices - Global variable indices
42+
* @param[in] y - Internal variables
43+
* @param[in] yp - Internal variable derivatives
44+
* @param[in] wb - Bus variables
45+
* @param[in] ws - Signal variables
46+
* @param[in,out] jac - Jacobian
47+
*/
48+
static void eval(ModelT* model,
49+
size_t n_res,
50+
size_t n_var,
51+
const std::vector<IdxT>& res_indices,
52+
const std::vector<IdxT>& var_indices,
53+
ScalarT* y,
54+
ScalarT* yp,
55+
ScalarT* wb,
56+
ScalarT* ws,
57+
MatrixT& jac)
58+
{
59+
if (n_res > 0 && n_var > 0)
60+
{
61+
std::vector<Triple<ScalarT>> triplets; //< @todo: Update once sparse storage format changes
62+
std::vector<ScalarT> elementary_v(n_var);
63+
for (size_t var_i = 0; var_i < n_var; ++var_i)
64+
{
65+
// Sparse storage. @see LowerSparseStorage.hpp
66+
ScalarT* output = __enzyme_todense<ScalarT*>((void*) ident_load<ScalarT>, (void*) ident_store<ScalarT>, var_i);
67+
ScalarT* d_output = __enzyme_todense<ScalarT*>((void*) sparse_load<ScalarT>, (void*) sparse_store<ScalarT>, var_i, &triplets);
68+
69+
// Elementary vector for Jacobian-vector product
70+
std::ranges::fill(elementary_v, 0.0);
71+
elementary_v[var_i] = 1.0;
72+
73+
// Core automatic differentiaation intrinsic that will be replaced by a derivative
74+
__enzyme_fwddiff<void>((void*) ModelWrapper<ModelT, function, ScalarT>::eval,
75+
enzyme_const,
76+
model,
77+
enzyme_const,
78+
y,
79+
enzyme_const,
80+
yp,
81+
enzyme_const,
82+
wb,
83+
enzyme_dup,
84+
ws,
85+
output,
86+
enzyme_dupnoneed,
87+
elementary_v.data(),
88+
d_output);
89+
}
90+
91+
// Store result
92+
std::vector<IdxT> ctemp{};
93+
std::vector<IdxT> rtemp{};
94+
std::vector<ScalarT> valtemp{};
95+
for (auto& tup : triplets)
96+
{
97+
IdxT row = res_indices[static_cast<size_t>(tup.row)];
98+
IdxT col = var_indices[static_cast<size_t>(tup.col)];
99+
if (col != INVALID_INDEX<IdxT>)
100+
{
101+
rtemp.push_back(row);
102+
ctemp.push_back(col);
103+
valtemp.push_back(tup.val);
104+
}
105+
}
106+
jac.setValues(rtemp, ctemp, valtemp); //< @todo: Update once sparse storage format changes
107+
}
108+
}
109+
};
110+
} // namespace Sparse
111+
} // namespace Enzyme
112+
} // namespace GridKit

0 commit comments

Comments
 (0)