From 8fc978ee2627fd23998cc9c3e9e76eddc49e94c8 Mon Sep 17 00:00:00 2001 From: Goran Jelic-Cizmek Date: Tue, 2 Sep 2025 18:33:07 +0200 Subject: [PATCH 1/6] Use `derivimplicit` if `cnexp` provides an unimplemented function Fixes #3579 --- src/nmodl/visitors/sympy_solver_visitor.cpp | 63 +++++++++++++++------ src/nmodl/visitors/sympy_solver_visitor.hpp | 3 + 2 files changed, 50 insertions(+), 16 deletions(-) diff --git a/src/nmodl/visitors/sympy_solver_visitor.cpp b/src/nmodl/visitors/sympy_solver_visitor.cpp index f57a70aaba..f8b8a8bcf9 100644 --- a/src/nmodl/visitors/sympy_solver_visitor.cpp +++ b/src/nmodl/visitors/sympy_solver_visitor.cpp @@ -402,6 +402,42 @@ void SympySolverVisitor::visit_var_name(ast::VarName& node) { // Skip visiting CVODE block void SympySolverVisitor::visit_cvode_block(ast::CvodeBlock& node) {} + +void SympySolverVisitor::collect_odes(ast::DiffEqExpression& node) { + const auto& lhs = node.get_expression()->get_lhs(); + auto lhs_name = std::dynamic_pointer_cast(lhs)->get_name(); + auto eq_str = to_nmodl_for_sympy(node); + auto var_name = lhs_name->get_node_name(); + if (lhs_name->is_indexed_name()) { + auto index_name = std::dynamic_pointer_cast(lhs_name); + var_name += "[" + + std::to_string( + std::dynamic_pointer_cast(index_name->get_length())->eval()) + + "]"; + } + logger->debug("SympySolverVisitor :: adding ODE system: {}", eq_str); + eq_system.push_back(eq_str); + logger->debug("SympySolverVisitor :: adding state var: {}", var_name); + state_vars_in_block.insert(var_name); + expression_statements.insert(current_expression_statement); + last_expression_statement = current_expression_statement; +} + +/// Check that the sympy solution can be implemented in C++. +/// Usecase: sympy uses the Lambert W function sometimes, which is not +/// implemented in NEURON nor any C++ standard. NOCMODL handles this by falling +/// back to the `derivimplicit` method, so we mimic this usage here +static bool check_solution_is_implementable(const std::string& solution) { + const std::unordered_set unsolvable_in_cpp = {"LambertW"}; + for (const auto& item: unsolvable_in_cpp) { + if (solution.find(item) != std::string::npos) { + return false; + } + } + return true; +} + + void SympySolverVisitor::visit_diff_eq_expression(ast::DiffEqExpression& node) { const auto& lhs = node.get_expression()->get_lhs(); @@ -435,24 +471,19 @@ void SympySolverVisitor::visit_diff_eq_expression(ast::DiffEqExpression& node) { // with analytic solution for x(t+dt) in terms of x(t) // x = ... logger->debug("SympySolverVisitor :: CNEXP - solving: {}", node_as_nmodl); + if (!check_solution_is_implementable(solution)) { + logger->warn( + fmt::format("Could not solve ODE {} using {} method; using {} method instead", + node_as_nmodl, + solve_method, + codegen::naming::DERIVIMPLICIT_METHOD)); + solve_method = codegen::naming::DERIVIMPLICIT_METHOD; + collect_odes(node); + return; + } } else { // for other solver methods: just collect the ODEs & return - std::string eq_str = to_nmodl_for_sympy(node); - std::string var_name = lhs_name->get_node_name(); - if (lhs_name->is_indexed_name()) { - auto index_name = std::dynamic_pointer_cast(lhs_name); - var_name += - "[" + - std::to_string( - std::dynamic_pointer_cast(index_name->get_length())->eval()) + - "]"; - } - logger->debug("SympySolverVisitor :: adding ODE system: {}", eq_str); - eq_system.push_back(eq_str); - logger->debug("SympySolverVisitor :: adding state var: {}", var_name); - state_vars_in_block.insert(var_name); - expression_statements.insert(current_expression_statement); - last_expression_statement = current_expression_statement; + collect_odes(node); return; } diff --git a/src/nmodl/visitors/sympy_solver_visitor.hpp b/src/nmodl/visitors/sympy_solver_visitor.hpp index a77373e2d4..ee3c90f2a6 100644 --- a/src/nmodl/visitors/sympy_solver_visitor.hpp +++ b/src/nmodl/visitors/sympy_solver_visitor.hpp @@ -164,6 +164,9 @@ class SympySolverVisitor: public AstVisitor { /// max number of state vars allowed for small system linear solver int SMALL_LINEAR_SYSTEM_MAX_STATES; + /// collect ODEs for numerical solving + void collect_odes(ast::DiffEqExpression& node); + public: explicit SympySolverVisitor(bool use_pade_approx = false, bool elimination = true, From 0be6d6ce6b3c53936e4f3cf75f8af1223cddce37 Mon Sep 17 00:00:00 2001 From: Goran Jelic-Cizmek Date: Tue, 2 Sep 2025 19:12:28 +0200 Subject: [PATCH 2/6] Add test --- src/nmodl/visitors/sympy_solver_visitor.cpp | 21 +++++++++ .../transpiler/unit/visitor/sympy_solver.cpp | 43 +++++++++++++++---- 2 files changed, 56 insertions(+), 8 deletions(-) diff --git a/src/nmodl/visitors/sympy_solver_visitor.cpp b/src/nmodl/visitors/sympy_solver_visitor.cpp index f8b8a8bcf9..e899e9a384 100644 --- a/src/nmodl/visitors/sympy_solver_visitor.cpp +++ b/src/nmodl/visitors/sympy_solver_visitor.cpp @@ -540,6 +540,11 @@ void SympySolverVisitor::visit_derivative_block(ast::DerivativeBlock& node) { // - otherwise, each equation is added to eq_system node.visit_children(*this); + // if there are changes after visiting the children, register them here to be consistent + if (solve_method != derivative_block_solve_method[node.get_node_name()]) { + derivative_block_solve_method[node.get_node_name()] = solve_method; + } + if (eq_system_is_valid && !eq_system.empty()) { // solve system of ODEs in eq_system logger->debug("SympySolverVisitor :: Solving {} system of ODEs", solve_method); @@ -719,6 +724,22 @@ void SympySolverVisitor::visit_program(ast::Program& node) { } node.visit_children(*this); + + // register any changes here for consistency + for (const auto& block: solve_block_nodes) { + if (auto block_ptr = std::dynamic_pointer_cast(block)) { + const auto& block_name = block_ptr->get_block_name()->get_value()->eval(); + if (block_ptr->get_method()) { + // Note: solve method name is an optional parameter + // LINEAR and NONLINEAR blocks do not have solve method specified + const auto& solve_method = block_ptr->get_method()->get_value()->eval(); + if (solve_method != derivative_block_solve_method[block_name]) { + block_ptr->set_method(std::make_shared( + std::make_shared(derivative_block_solve_method[block_name]))); + } + } + } + } } } // namespace visitor diff --git a/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp b/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp index ff6819e944..3f400be0a6 100644 --- a/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp +++ b/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp @@ -45,14 +45,10 @@ using nmodl::parser::NmodlDriver; // SympySolver visitor tests //============================================================================= -std::vector run_sympy_solver_visitor( - const std::string& text, - bool pade = false, - bool cse = false, - AstNodeType ret_nodetype = AstNodeType::DIFF_EQ_EXPRESSION, - bool kinetic = false) { - std::vector results; - +auto run_sympy_solver_visitor_ast(const std::string& text, + bool pade = false, + bool cse = false, + bool kinetic = false) { // construct AST from text NmodlDriver driver; const auto& ast = driver.parse_string(text); @@ -76,6 +72,19 @@ std::vector run_sympy_solver_visitor( // check that, after visitor rearrangement, parents are still up-to-date CheckParentVisitor().check_ast(*ast); + return ast; +} + +std::vector run_sympy_solver_visitor( + const std::string& text, + bool pade = false, + bool cse = false, + AstNodeType ret_nodetype = AstNodeType::DIFF_EQ_EXPRESSION, + bool kinetic = false) { + std::vector results; + + const auto& ast = run_sympy_solver_visitor_ast(text, pade, cse, kinetic); + // run lookup visitor to extract results from AST for (const auto& eq: collect_nodes(*ast, {ret_nodetype})) { results.push_back(to_nmodl(eq)); @@ -1935,3 +1944,21 @@ SCENARIO("Solve KINETIC block using SympySolver Visitor", "[visitor][solver][sym } } } + +SCENARIO("Replace unimplementable cnexp solution with derivimplicit solution", + "[visitor][sympy][cnexp][derivimplicit]") { + GIVEN("Derivative block that has a LambertW analytic solution") { + std::string nmodl_text = R"( + BREAKPOINT { + SOLVE states METHOD cnexp + } + DERIVATIVE states { + a' = -a/(1 + a) + } + )"; + THEN("The method has been replaced with derivimplicit") { + const auto& result = run_sympy_solver_visitor_ast(nmodl_text); + REQUIRE_THAT(to_nmodl(result), Catch::Matchers::ContainsSubstring("derivimplicit")); + } + } +} From 956ace095c01d169a58e20201fe5a1e66c96f299 Mon Sep 17 00:00:00 2001 From: Goran Jelic-Cizmek Date: Wed, 3 Sep 2025 09:09:57 +0200 Subject: [PATCH 3/6] Add STATE so sympy doesn't raise an error --- test/nmodl/transpiler/unit/visitor/sympy_solver.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp b/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp index 3f400be0a6..86a07ac70b 100644 --- a/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp +++ b/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp @@ -1949,6 +1949,9 @@ SCENARIO("Replace unimplementable cnexp solution with derivimplicit solution", "[visitor][sympy][cnexp][derivimplicit]") { GIVEN("Derivative block that has a LambertW analytic solution") { std::string nmodl_text = R"( + STATE { + a + } BREAKPOINT { SOLVE states METHOD cnexp } From b18656ac2ea0fe13a9e34ad67df5c3aeee623f29 Mon Sep 17 00:00:00 2001 From: Goran Jelic-Cizmek Date: Wed, 3 Sep 2025 17:10:54 +0200 Subject: [PATCH 4/6] Add usecase test --- .../usecases/solve/cnexp_to_derivimplicit.mod | 19 +++++++++ .../solve/test_cnexp_to_derivimplicit.py | 41 +++++++++++++++++++ 2 files changed, 60 insertions(+) create mode 100644 test/nmodl/transpiler/usecases/solve/cnexp_to_derivimplicit.mod create mode 100644 test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py diff --git a/test/nmodl/transpiler/usecases/solve/cnexp_to_derivimplicit.mod b/test/nmodl/transpiler/usecases/solve/cnexp_to_derivimplicit.mod new file mode 100644 index 0000000000..a033bc062b --- /dev/null +++ b/test/nmodl/transpiler/usecases/solve/cnexp_to_derivimplicit.mod @@ -0,0 +1,19 @@ +NEURON { + SUFFIX cnexp_to_derivimplicit +} + +STATE { + x +} + +INITIAL { + x = 42 +} + +BREAKPOINT { + SOLVE dX METHOD cnexp +} + +DERIVATIVE dX { + x' = -x/(1 + x) +} diff --git a/test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py b/test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py new file mode 100644 index 0000000000..0b32160497 --- /dev/null +++ b/test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py @@ -0,0 +1,41 @@ +from typing import Optional + +from scipy.special import lambertw +import numpy as np +from neuron import h, gui +from neuron.units import ms + + +def test_cnexp_to_derivimplicit( + mech: str, + rtol: float, + dt: Optional[float] = None, +): + """ + Test that NMODL changes the solver from cnexp to derivimplicit if it + detects the Lambert W function in the solution + """ + nseg = 1 + + s = h.Section() + s.insert(mech) + s.nseg = nseg + + x_hoc = h.Vector().record(getattr(s(0.5), f"_ref_x_{mech}")) + t_hoc = h.Vector().record(h._ref_t) + + h.stdinit() + if dt is not None: + h.dt = dt * ms + h.tstop = 5.0 * ms + h.run() + + x = np.array(x_hoc.as_numpy()) + t = np.array(t_hoc.as_numpy()) + + x_exact = lambertw(42 * np.exp(-t) * np.exp(42)) + np.testing.assert_allclose(x, x_exact, rtol=rtol) + + +if __name__ == "__main__": + test_cnexp_to_derivimplicit("cnexp_to_derivimplicit", rtol=1.1e-6) From c54c010718cf26d94e0a58050fd8f16e8a6ea3aa Mon Sep 17 00:00:00 2001 From: Goran Jelic-Cizmek Date: Wed, 3 Sep 2025 17:40:22 +0200 Subject: [PATCH 5/6] Add explanations for possible ambiguities --- .../usecases/solve/test_cnexp_to_derivimplicit.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py b/test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py index 0b32160497..e39e219532 100644 --- a/test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py +++ b/test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py @@ -33,9 +33,18 @@ def test_cnexp_to_derivimplicit( x = np.array(x_hoc.as_numpy()) t = np.array(t_hoc.as_numpy()) + # solution to: + # x'(t) = - x / (x + 1) + # with x(t=0) = C1 is: + # x(t) = lambertw(C1 * exp(-t) * exp(C1)) x_exact = lambertw(42 * np.exp(-t) * np.exp(42)) np.testing.assert_allclose(x, x_exact, rtol=rtol) if __name__ == "__main__": - test_cnexp_to_derivimplicit("cnexp_to_derivimplicit", rtol=1.1e-6) + test_cnexp_to_derivimplicit( + "cnexp_to_derivimplicit", + # by trial and error, the derivimplicit solver seems to be accurate + # down to almost 1e-6, but not quite, hence the seemingly magic number + rtol=1.1e-6, + ) From 4ff6f369b84c538a2ce59bc5452c5372f9796395 Mon Sep 17 00:00:00 2001 From: Goran Jelic-Cizmek Date: Wed, 3 Sep 2025 17:45:09 +0200 Subject: [PATCH 6/6] Better comments --- src/nmodl/visitors/sympy_solver_visitor.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/nmodl/visitors/sympy_solver_visitor.cpp b/src/nmodl/visitors/sympy_solver_visitor.cpp index e899e9a384..07c27ef9bd 100644 --- a/src/nmodl/visitors/sympy_solver_visitor.cpp +++ b/src/nmodl/visitors/sympy_solver_visitor.cpp @@ -540,7 +540,7 @@ void SympySolverVisitor::visit_derivative_block(ast::DerivativeBlock& node) { // - otherwise, each equation is added to eq_system node.visit_children(*this); - // if there are changes after visiting the children, register them here to be consistent + // if there are changes to the solver method, register them to be consistent if (solve_method != derivative_block_solve_method[node.get_node_name()]) { derivative_block_solve_method[node.get_node_name()] = solve_method; } @@ -725,7 +725,7 @@ void SympySolverVisitor::visit_program(ast::Program& node) { node.visit_children(*this); - // register any changes here for consistency + // if the solver method changed, also change the AST to be consistent for (const auto& block: solve_block_nodes) { if (auto block_ptr = std::dynamic_pointer_cast(block)) { const auto& block_name = block_ptr->get_block_name()->get_value()->eval();