Skip to content

Commit

Permalink
Avoid transposing J.
Browse files Browse the repository at this point in the history
  • Loading branch information
@1uc
1uc committed Sep 8, 2023
1 parent 9bf29f6 commit 2c29610
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 23 deletions.
1 change: 1 addition & 0 deletions nmodl/ode.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -315,6 +315,7 @@ def solve_non_lin_system(eq_strings, vars, constants, function_calls):
vecJcode = []
for i, j in itertools.product(range(jacobian.rows), range(jacobian.cols)):
flat_index = i + jacobian.rows * j
flat_index = i*jacobian.cols + j

rhs = sp.ccode(jacobian[i,j].simplify().subs(X_vec_map).evalf(), user_functions=custom_fcts)
vecJcode.append(f"J[{flat_index}] = {rhs}")
Expand Down
2 changes: 1 addition & 1 deletion src/codegen/codegen_cpp_visitor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1853,7 +1853,7 @@ void CodegenCppVisitor::print_functor_definition(const ast::EigenNewtonSolverBlo
printer->fmt_text(
"void operator()(const Eigen::Matrix<{0}, {1}, 1>& nmodl_eigen_xm, Eigen::Matrix<{0}, {1}, "
"1>& nmodl_eigen_fm, "
"Eigen::Matrix<{0}, {1}, {1}>& nmodl_eigen_jm) {2}",
"Eigen::Matrix<{0}, {1}, {1}, Eigen::RowMajor>& nmodl_eigen_jm) {2}",
float_type,
N,
is_functor_const(variable_block, functor_block) ? "const " : "");
Expand Down
10 changes: 2 additions & 8 deletions src/solver/newton/newton.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ EIGEN_DEVICE_FUNC int newton_solver(Eigen::Matrix<double, N, 1>& X,
// Vector to store result of function F(X):
Eigen::Matrix<double, N, 1> F;
// Matrix to store jacobian of F(X):
Eigen::Matrix<double, N, N> J;
Eigen::Matrix<double, N, N, Eigen::RowMajor> J;
// Solver iteration count:
int iter = -1;
while (++iter < max_iter) {
Expand All @@ -72,12 +72,6 @@ EIGEN_DEVICE_FUNC int newton_solver(Eigen::Matrix<double, N, 1>& X,
// we have converged: return iteration count
return iter;
}
// In Eigen the default storage order is ColMajor.
// Crout's implementation requires matrices stored in RowMajor order (C-style arrays).
// Therefore, the transposeInPlace is critical such that the data() method to give the rows
// instead of the columns.
if (!J.IsRowMajor)
J.transposeInPlace();
Eigen::Matrix<int, N, 1> pivot;
Eigen::Matrix<double, N, 1> rowmax;
// Check if J is singular
Expand Down Expand Up @@ -184,7 +178,7 @@ EIGEN_DEVICE_FUNC int newton_solver_small_N(Eigen::Matrix<double, N, 1>& X,
int max_iter) {
bool invertible;
Eigen::Matrix<double, N, 1> F;
Eigen::Matrix<double, N, N> J, J_inv;
Eigen::Matrix<double, N, N, Eigen::RowMajor> J, J_inv;
int iter = -1;
while (++iter < max_iter) {
functor(X, F, J);
Expand Down
28 changes: 14 additions & 14 deletions test/unit/newton/newton.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
struct functor {
void operator()(const Eigen::Matrix<double, 1, 1>& X,
Eigen::Matrix<double, 1, 1>& F,
Eigen::Matrix<double, 1, 1>& J) const {
Eigen::Matrix<double, 1, 1, Eigen::RowMajor>& J) const {
// Function F(X) to find F(X)=0 solution
F[0] = -3.0 * X[0] + 3.0;
// Jacobian of F(X), i.e. the matrix dF(X)_i/dX_j
Expand All @@ -215,7 +215,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
};
Eigen::Matrix<double, 1, 1> X{22.2536};
Eigen::Matrix<double, 1, 1> F;
Eigen::Matrix<double, 1, 1> J;
Eigen::Matrix<double, 1, 1, Eigen::RowMajor> J;
functor fn;
int iter_newton = newton::newton_solver(X, fn);
fn(X, F, J);
Expand All @@ -232,14 +232,14 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
struct functor {
void operator()(const Eigen::Matrix<double, 1, 1>& X,
Eigen::Matrix<double, 1, 1>& F,
Eigen::Matrix<double, 1, 1>& J) const {
Eigen::Matrix<double, 1, 1, Eigen::RowMajor>& J) const {
F[0] = -3.0 * X[0] + std::sin(X[0]) + std::log(X[0] * X[0] + 11.435243) + 3.0;
J(0, 0) = -3.0 + std::cos(X[0]) + 2.0 * X[0] / (X[0] * X[0] + 11.435243);
}
};
Eigen::Matrix<double, 1, 1> X{-0.21421};
Eigen::Matrix<double, 1, 1> F;
Eigen::Matrix<double, 1, 1> J;
Eigen::Matrix<double, 1, 1, Eigen::RowMajor> J;
functor fn;
int iter_newton = newton::newton_solver(X, fn);
fn(X, F, J);
Expand All @@ -256,7 +256,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
struct functor {
void operator()(const Eigen::Matrix<double, 2, 1>& X,
Eigen::Matrix<double, 2, 1>& F,
Eigen::Matrix<double, 2, 2>& J) const {
Eigen::Matrix<double, 2, 2, Eigen::RowMajor>& J) const {
F[0] = -3.0 * X[0] * X[1] + X[0] + 2.0 * X[1] - 1.0;
F[1] = 4.0 * X[0] - 0.29999999999999999 * std::pow(X[1], 2) + X[1] + 0.4;
J(0, 0) = -3.0 * X[1] + 1.0;
Expand All @@ -267,7 +267,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
};
Eigen::Matrix<double, 2, 1> X{0.2, 0.4};
Eigen::Matrix<double, 2, 1> F;
Eigen::Matrix<double, 2, 2> J;
Eigen::Matrix<double, 2, 2, Eigen::RowMajor> J;
functor fn;
int iter_newton = newton::newton_solver(X, fn);
fn(X, F, J);
Expand All @@ -292,7 +292,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
double z = 0.99;
void operator()(const Eigen::Matrix<double, 3, 1>& X,
Eigen::Matrix<double, 3, 1>& F,
Eigen::Matrix<double, 3, 3>& J) const {
Eigen::Matrix<double, 3, 3, Eigen::RowMajor>& J) const {
F(0) = -(-_x_old - dt * (a * std::pow(X[0], 2) + X[1]) + X[0]);
F(1) = -(_y_old - dt * (a * X[0] + b * X[1] + d) + X[1]);
F(2) = -(_z_old - dt * (e * z - 3.0 * X[0] + 2.0 * X[1]) + X[2]);
Expand All @@ -309,7 +309,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
};
Eigen::Matrix<double, 3, 1> X{0.21231, 0.4435, -0.11537};
Eigen::Matrix<double, 3, 1> F;
Eigen::Matrix<double, 3, 3> J;
Eigen::Matrix<double, 3, 3, Eigen::RowMajor> J;
functor fn;
int iter_newton = newton::newton_solver(X, fn);
fn(X, F, J);
Expand All @@ -330,7 +330,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
double dt = 0.2;
void operator()(const Eigen::Matrix<double, 4, 1>& X,
Eigen::Matrix<double, 4, 1>& F,
Eigen::Matrix<double, 4, 4>& J) const {
Eigen::Matrix<double, 4, 4, Eigen::RowMajor>& J) const {
F[0] = -(-3.0 * X[0] * X[2] * dt + X[0] - X0_old + 2.0 * dt / X[1]);
F[1] = -(X[1] - X1_old + dt * (4.0 * X[0] - 6.2 * X[1] + X[3]));
F[2] = -((X[2] * (X[2] - X2_old) - dt * (X[2] * (-1.2 * X[1] + 3.0) + 0.3)) / X[2]);
Expand All @@ -355,7 +355,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
};
Eigen::Matrix<double, 4, 1> X{0.21231, 0.4435, -0.11537, -0.8124312};
Eigen::Matrix<double, 4, 1> F;
Eigen::Matrix<double, 4, 4> J;
Eigen::Matrix<double, 4, 4, Eigen::RowMajor> J;
functor fn;
int iter_newton = newton::newton_solver(X, fn);
fn(X, F, J);
Expand All @@ -371,7 +371,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
struct functor {
void operator()(const Eigen::Matrix<double, 5, 1>& X,
Eigen::Matrix<double, 5, 1>& F,
Eigen::Matrix<double, 5, 5>& J) const {
Eigen::Matrix<double, 5, 5, Eigen::RowMajor>& J) const {
F[0] = -3.0 * X[0] * X[2] + X[0] + 2.0 / X[1];
F[1] = 4.0 * X[0] - 5.2 * X[1] + X[3];
F[2] = 1.2 * X[1] + X[2] - 3.0 - 0.3 / X[2];
Expand Down Expand Up @@ -409,7 +409,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
Eigen::Matrix<double, 5, 1> X;
X << 8.234, -245.46, 123.123, 0.8343, 5.555;
Eigen::Matrix<double, 5, 1> F;
Eigen::Matrix<double, 5, 5> J;
Eigen::Matrix<double, 5, 5, Eigen::RowMajor> J;
functor fn;
int iter_newton = newton::newton_solver(X, fn);
fn(X, F, J);
Expand All @@ -425,7 +425,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
struct functor {
void operator()(const Eigen::Matrix<double, 10, 1>& X,
Eigen::Matrix<double, 10, 1>& F,
Eigen::Matrix<double, 10, 10>& J) const {
Eigen::Matrix<double, 10, 10, Eigen::RowMajor>& J) const {
F[0] = -3.0 * X[0] * X[1] + X[0] + 2.0 * X[1];
F[1] = 4.0 * X[0] - 0.29999999999999999 * std::pow(X[1], 2) + X[1];
F[2] = 2.0 * X[1] + X[2] + 2.0 * X[3] * X[5] * X[7] - 3.0 * X[4] * X[8] - X[5];
Expand Down Expand Up @@ -545,7 +545,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]")
Eigen::Matrix<double, 10, 1> X;
X << 8.234, -5.46, 1.123, 0.8343, 5.555, 18.234, -2.46, 0.123, 10.8343, -4.685;
Eigen::Matrix<double, 10, 1> F;
Eigen::Matrix<double, 10, 10> J;
Eigen::Matrix<double, 10, 10, Eigen::RowMajor> J;
functor fn;
int iter_newton = newton::newton_solver(X, fn);
fn(X, F, J);
Expand Down

0 comments on commit 2c29610

Please sign in to comment.