From 8a2cd24ea7b71a2846bbc4d79fe8f93273e1e2f3 Mon Sep 17 00:00:00 2001
From: Goran Jelic-Cizmek <goran.jelic-cizmek@epfl.ch>
Date: Mon, 2 Dec 2024 17:25:22 +0100
Subject: [PATCH 1/3] Use regex to remove units

---
 src/visitors/cvode_visitor.cpp | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/src/visitors/cvode_visitor.cpp b/src/visitors/cvode_visitor.cpp
index 617e1e6c6..3ac4be3a3 100644
--- a/src/visitors/cvode_visitor.cpp
+++ b/src/visitors/cvode_visitor.cpp
@@ -13,6 +13,7 @@
 #include "utils/logger.hpp"
 #include "visitors/visitor_utils.hpp"
 #include <optional>
+#include <regex>
 #include <utility>
 
 namespace pywrap = nmodl::pybind_wrappers;
@@ -35,6 +36,22 @@ static void remove_conserve_statements(ast::StatementBlock& node) {
     }
 }
 
+// remove units using state-of-the-art technology
+static void remove_units(ast::BinaryExpression& node) {
+    std::regex unit_pattern(R"((\d*\.?\d+)\([a-zA-Z]+\))");
+    auto rhs_string = to_nmodl(node.get_rhs());
+    auto rhs_string_no_units = fmt::format("{} = {}",
+                                           to_nmodl(node.get_lhs()),
+                                           std::regex_replace(rhs_string, unit_pattern, "$1"));
+    logger->debug("CvodeVisitor :: removing units from statement {}", to_nmodl(node));
+    logger->debug("CvodeVisitor :: result: {}", rhs_string_no_units);
+    auto expr_statement = std::dynamic_pointer_cast<ast::ExpressionStatement>(
+        create_statement(rhs_string_no_units));
+    const auto bin_expr = std::dynamic_pointer_cast<const ast::BinaryExpression>(
+        expr_statement->get_expression());
+    node.set_rhs(std::shared_ptr<ast::Expression>(bin_expr->get_rhs()->clone()));
+}
+
 static std::pair<std::string, std::optional<int>> parse_independent_var(
     std::shared_ptr<ast::Identifier> node) {
     auto variable = std::make_pair(node->get_node_name(), std::optional<int>());
@@ -152,7 +169,10 @@ class StiffVisitor: public CvodeHelperVisitor {
             program_symtab->insert(symbol);
         }
 
+        remove_units(node);
+
         auto rhs = node.get_rhs();
+
         // all indexed variables (need special treatment in SymPy)
         auto indexed_variables = get_indexed_variables(*rhs, name->get_node_name());
         auto diff2c = pywrap::EmbeddedPythonLoader::get_instance().api().diff2c;

From 34824f67b4df62160a64c32c09fa5609f9a25a5e Mon Sep 17 00:00:00 2001
From: Goran Jelic-Cizmek <goran.jelic-cizmek@epfl.ch>
Date: Tue, 3 Dec 2024 10:10:10 +0100
Subject: [PATCH 2/3] Generalize parsing and add tests

---
 src/visitors/cvode_visitor.cpp     | 7 +++++--
 test/usecases/cvode/derivative.mod | 8 +++++++-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/visitors/cvode_visitor.cpp b/src/visitors/cvode_visitor.cpp
index 3ac4be3a3..13ef1e819 100644
--- a/src/visitors/cvode_visitor.cpp
+++ b/src/visitors/cvode_visitor.cpp
@@ -36,9 +36,12 @@ static void remove_conserve_statements(ast::StatementBlock& node) {
     }
 }
 
-// remove units using state-of-the-art technology
+// remove units from CVODE block so sympy can parse it properly
 static void remove_units(ast::BinaryExpression& node) {
-    std::regex unit_pattern(R"((\d*\.?\d+)\([a-zA-Z]+\))");
+    // matches either an int or a float, followed by any (including zero)
+    // number of spaces, followed by an expression in parentheses, that only
+    // has letters of the alphabet
+    std::regex unit_pattern(R"((\d+\.?\d*|\.\d+)\s*\([a-zA-Z]+\))");
     auto rhs_string = to_nmodl(node.get_rhs());
     auto rhs_string_no_units = fmt::format("{} = {}",
                                            to_nmodl(node.get_lhs()),
diff --git a/test/usecases/cvode/derivative.mod b/test/usecases/cvode/derivative.mod
index d3715352f..2a8ba6ca6 100644
--- a/test/usecases/cvode/derivative.mod
+++ b/test/usecases/cvode/derivative.mod
@@ -2,6 +2,10 @@ NEURON {
     SUFFIX scalar
 }
 
+UNITS {
+    (um) = (micron)
+}
+
 PARAMETER {
     freq = 10
     a = 5
@@ -14,7 +18,7 @@ PARAMETER {
     k = 0.2
 }
 
-STATE {var1 var2 var3}
+STATE {var1 var2 var3 var4}
 
 INITIAL {
     var1 = v1
@@ -34,4 +38,6 @@ DERIVATIVE equation {
     var2' = -var2 * a
     : logistic ODE
     var3' = r * var3 * (1 - var3 / k)
+    : ODE with some units
+    var4' = 1(um) * var4 + a * .1(um) + r * 1.(um) + 1.0    (um)
 }

From 75994aa2c0fcd7f1f7633a3d9453fe77fba17171 Mon Sep 17 00:00:00 2001
From: Goran Jelic-Cizmek <goran.jelic-cizmek@epfl.ch>
Date: Tue, 3 Dec 2024 10:12:22 +0100
Subject: [PATCH 3/3] CI, would you mind running please?