Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid transposing J. #1070

Closed
wants to merge 2 commits into from
Closed

Avoid transposing J. #1070

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9bf29f6
Simplify flat index computation.
1uc Sep 5, 2023
2c29610
Avoid transposing `J`.
1uc Sep 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions nmodl/ode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from importlib import import_module

import itertools
import sympy as sp
import re

Expand Down Expand Up @@ -310,13 +311,15 @@ 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
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}")

# interweave
code = _interweave_eqs(vecFcode, vecJcode)

Expand Down
2 changes: 1 addition & 1 deletion src/codegen/codegen_cpp_visitor.cpp
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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