include_flux for arithmeticed object

share/lib/python/neuron/rxd/node.py

@@ -1,7 +1,7 @@
 import neuron
 from neuron import h, nrn, hoc, nrn_dll_sym
-from . import region, constants
-from . import rxdsection
+from . import region, constants, species
+from . import rxdsection, rxdmath
 import numpy
 import weakref
 from .rxdException import RxDException
@@ -121,6 +121,15 @@ def _replace(old_offset, old_nseg, new_offset, new_nseg):
 _numpy_element_ref = neuron.numpy_element_ref
 
 
+def eval_arith_flux(arith, nregion, node):
+    func, _species = rxdmath._compile(arith, [nregion])
+    c = compile(list(func.values())[0][0], "f", "eval")
+    s = [[None] * region._region_count for _ in range(species._species_count)]
+    for specie in _species:
+        s[specie()._id][nregion._id] = float(specie().nodes(node.segment).value[0])
+    return eval(c, {"species": s})
+
+
 class Node(object):
     def satisfies(self, condition):
         """Tests if a Node satisfies a given condition.
@@ -350,8 +359,11 @@ def include_flux(self, *args, **kwargs):
                     f = float(args[0])
                     source = f
                     success = True
-                except:
-                    pass
+                except TypeError:
+                    arith = rxdmath._ensure_arithmeticed(args[0])
+                    source = lambda: eval_arith_flux(arith, self.region, self)
+                    scale = 1 / self.volume
+                    success = True
             if not success:
                 raise RxDException("unsupported flux form")
         _node_fluxes["index"].append(self._index)

test/rxd/test_arith_include_flux.py

@@ -0,0 +1,26 @@
+def test_arith_include_flux(neuron_nosave_instance):
+    diff = 1e-15
+    h, rxd, _ = neuron_nosave_instance
+    dend1 = h.Section("dend1")
+    dend2 = h.Section("dend2")
+    dend1.nseg = 4
+    dend2.nseg = 4
+
+    cyt1 = rxd.Region(dend1.wholetree(), nrn_region="i")
+    cyt2 = rxd.Region(dend2.wholetree(), nrn_region="i")
+    ca1 = rxd.Species(cyt1, name="ca1", charge=2, initial=1e-12)
+    ca2 = rxd.Species(cyt2, name="ca1", charge=2, initial=1e-12)
+
+    node1 = ca1.nodes(dend1(0.5))[0]
+    node2 = ca2.nodes(dend2(0.5))[0]
+    node1.include_flux(ca1 * (1 - ca1))
+    r = rxd.Rate(ca2, ca2 * (1 - ca2))
+
+    h.finitialize(-65)
+    h.dt /= 512
+    h.continuerun(0.025 / 4)
+    assert abs(node1.concentration - node2.concentration) < diff
+    h.continuerun(0.025 / 2)
+    assert abs(node1.concentration - node2.concentration) < diff
+    h.continuerun(0.025)
+    assert abs(node1.concentration - node2.concentration) < diff

test/rxd/test_nodelist_include_flux.py

@@ -1,7 +1,7 @@
 def test_nodelist_include_flux(neuron_nosave_instance):
+    diff = 1e-15
     h, rxd, _ = neuron_nosave_instance
     dend1 = h.Section("dend1")
-    diff = 1e-15
     cyt = rxd.Region(dend1.wholetree(), nrn_region="i")
     ca1 = rxd.Species(cyt, name="ca1", charge=2, initial=0)
     ca2 = rxd.Species(cyt, name="ca2", charge=2, initial=0)