From 58bb7fd938aee2d0066eb3ba47159ad7fe9419fe Mon Sep 17 00:00:00 2001
From: Luc Grosheintz <luc.grosheintz@gmail.com>
Date: Mon, 8 Jan 2024 10:55:53 +0100
Subject: [PATCH] Array variables in DERIVATIVE block. (#1126)

* Array variables in DERIVATIVE block.

* Check RANGE array variables.
---
 src/codegen/codegen_neuron_cpp_visitor.cpp | 13 +++++++---
 test/usecases/CMakeLists.txt               |  2 +-
 test/usecases/cnexp_array/leonhard.mod     | 30 ++++++++++++++++++++++
 test/usecases/cnexp_array/simulate.py      | 27 +++++++++++++++++++
 4 files changed, 68 insertions(+), 4 deletions(-)
 create mode 100644 test/usecases/cnexp_array/leonhard.mod
 create mode 100644 test/usecases/cnexp_array/simulate.py

diff --git a/src/codegen/codegen_neuron_cpp_visitor.cpp b/src/codegen/codegen_neuron_cpp_visitor.cpp
index 1f3fb893f5..edd93d473c 100644
--- a/src/codegen/codegen_neuron_cpp_visitor.cpp
+++ b/src/codegen/codegen_neuron_cpp_visitor.cpp
@@ -468,9 +468,16 @@ void CodegenNeuronCppVisitor::print_make_instance() const {
     const auto codegen_float_variables_size = codegen_float_variables.size();
     for (int i = 0; i < codegen_float_variables_size; ++i) {
         const auto& float_var = codegen_float_variables[i];
-        printer->fmt_line("&_ml.template fpfield<{}>(0){}",
-                          i,
-                          i < codegen_float_variables_size - 1 ? "," : "");
+        if (float_var->is_array()) {
+            printer->fmt_line("_ml.template data_array<{}, {}>(0){}",
+                              i,
+                              float_var->get_length(),
+                              i < codegen_float_variables_size - 1 ? "," : "");
+        } else {
+            printer->fmt_line("&_ml.template fpfield<{}>(0){}",
+                              i,
+                              i < codegen_float_variables_size - 1 ? "," : "");
+        }
     }
     printer->pop_block(";");
     printer->pop_block();
diff --git a/test/usecases/CMakeLists.txt b/test/usecases/CMakeLists.txt
index 82e53102d7..244dd822cd 100644
--- a/test/usecases/CMakeLists.txt
+++ b/test/usecases/CMakeLists.txt
@@ -1,4 +1,4 @@
-set(NMODL_USECASE_DIRS cnexp_scalar)
+set(NMODL_USECASE_DIRS cnexp_scalar cnexp_array)
 
 foreach(usecase ${NMODL_USECASE_DIRS})
   add_test(NAME usecase_${usecase}
diff --git a/test/usecases/cnexp_array/leonhard.mod b/test/usecases/cnexp_array/leonhard.mod
new file mode 100644
index 0000000000..35d208d5e8
--- /dev/null
+++ b/test/usecases/cnexp_array/leonhard.mod
@@ -0,0 +1,30 @@
+NEURON {
+    SUFFIX leonhard
+    RANGE z
+}
+
+ASSIGNED {
+  z[3]
+}
+
+STATE {
+  x
+  y[2]
+}
+
+INITIAL {
+    x = 42.0
+    y[0] = 0.1
+    y[1] = -1.0
+    z[0] = 0.7
+    z[1] = 0.8
+    z[2] = 0.9
+}
+
+BREAKPOINT {
+    SOLVE dX METHOD cnexp
+}
+
+DERIVATIVE dX {
+    x' = (y[0] + y[1])*(z[0]*z[1]*z[2])*x
+}
diff --git a/test/usecases/cnexp_array/simulate.py b/test/usecases/cnexp_array/simulate.py
new file mode 100644
index 0000000000..fdcd1ea75f
--- /dev/null
+++ b/test/usecases/cnexp_array/simulate.py
@@ -0,0 +1,27 @@
+import numpy as np
+
+from neuron import h, gui
+from neuron.units import ms
+
+nseg = 1
+
+s = h.Section()
+s.insert("leonhard")
+s.nseg = nseg
+
+x_hoc = h.Vector().record(s(0.5)._ref_x_leonhard)
+t_hoc = h.Vector().record(h._ref_t)
+
+h.stdinit()
+h.tstop = 5.0 * ms
+h.run()
+
+x = np.array(x_hoc.as_numpy())
+t = np.array(t_hoc.as_numpy())
+
+rate = (0.1 - 1.0) * (0.7 * 0.8 * 0.9)
+x_exact = 42.0 * np.exp(rate*t)
+rel_err = np.abs(x - x_exact) / x_exact
+
+assert np.all(rel_err < 1e-12)
+print("leonhard: success")