Fix issue with _mult ids when Parameters are used

Support Kineitic Block AST compilation.

Author: adam

share/lib/python/neuron/nmodl/ode.py

@@ -507,12 +507,18 @@ def optimize_odes(
         - new_local_vars: list of new temporary variable name strings from CSE
     """
     from sympy.codegen.rewriting import create_expand_pow_optimization
+    from sympy import Abs
+
     custom_fcts = _get_custom_functions(function_calls)
+    custom_fcts["Abs"] = "fabs"
 
     # Build sympy symbols the variables
     sympy_vars = {}
     for var in set(var_names + constants + list(local_consts)):
-        sympy_vars[var] = sp.Symbol(var, real=True)
+        if "_absmult_" in var:
+            sympy_vars[var] = Abs(sp.Symbol(f"_mult_{int(var.split('_')[-1])}"))
+        else:
+            sympy_vars[var] = sp.Symbol(var, real=True)
     local_consts = {sympy_vars[var]: value for var, value in local_consts.items()}
     # Parse each RHS expression
     expand_pow = create_expand_pow_optimization(10)
@@ -522,7 +528,6 @@ def optimize_odes(
 
     code = []
     local_vars = []
-
     if do_cse:
         my_symbols = sp.utilities.iterables.numbered_symbols(prefix="tmp_")
         sub_exprs, reduced = sp.cse(
@@ -539,7 +544,9 @@ def optimize_odes(
         rhs_exprs = reduced
 
     for i, expr in enumerate(rhs_exprs):
-        rhs = sp.ccode(expand_pow(expr).subs(local_consts).evalf(), user_functions=custom_fcts)
+        rhs = sp.ccode(
+            expand_pow(expr).subs(local_consts).evalf(), user_functions=custom_fcts
+        )
         if rhs_ids:
             code.append(f"{rhs_ids[i]} = {rhs}")
         else:

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

@@ -472,9 +472,7 @@ def ast(self, regions=None):
             ):
                 species.append(sp().ast().get_node_name())
 
-        if kinetic_block == "off" or (
-            kinetic_block == "mass_action" and self._custom_dynamics
-        ):
+        if kinetic_block == "off" or self._custom_dynamics:
             # represent the reaction in a derivative block
             rates = []
 
@@ -505,7 +503,7 @@ def ast(self, regions=None):
             return rates, species
         else:
             # represent the reaction in a kinetic block
-            rast = self._scheme.ast()
+            rast = self._scheme.ast(use_react_var=True)
 
             # fill in the correct rates
             if ">" in self._dir:

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

@@ -43,6 +43,13 @@
             ReactionStatement,
         )
 
+        def _int_or_double(c):
+            return (
+                Integer(c, Name(String(str(c))))
+                if isinstance(c, int)
+                else Double(str(c))
+            )
+
         # used for optimization
         from neuron.nmodl.visitor import AstLookupVisitor
         from neuron.nmodl import dsl
@@ -330,6 +337,7 @@ def ast(self):
         from . import Rate, rxd, MultiCompartmentReaction
         from .species import (
             SpeciesOnRegion,
+            Parameter,
             ParameterOnRegion,
             SpeciesOnExtracellular,
             ParameterOnExtracellular,
@@ -369,13 +377,18 @@ def parse_kinetic_block(reactions, blocks):
             s = sptr()
             sid = self._species_ids[s._id]
             for r in s._regions:
+                # skip regions not local here
+                if r._id not in self._region_ids:
+                    continue
                 rid = self._region_ids[r._id]
                 flat_name = s.ast(r).get_node_name()
                 self._name_to_index[flat_name] = ("species", sid, rid)
         for sptr in self._react_params:
             s = sptr()
             sid = self._params_ids[s._id]
             for r in s._regions:
+                if r._id not in self._region_ids:
+                    continue
                 rid = self._region_ids[r._id]
                 flat_name = s.ast(r).get_node_name()
                 self._name_to_index[flat_name] = ("params", sid, rid)
@@ -414,32 +427,62 @@ def parse_kinetic_block(reactions, blocks):
         rates = []
         multicompartmentReactions = []
         states = []
-        mc_flux = []
+        # maps for kinetic-blocks
         mc_kinetic_blocks = []
         local_consts = {}
+        mc_mult_count = 0
         for rptr, rlst in self._react_regions.items():
             r = rptr()
             rast, sp = r.ast(rlst)
             species += sp
             rast = rast if hasattr(rast, "__len__") else [rast]
             mc_rates_ast = []
             if isinstance(r, MultiCompartmentReaction):
-                for react in rast:
-                    if isinstance(react, ReactionStatement):
-                        mc_kinetic_blocks.append(react)
-                    else:
-                        mc_rates_ast.append(react)
-                multicompartmentReactions += mc_rates_ast
-                flux = []
+                # keep track for the fluxes and area/volume multipliers
+                mults = []
+                fluxes = []
                 for i, sp in enumerate(r._sources + r._dests):
                     s = sp()
-                    if r._membrane_flux and isinstance(s, SpeciesOnRegion):
-                        sid = self._species_ids[s._id]
-                        rid = self._region_ids[s._region()._id]
-                        flux.append((sid, rid, r._cur_charges[i]))
+                    if r._membrane_flux:
+                        if isinstance(s, SpeciesOnRegion) and not isinstance(
+                            s, ParameterOnRegion
+                        ):
+                            sid = self._species_ids[s._id]
+                            rid = self._region_ids[s._region()._id]
+                            flux = (sid, rid, r._cur_charges[i])
+
+                            # keep track of direction
+                            kin_coeff = (
+                                -r._cur_charges[i]
+                                if i < len(r._sources)
+                                else r._cur_charges[i]
+                            )
+                            flux = (sid, rid, r._cur_charges[i], kin_coeff)
+                        else:
+                            flux = None
                     else:
-                        flux.append(None)
-                mc_flux += flux
+                        flux = None
+                    if not isinstance(
+                        s, (Parameter, ParameterOnRegion, ParameterOnExtracellular)
+                    ):
+                        mults.append(mc_mult_count)
+                        fluxes.append(flux)
+                    mc_mult_count += 1
+                for react, mid, flux in zip(rast, mults, fluxes):
+                    if isinstance(react, ReactionStatement):
+                        # ReactionStatement and mult with both account for direction
+                        flx = [
+                            (f[0], f[1], f[3]) if f is not None else f for f in fluxes
+                        ]
+                        mc_kinetic_blocks.append((react, mults, flx))
+                    else:
+                        if flux is None:
+                            flx = None
+                        else:
+                            sid, rid, charge, _ = flux
+                            flx = (sid, rid, charge)
+                        mc_rates_ast.append((react, mid, flx))
+                multicompartmentReactions += mc_rates_ast
             elif (
                 isinstance(rptr(), Rate)
                 or _ast_config["kinetic_block"] == "off"
@@ -460,67 +503,84 @@ def parse_kinetic_block(reactions, blocks):
             )
         if states != []:
             blocks.append(StateBlock(states))
-
         # convert KineticBlock to DerivativeBlock to compile
         if reactions != []:
             rast, lc = parse_kinetic_block(reactions, blocks)
             rates += rast
             local_consts |= lc
-        if mc_kinetic_blocks != []:
-            rast, lc = parse_kinetic_block(mc_kinetic_blocks, blocks)
-            multicompartmentReactions += rast
+        mc_div_block_count = len(multicompartmentReactions)
+        mc_kinetic_blocks_mults = []
+        for react, mult, flux in mc_kinetic_blocks:
+            rast, lc = parse_kinetic_block(
+                [react],
+                blocks,
+            )
+
+            for react, mid, flx in zip(rast, mult, flux):
+                multicompartmentReactions.append((react, mid, flx))
             local_consts |= lc
+            mc_kinetic_blocks_mults += mult
         # Merge rates for common species or states
         if rates != [] or multicompartmentReactions != []:
             lookup = AstLookupVisitor()
             merged = {}
             constants = set()
-            mult_id = 0
-            for rid, stmt in enumerate(rates + multicompartmentReactions):
+            mc_stmt = 0
+
+            def _add_flux(flx, frhs):
+                # accumulate a membrane-flux contribution under key `flx`
+                if flx in merged:
+                    merged[flx] = (
+                        Name(String(flx)),
+                        BinaryExpression(
+                            merged[flx][1],
+                            BinaryOperator(BinaryOp.BOP_ADDITION),
+                            frhs,
+                        ),
+                    )
+                else:
+                    merged[flx] = (Name(String(flx)), frhs)
+                constants.add(flx)
+
+            for rid, statement in enumerate(rates + multicompartmentReactions):
+                if rid < len(rates):
+                    stmt = statement
+                else:
+                    stmt, mult_id, flux = statement
+
                 diffeq = stmt.expression
                 binexpr = diffeq.expression
                 var_name = binexpr.lhs.get_node_name()
                 rhs = ParenExpression(binexpr.rhs)
                 if rid >= len(rates):
                     # MCR have to be multiplied by mult[]
-                    mult = Name(String(f"_mult_{mult_id}"))
-                    if mc_flux[mult_id] is not None:
-                        sid, rid, charge = mc_flux[mult_id]
-                        flx = f"_flux_{sid}_{rid}_"
-                        fast = Name(String(flx))
-                        cast = (
-                            Integer(charge, Name(String(str(charge))))
-                            if isinstance(charge, int)
-                            else Double(str(charge))
-                        )
+                    if mult_id in mc_kinetic_blocks_mults:
+                        # Kinetic-block MC reaction
+                        # sign is already in d/dt so use abs(_mult[id])
+                        mult = Name(String(f"_absmult_{mult_id}"))
+                        constants.add(f"_absmult_{mult_id}")
+                    else:
+                        # Derivative MC reaction
+                        mult = Name(String(f"_mult_{mult_id}"))
+                        constants.add(f"_mult_{mult_id}")
+                    if flux is not None:
+                        sid, frid, charge = flux
                         if charge == 1:
                             frhs = rhs
                         else:
                             frhs = BinaryExpression(
-                                cast,
+                                _int_or_double(charge),
                                 BinaryOperator(BinaryOp.BOP_MULTIPLICATION),
                                 rhs,
                             )
-                        if flx in merged:
-                            merged[flx] = (
-                                fast,
-                                BinaryExpression(
-                                    merged[flx][1],
-                                    BinaryOperator(BinaryOp.BOP_ADDITION),
-                                    frhs,
-                                ),
-                            )
-                        else:
-                            merged[flx] = (fast, frhs)
-                        constants.add(flx)
-                    # multiple by a constant to scale units
+                        _add_flux(f"_flux_{sid}_{frid}_", frhs)
+
+                    # multiply by a constant to scale units
                     rhs = BinaryExpression(
                         mult,
                         BinaryOperator(BinaryOp.BOP_MULTIPLICATION),
                         rhs,
                     )
-                    constants.add(f"_mult_{mult_id}")
-                    mult_id += 1
                 if var_name in merged:
                     merged[var_name] = (
                         binexpr.lhs,
@@ -587,6 +647,7 @@ def parse_kinetic_block(reactions, blocks):
                 "var_names": list(merged.keys()),
                 "code": code,
                 "tmp_vars": tmp_vars,
+                "merged": merged,
             }
 
             blocks.append(

src/nrnpython/rxd.cpp

@@ -671,6 +671,7 @@ extern "C" NRN_EXPORT void setup_currents(int num_currents,
     // initialize memory here to allow currents from an intracellular species
     // with no corresponding nrn_region='o' or Extracellular species
     memset(induced_currents_ecs_idx, SPECIES_ABSENT, sizeof(int) * _memb_curr_total);
+    memset(induced_currents_grid_id, SPECIES_ABSENT, sizeof(int) * _memb_curr_total);
 
     for (i = 0, k = 0; i < num_currents; i++) {
         _memb_cur_ptrs[i].resize(num_species[i]);
@@ -727,7 +728,7 @@ extern "C" NRN_EXPORT void setup_currents(int num_currents,
 
             for (i = 0, k = 0; k < _memb_curr_total; k++) {
                 if (induced_currents_grid_id[k] == id)
-                    _rxd_induced_currents_scale[k] = current_scales[i];
+                    _rxd_induced_currents_scale[k] = current_scales[i++];
             }
         }
     }

test/rxd/test_ast.py

@@ -283,9 +283,19 @@ def test_multicompartment_reaction_ast(setup_section):
 
     ast = {
         "ReactionStatement": [
-            {"VarName": [{"Name": [{"String": [{"name": "ca_cyt_0"}]}]}]},
+            {
+                "ReactVarName": [
+                    {"Integer": [{"Name": [{"String": [{"name": "1"}]}]}]},
+                    {"VarName": [{"Name": [{"String": [{"name": "ca_cyt_0"}]}]}]},
+                ]
+            },
             {"ReactionOperator": [{"name": "<->"}]},
-            {"VarName": [{"Name": [{"String": [{"name": "ca_er_2"}]}]}]},
+            {
+                "ReactVarName": [
+                    {"Integer": [{"Name": [{"String": [{"name": "1"}]}]}]},
+                    {"VarName": [{"Name": [{"String": [{"name": "ca_er_2"}]}]}]},
+                ]
+            },
             {"Double": [{"name": "0.1"}]},
             {"Double": [{"name": "0.05"}]},
         ]