From 9bf29f64ce21a099c4e726fa6b14f227bcb26c2e Mon Sep 17 00:00:00 2001 From: Luc Grosheintz Date: Tue, 5 Sep 2023 17:34:49 +0200 Subject: [PATCH 1/2] Simplify flat index computation. --- nmodl/ode.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/nmodl/ode.py b/nmodl/ode.py index 283038f7ec..e25da9337b 100644 --- a/nmodl/ode.py +++ b/nmodl/ode.py @@ -7,6 +7,7 @@ from importlib import import_module +import itertools import sympy as sp import re @@ -310,13 +311,14 @@ def solve_non_lin_system(eq_strings, vars, constants, function_calls): vecFcode = [] for i, eq in enumerate(eqs): vecFcode.append(f"F[{i}] = {sp.ccode(eq.simplify().subs(X_vec_map).evalf(), user_functions=custom_fcts)}") + vecJcode = [] - for i, jac in enumerate(jacobian): - # todo: fix indexing to be ascending order - flat_index = jacobian.rows * (i % jacobian.rows) + (i // jacobian.rows) - vecJcode.append( - f"J[{flat_index}] = {sp.ccode(jac.simplify().subs(X_vec_map).evalf(), user_functions=custom_fcts)}" - ) + for i, j in itertools.product(range(jacobian.rows), range(jacobian.cols)): + flat_index = i + jacobian.rows * j + + rhs = sp.ccode(jacobian[i,j].simplify().subs(X_vec_map).evalf(), user_functions=custom_fcts) + vecJcode.append(f"J[{flat_index}] = {rhs}") + # interweave code = _interweave_eqs(vecFcode, vecJcode) From 2c29610d306d252d40caae23176ecde2cd4433f1 Mon Sep 17 00:00:00 2001 From: Luc Grosheintz Date: Wed, 6 Sep 2023 08:35:40 +0200 Subject: [PATCH 2/2] Avoid transposing `J`. --- nmodl/ode.py | 1 + src/codegen/codegen_cpp_visitor.cpp | 2 +- src/solver/newton/newton.hpp | 10 ++-------- test/unit/newton/newton.cpp | 28 ++++++++++++++-------------- 4 files changed, 18 insertions(+), 23 deletions(-) diff --git a/nmodl/ode.py b/nmodl/ode.py index e25da9337b..0e8ce31a08 100644 --- a/nmodl/ode.py +++ b/nmodl/ode.py @@ -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}") diff --git a/src/codegen/codegen_cpp_visitor.cpp b/src/codegen/codegen_cpp_visitor.cpp index 946fca68ff..993bc7897b 100644 --- a/src/codegen/codegen_cpp_visitor.cpp +++ b/src/codegen/codegen_cpp_visitor.cpp @@ -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 " : ""); diff --git a/src/solver/newton/newton.hpp b/src/solver/newton/newton.hpp index 973b4cd332..2a1b5ee1f3 100644 --- a/src/solver/newton/newton.hpp +++ b/src/solver/newton/newton.hpp @@ -60,7 +60,7 @@ EIGEN_DEVICE_FUNC int newton_solver(Eigen::Matrix& X, // Vector to store result of function F(X): Eigen::Matrix F; // Matrix to store jacobian of F(X): - Eigen::Matrix J; + Eigen::Matrix J; // Solver iteration count: int iter = -1; while (++iter < max_iter) { @@ -72,12 +72,6 @@ EIGEN_DEVICE_FUNC int newton_solver(Eigen::Matrix& 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 pivot; Eigen::Matrix rowmax; // Check if J is singular @@ -184,7 +178,7 @@ EIGEN_DEVICE_FUNC int newton_solver_small_N(Eigen::Matrix& X, int max_iter) { bool invertible; Eigen::Matrix F; - Eigen::Matrix J, J_inv; + Eigen::Matrix J, J_inv; int iter = -1; while (++iter < max_iter) { functor(X, F, J); diff --git a/test/unit/newton/newton.cpp b/test/unit/newton/newton.cpp index fec200beb3..b1bdaff774 100644 --- a/test/unit/newton/newton.cpp +++ b/test/unit/newton/newton.cpp @@ -206,7 +206,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") struct functor { void operator()(const Eigen::Matrix& X, Eigen::Matrix& F, - Eigen::Matrix& J) const { + Eigen::Matrix& 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 @@ -215,7 +215,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") }; Eigen::Matrix X{22.2536}; Eigen::Matrix F; - Eigen::Matrix J; + Eigen::Matrix J; functor fn; int iter_newton = newton::newton_solver(X, fn); fn(X, F, J); @@ -232,14 +232,14 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") struct functor { void operator()(const Eigen::Matrix& X, Eigen::Matrix& F, - Eigen::Matrix& J) const { + Eigen::Matrix& 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 X{-0.21421}; Eigen::Matrix F; - Eigen::Matrix J; + Eigen::Matrix J; functor fn; int iter_newton = newton::newton_solver(X, fn); fn(X, F, J); @@ -256,7 +256,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") struct functor { void operator()(const Eigen::Matrix& X, Eigen::Matrix& F, - Eigen::Matrix& J) const { + Eigen::Matrix& 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; @@ -267,7 +267,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") }; Eigen::Matrix X{0.2, 0.4}; Eigen::Matrix F; - Eigen::Matrix J; + Eigen::Matrix J; functor fn; int iter_newton = newton::newton_solver(X, fn); fn(X, F, J); @@ -292,7 +292,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") double z = 0.99; void operator()(const Eigen::Matrix& X, Eigen::Matrix& F, - Eigen::Matrix& J) const { + Eigen::Matrix& 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]); @@ -309,7 +309,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") }; Eigen::Matrix X{0.21231, 0.4435, -0.11537}; Eigen::Matrix F; - Eigen::Matrix J; + Eigen::Matrix J; functor fn; int iter_newton = newton::newton_solver(X, fn); fn(X, F, J); @@ -330,7 +330,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") double dt = 0.2; void operator()(const Eigen::Matrix& X, Eigen::Matrix& F, - Eigen::Matrix& J) const { + Eigen::Matrix& 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]); @@ -355,7 +355,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") }; Eigen::Matrix X{0.21231, 0.4435, -0.11537, -0.8124312}; Eigen::Matrix F; - Eigen::Matrix J; + Eigen::Matrix J; functor fn; int iter_newton = newton::newton_solver(X, fn); fn(X, F, J); @@ -371,7 +371,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") struct functor { void operator()(const Eigen::Matrix& X, Eigen::Matrix& F, - Eigen::Matrix& J) const { + Eigen::Matrix& 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]; @@ -409,7 +409,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") Eigen::Matrix X; X << 8.234, -245.46, 123.123, 0.8343, 5.555; Eigen::Matrix F; - Eigen::Matrix J; + Eigen::Matrix J; functor fn; int iter_newton = newton::newton_solver(X, fn); fn(X, F, J); @@ -425,7 +425,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") struct functor { void operator()(const Eigen::Matrix& X, Eigen::Matrix& F, - Eigen::Matrix& J) const { + Eigen::Matrix& 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]; @@ -545,7 +545,7 @@ SCENARIO("Non-linear system to solve with Newton Solver", "[analytic][solver]") Eigen::Matrix X; X << 8.234, -5.46, 1.123, 0.8343, 5.555, 18.234, -2.46, 0.123, 10.8343, -4.685; Eigen::Matrix F; - Eigen::Matrix J; + Eigen::Matrix J; functor fn; int iter_newton = newton::newton_solver(X, fn); fn(X, F, J);