Add CvodeBlock and CvodeVisitor

src/language/code_generator.cmake

@@ -74,6 +74,7 @@ set(AST_GENERATED_SOURCES
     ${PROJECT_BINARY_DIR}/src/ast/constructor_block.hpp
     ${PROJECT_BINARY_DIR}/src/ast/define.hpp
     ${PROJECT_BINARY_DIR}/src/ast/derivative_block.hpp
+    ${PROJECT_BINARY_DIR}/src/ast/derivative_original_function_block.hpp
     ${PROJECT_BINARY_DIR}/src/ast/derivimplicit_callback.hpp
     ${PROJECT_BINARY_DIR}/src/ast/destructor_block.hpp
     ${PROJECT_BINARY_DIR}/src/ast/diff_eq_expression.hpp

src/language/codegen.yaml

@@ -87,7 +87,30 @@
                                   type: StatementBlock
                               - finalize_block:
                                   brief: "Statement block to be executed after calling linear solver"
-                                  type: StatementBlock                                  
+                                  type: StatementBlock
+                        - DerivativeOriginalFunctionBlock:
+                            nmodl: "DERIVATIVE_ORIGINAL_FUNCTION "
+                            members:
+                              - name:
+                                  brief: "Name of the derivative block"
+                                  type: Name
+                                  node_name: true
+                                  suffix: {value: " "}
+                              - statement_block:
+                                  brief: "Block with statements vector"
+                                  type: StatementBlock
+                                  getter: {override: true}
+                            brief: "Represents the original, unmodified `DERIVATIVE` block in the NMODL"
+                            description: |
+                                  The original `DERIVATIVE` block in NMODL is
+                                  replaced in-place if the system of ODEs is
+                                  solvable analytically. Therefore, this
+                                  block's sole purpose is to keep the
+                                  original, unsolved block in the AST. This is
+                                  primarily useful when we need to solve the
+                                  ODE system using implicit methods, for
+                                  instance, CVode.
+
                   - WrappedExpression:
                       brief: "Wrap any other expression type"
                       members:

src/main.cpp

@@ -25,6 +25,7 @@
 #include "visitors/after_cvode_to_cnexp_visitor.hpp"
 #include "visitors/ast_visitor.hpp"
 #include "visitors/constant_folder_visitor.hpp"
+#include "visitors/derivative_original_visitor.hpp"
 #include "visitors/function_callpath_visitor.hpp"
 #include "visitors/global_var_visitor.hpp"
 #include "visitors/implicit_argument_visitor.hpp"
@@ -497,6 +498,13 @@ int run_nmodl(int argc, const char* argv[]) {
         const bool sympy_linear = node_exists(*ast, ast::AstNodeType::LINEAR_BLOCK);
         const bool sympy_sparse = solver_exists(*ast, "sparse");
 
+        if (neuron_code) {
+            logger->info("Running derivative visitor");
+            DerivativeOriginalVisitor().visit_program(*ast);
+            SymtabVisitor(update_symtab).visit_program(*ast);
+            ast_to_nmodl(*ast, filepath("derivative_original"));
+        }
+
         if (sympy_conductance || sympy_analytic || sympy_sparse || sympy_derivimplicit ||
             sympy_linear) {
             nmodl::pybind_wrappers::EmbeddedPythonLoader::get_instance()

src/visitors/CMakeLists.txt

@@ -11,6 +11,7 @@ add_library(
   visitor STATIC
   after_cvode_to_cnexp_visitor.cpp
   constant_folder_visitor.cpp
+  derivative_original_visitor.cpp
   defuse_analyze_visitor.cpp
   function_callpath_visitor.cpp
   global_var_visitor.cpp

src/visitors/derivative_original_visitor.cpp

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2023 Blue Brain Project, EPFL.
+ * See the top-level LICENSE file for details.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "visitors/derivative_original_visitor.hpp"
+
+#include "ast/all.hpp"
+#include "lexer/token_mapping.hpp"
+#include "pybind/pyembed.hpp"
+#include "utils/logger.hpp"
+#include "visitors/visitor_utils.hpp"
+#include <optional>
+#include <utility>
+
+namespace pywrap = nmodl::pybind_wrappers;
+
+namespace nmodl {
+namespace visitor {
+
+
+void DerivativeOriginalVisitor::visit_derivative_block(ast::DerivativeBlock& node) {
+    node.visit_children(*this);
+    der_block_function = node.clone();
+}
+
+
+void DerivativeOriginalVisitor::visit_derivative_original_function_block(
+    ast::DerivativeOriginalFunctionBlock& node) {
+    derivative_block = true;
+    node_type = node.get_node_type();
+    node.visit_children(*this);
+    node_type = ast::AstNodeType::NODE;
+    derivative_block = false;
+}
+
+void DerivativeOriginalVisitor::visit_diff_eq_expression(ast::DiffEqExpression& node) {
+    differential_equation = true;
+    node.visit_children(*this);
+    differential_equation = false;
+}
+
+
+void DerivativeOriginalVisitor::visit_binary_expression(ast::BinaryExpression& node) {
+    const auto& lhs = node.get_lhs();
+
+    /// we have to only solve ODEs under original derivative block where lhs is variable
+    if (!derivative_block || !differential_equation || !lhs->is_var_name()) {
+        return;
+    }
+
+    auto name = std::dynamic_pointer_cast<ast::VarName>(lhs)->get_name();
+
+    if (name->is_prime_name() || name->is_indexed_name()) {
+        std::string varname;
+        if (name->is_prime_name()) {
+            varname = "D" + name->get_node_name();
+            logger->debug("DerivativeOriginalVisitor :: replacing {} with {} on LHS of {}",
+                          name->get_node_name(),
+                          varname,
+                          to_nmodl(node));
+            node.set_lhs(std::make_shared<ast::Name>(new ast::String(varname)));
+            if (program_symtab->lookup(varname) == nullptr) {
+                auto symbol = std::make_shared<symtab::Symbol>(varname, ModToken());
+                symbol->set_original_name(name->get_node_name());
+                program_symtab->insert(symbol);
+            }
+        } else {
+            varname = "D" + stringutils::remove_character(to_nmodl(node.get_lhs()), '\'');
+            // we discard the RHS here so it can be anything (as long as NMODL considers it valid)
+            auto statement = fmt::format("{} = {}", varname, varname);
+            logger->debug("DerivativeOriginalVisitor :: replacing {} with {} on LHS of {}",
+                          to_nmodl(node.get_lhs()),
+                          varname,
+                          to_nmodl(node));
+            auto expr_statement = std::dynamic_pointer_cast<ast::ExpressionStatement>(
+                create_statement(statement));
+            const auto bin_expr = std::dynamic_pointer_cast<const ast::BinaryExpression>(
+                expr_statement->get_expression());
+            node.set_lhs(std::shared_ptr<ast::Expression>(bin_expr->get_lhs()->clone()));
+            // TODO add symbol?
+        }
+    }
+}
+
+void DerivativeOriginalVisitor::visit_program(ast::Program& node) {
+    program_symtab = node.get_symbol_table();
+    node.visit_children(*this);
+    if (der_block_function) {
+        auto der_node =
+            new ast::DerivativeOriginalFunctionBlock(der_block_function->get_name(),
+                                                     der_block_function->get_statement_block());
+        node.emplace_back_node(der_node);
+    }
+
+    // re-visit the AST since we now inserted the DERIVATIVE_ORIGINAL block
+    node.visit_children(*this);
+}
+
+}  // namespace visitor
+}  // namespace nmodl

src/visitors/derivative_original_visitor.hpp

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2023 Blue Brain Project, EPFL.
+ * See the top-level LICENSE file for details.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+/**
+ * \file
+ * \brief \copybrief nmodl::visitor::DerivativeOriginalVisitor
+ */
+
+#include "symtab/decl.hpp"
+#include "visitors/ast_visitor.hpp"
+#include <string>
+
+namespace nmodl {
+namespace visitor {
+
+/**
+ * \addtogroup visitor_classes
+ * \{
+ */
+
+/**
+ * \class DerivativeOriginalVisitor
+ * \brief Make a copy of the `DERIVATIVE` block (if it exists), and insert back as
+ * `DERIVATIVE_ORIGINAL_FUNCTION` block.
+ *
+ * If \ref SympySolverVisitor runs successfully, it replaces the original
+ * solution. This block is inserted before that to prevent losing access to
+ * information about the block.
+ */
+class DerivativeOriginalVisitor: public AstVisitor {
+  private:
+    /// The copy of the derivative block we are solving
+    ast::DerivativeBlock* der_block_function = nullptr;
+
+    /// true while visiting differential equation
+    bool differential_equation = false;
+
+    /// global symbol table
+    symtab::SymbolTable* program_symtab = nullptr;
+
+    /// visiting derivative block
+    bool derivative_block = false;
+
+    ast::AstNodeType node_type = ast::AstNodeType::NODE;
+
+  public:
+    void visit_derivative_block(ast::DerivativeBlock& node) override;
+    void visit_program(ast::Program& node) override;
+    void visit_derivative_original_function_block(
+        ast::DerivativeOriginalFunctionBlock& node) override;
+    void visit_diff_eq_expression(ast::DiffEqExpression& node) override;
+    void visit_binary_expression(ast::BinaryExpression& node) override;
+};
+
+/** \} */  // end of visitor_classes
+
+}  // namespace visitor
+}  // namespace nmodl

src/visitors/sympy_solver_visitor.cpp

@@ -399,6 +399,10 @@ void SympySolverVisitor::visit_var_name(ast::VarName& node) {
     }
 }
 
+// Skip visiting DERIVATIVE_ORIGINAL block
+void SympySolverVisitor::visit_derivative_original_function_block(
+    ast::DerivativeOriginalFunctionBlock& node) {}
+
 void SympySolverVisitor::visit_diff_eq_expression(ast::DiffEqExpression& node) {
     const auto& lhs = node.get_expression()->get_lhs();
 

src/visitors/sympy_solver_visitor.hpp

@@ -185,6 +185,8 @@ class SympySolverVisitor: public AstVisitor {
     void visit_expression_statement(ast::ExpressionStatement& node) override;
     void visit_statement_block(ast::StatementBlock& node) override;
     void visit_program(ast::Program& node) override;
+    void visit_derivative_original_function_block(
+        ast::DerivativeOriginalFunctionBlock& node) override;
 };
 
 /** @} */  // end of visitor_classes