Use derivimplicit if cnexp solution contains an unimplemented function

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<ast::VarName>(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<ast::IndexedName>(lhs_name);
+        var_name += "[" +
+                    std::to_string(
+                        std::dynamic_pointer_cast<ast::Integer>(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<std::string> 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<ast::IndexedName>(lhs_name);
-            var_name +=
-                "[" +
-                std::to_string(
-                    std::dynamic_pointer_cast<ast::Integer>(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<ast::SolveBlock>(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<ast::Name>(
+                        std::make_shared<ast::String>(derivative_block_solve_method[block_name])));
+                }
+            }
+        }
+    }
 }
 
 }  // namespace visitor

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,

test/nmodl/transpiler/unit/visitor/sympy_solver.cpp

@@ -45,14 +45,10 @@ using nmodl::parser::NmodlDriver;
 // SympySolver visitor tests
 //=============================================================================
 
-std::vector<std::string> 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<std::string> 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<std::string> run_sympy_solver_visitor(
     // check that, after visitor rearrangement, parents are still up-to-date
     CheckParentVisitor().check_ast(*ast);
 
+    return ast;
+}
+
+std::vector<std::string> 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<std::string> 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"));
+        }
+    }
+}

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)
+}

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,
+    )