Fix issues with wrong # of ODEs

BlueBrain · Oct 21, 2024 · 6bee9e0 · 6bee9e0
1 parent 052a8b7
commit 6bee9e0
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/src/codegen/codegen_helper_visitor.cpp b/src/codegen/codegen_helper_visitor.cpp
@@ -632,6 +632,7 @@ void CodegenHelperVisitor::visit_nrn_state_block(const ast::NrnStateBlock& node)
 
 void CodegenHelperVisitor::visit_cvode_block(const ast::CvodeBlock& node) {
     info.cvode_block = &node;
+    info.num_equations = node.get_n_odes()->get_value();
     node.visit_children(*this);
 }
 
@@ -745,7 +746,6 @@ void CodegenHelperVisitor::visit_statement_block(const ast::StatementBlock& node
                                          return sym->get_name() == symbol->get_name();
                                      }) == info.prime_variables_by_order.end()) {
                         info.prime_variables_by_order.push_back(symbol);
-                        info.num_equations++;
                     }
                 }
             }

diff --git a/src/language/codegen.yaml b/src/language/codegen.yaml
@@ -96,6 +96,11 @@
                                   type: Name
                                   node_name: true
                                   suffix: {value: " "}
+                              - n_odes:
+                                  brief: "number of ODEs to solve"
+                                  type: Integer
+                                  prefix: {value: "["}
+                                  suffix: {value: "]"}
                               - non_stiff_block:
                                   brief: "Block with statements of the form Dvar = f(var), used for updating non-stiff systems"
                                   type: StatementBlock

diff --git a/src/visitors/cvode_visitor.cpp b/src/visitors/cvode_visitor.cpp
@@ -187,10 +187,12 @@ void CvodeVisitor::visit_program(ast::Program& node) {
 
         NonStiffVisitor(node.get_symbol_table()).visit_statement_block(*non_stiff_block);
         StiffVisitor(node.get_symbol_table()).visit_statement_block(*stiff_block);
-        node.emplace_back_node(
-            new ast::CvodeBlock(der_block->get_name(),
-                                std::shared_ptr<ast::StatementBlock>(non_stiff_block),
-                                std::shared_ptr<ast::StatementBlock>(stiff_block)));
+        auto prime_vars = collect_nodes(*der_block, {ast::AstNodeType::PRIME_NAME});
+        node.emplace_back_node(new ast::CvodeBlock(
+            der_block->get_name(),
+            std::shared_ptr<ast::Integer>(new ast::Integer(prime_vars.size(), nullptr)),
+            std::shared_ptr<ast::StatementBlock>(non_stiff_block),
+            std::shared_ptr<ast::StatementBlock>(stiff_block)));
     }
 }