diff --git a/src/nmodl/visitors/sympy_solver_visitor.cpp b/src/nmodl/visitors/sympy_solver_visitor.cpp index f57a70aaba..07c27ef9bd 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; } @@ -509,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 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; + } + if (eq_system_is_valid && !eq_system.empty()) { // solve system of ODEs in eq_system logger->debug("SympySolverVisitor :: Solving {} system of ODEs", solve_method); @@ -688,6 +724,22 @@ void SympySolverVisitor::visit_program(ast::Program& node) { } node.visit_children(*this); + + // 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(); + 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/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, diff --git a/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp b/test/nmodl/transpiler/unit/visitor/sympy_solver.cpp index ff6819e944..86a07ac70b 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,24 @@ 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"( + STATE { + a + } + 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")); + } + } +} 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..e39e219532 --- /dev/null +++ b/test/nmodl/transpiler/usecases/solve/test_cnexp_to_derivimplicit.py @@ -0,0 +1,50 @@ +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()) + + # 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", + # 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, + )