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.
@@ -351,7 +360,11 @@ def include_flux(self, *args, **kwargs):
                     source = f
                     success = True
                 except:
-                    pass
+                    arith = args[0]
+                    if isinstance(arith, rxdmath._Arithmeticed):
+                        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,50 @@
+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 = 5
+    dend2.nseg = 5
+
+    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
+
+
+def test_invalid_include_flux(neuron_nosave_instance):
+    h, rxd, _ = neuron_nosave_instance
+    dend1 = h.Section("dend1")
+    cyt1 = rxd.Region(dend1.wholetree(), nrn_region="i")
+    ca1 = rxd.Species(cyt1, name="ca1", charge=2, initial=1e-12)
+    node1 = ca1.nodes(dend1(0.5))[0]
+    try:
+        node1.include_flux("wow")
+    except Exception as e:
+        assert isinstance(e, rxd.RxDException)
+    try:
+        node1.include_flux([])
+    except Exception as e:
+        assert isinstance(e, rxd.RxDException)
+    try:
+        node1.include_flux(None)
+    except Exception as e:
+        assert isinstance(e, rxd.RxDException)
+    try:
+        node1.include_flux({"a": 1})
+    except Exception as e:
+        assert isinstance(e, rxd.RxDException)

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)