Generate AST from rxd rates, reactions and mulitcompartment reactions.

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

@@ -474,6 +474,87 @@ def needs_finite_differences(mat) -> bool:
     return any(isinstance(expr, sp.Derivative) for expr in sp.preorder_traversal(mat))
 
 
+def optimize_odes(
+    rhs_strings,
+    var_names,
+    constants,
+    local_consts,
+    function_calls,
+    rhs_ids=None,
+    do_cse=True,
+):
+    """Apply SymPy optimizations to ODE rate expressions.
+
+    Unlike solve_non_lin_system, this does not perform any time integration.
+    It takes the RHS of rate expressions f(x) from x' = f(x) and returns
+    optimised C code, suitable for operator-split solvers that handle
+    time-stepping externally. Used by the rxd module.
+
+    Args:
+        rhs_strings: list of RHS expression strings,
+                     e.g. ["-0.005*ca*cam + 0.01*cacam", ...]
+        var_names: list of all variable names (state vars + constants),
+                   e.g. ["ca", "cam", "cacam", "kf", "kb"]
+        constants: set of any other symbolic names used
+        local_consts: dictionary of values to substitute in in place of variable names
+        function_calls: set of function call names used in the ODEs
+        do_cse: if True (default), apply Common Subexpression Elimination
+
+    Returns:
+        Tuple of (code, new_local_vars) where:
+        - code: list of C assignment strings,
+                e.g. ["tmp_0 = ca*cam", "rhs[0] = -0.005*tmp_0 + ..."]
+        - 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)):
+        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)
+    rhs_exprs = [sp.sympify(rhs, locals=sympy_vars) for rhs in rhs_strings]
+    # Don't apply sympy simplifications, they don't seem to improve performance
+    # and they can hang or crash for more complex reactions.
+
+    code = []
+    local_vars = []
+    if do_cse:
+        my_symbols = sp.utilities.iterables.numbered_symbols(prefix="tmp_")
+        sub_exprs, reduced = sp.cse(
+            rhs_exprs,
+            symbols=my_symbols,
+            optimizations="basic",
+            order="canonical",
+        )
+        for var, expr in sub_exprs:
+            local_vars.append(sp.ccode(var))
+            code.append(
+                f"{var} = {sp.ccode(expand_pow(expr).subs(local_consts).evalf(), user_functions=custom_fcts)}"
+            )
+        rhs_exprs = reduced
+
+    for i, expr in enumerate(rhs_exprs):
+        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:
+            code.append(f"rhs[{i}] = {rhs}")
+
+    return code, local_vars
+
+
 def solve_non_lin_system(
     eq_strings: Iterable[str],
     vars: Iterable[str],

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

@@ -7,8 +7,8 @@
 from .reaction import Reaction
 from . import geometry
 from .multiCompartmentReaction import MultiCompartmentReaction
-from .rxd import re_init, set_solve_type, nthread
-from .rxdmath import v
+from .rxd import re_init, set_solve_type, nthread, ast
+from .rxdmath import v, _ast_config
 
 try:
     from . import dimension3

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

@@ -9,6 +9,20 @@
 from neuron import h
 import itertools
 from .rxdException import RxDException
+from .rxdmath import _ast_config, _ast_check
+
+if _ast_config["nmodl_support"]:
+    try:
+        from neuron.nmodl.ast import (
+            ExpressionStatement,
+            DiffEqExpression,
+            BinaryExpression,
+            BinaryOperator,
+            BinaryOp,
+        )
+    except ModuleNotFoundError as e:
+        _ast_config["nmodl_support"] = False
+        _ast_config["exception"] = e
 from typing import Any, Optional
 
 
@@ -224,6 +238,7 @@
         #        regs.append(sptr()._extracellular()._region)
         #    else:
         #        regs.append(sptr()._region())
+        self._arithmetic_rate = rxdmath._ensure_arithmeticed(rate)
         self._rate, self._involved_species = rxdmath._compile(rate, regs)
 
     @property
@@ -402,3 +417,93 @@
                 self._cur_ptrs.append(tuple(local_ptrs))
                 self._cur_mapped.append(tuple(local_mapped))
                 self._cur_mapped_ecs.append(local_mapped_ecs)
+
+    def ast(self, regions=None):
+        """Provide an AST representation of the mutlicompartment reactions.
+
+        Args:
+            regions (optional): argument for compatability -- it is ignored
+            the regions are specified when defining the multicompartment reaction.
+
+        Depending on rxd._ast_config["kinetic_block"] if 'off' (default) or
+        if rxd._ast_config["kinetic_block"] if 'mass_action' and the reaction has
+        custom dynamics then the reaction will be represented as as list of
+        DiffEqExpression each wrapped in ExpressionStatement.
+        If rxd._ast_config["kinetic_block"] if 'on' or
+        if rxd._ast_config["kinetic_block"] if 'mass_action' and the reaction
+        has mass action kinetics then the reaction will be represent as a
+        ReactionStatement.
+
+
+        Returns:
+            List[nmodl.ast]: A list of ExpressionStatement or single ReactionStatement
+            List[str]:  A list of the species (AST state names)
+        """
+        from .species import Parameter, ParameterOnRegion, ParameterOnExtracellular
+
+        kinetic_block = _ast_check()
+
+        if not initializer.is_initialized():
+            initializer._do_init()
+
+        # assume all source share the same region
+        src = self._sources[0]()
+        # assume all dests share the same region
+        dst = self._dests[0]()
+
+        lreg = (
+            src._region() if hasattr(src, "_region") else src._extracellular()._region
+        )
+        rreg = (
+            dst._region() if hasattr(dst, "_region") else dst._extracellular()._region
+        )
+
+        species = []
+        for sp in self._sources + self._dests:
+            if not isinstance(
+                sp, (Parameter, ParameterOnRegion, ParameterOnExtracellular)
+            ):
+                species.append(sp().ast().get_node_name())
+
+        if kinetic_block == "off" or self._custom_dynamics:
+            # represent the reaction in a derivative block
+            rates = []
+
+            for idx, sptr in enumerate(self._sources + self._dests):
+                sp = sptr()
+                if isinstance(
+                    sp, (Parameter, ParameterOnRegion, ParameterOnExtracellular)
+                ):
+                    continue
+                dx = sptr().ast(prime=True)
+                if idx < len(self._sources):
+                    rast = self._arithmetic_rate.ast([lreg, rreg])
+                else:
+                    rast = (self._arithmetic_rate).ast([rreg, lreg])
+                rates.append(
+                    ExpressionStatement(
+                        DiffEqExpression(
+                            (
+                                BinaryExpression(
+                                    dx, BinaryOperator(BinaryOp.BOP_ASSIGN), rast
+                                )
+                            )
+                        )
+                    )
+                )
+            return rates, species
+        else:
+            # represent the reaction in a kinetic block
+            rast = self._scheme.ast(use_react_var=True)
+
+            # fill in the correct rates
+            if ">" in self._dir:
+                rast.expression1 = rxdmath._ensure_arithmeticed(self.f_rate).ast(
+                    [lreg, rreg]
+                )
+            if "<" in self._dir:
+                rast.expression2 = rxdmath._ensure_arithmeticed(self.b_rate).ast(
+                    [rreg, lreg]
+                )
+
+        return rast, species

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

@@ -5,6 +5,20 @@
 import itertools
 import warnings
 from .generalizedReaction import GeneralizedReaction
+from .rxdmath import _ast_config, _ast_check
+
+if _ast_config["nmodl_support"]:
+    try:
+        from neuron.nmodl.ast import (
+            BinaryOp,
+            BinaryOperator,
+            BinaryExpression,
+            DiffEqExpression,
+            ExpressionStatement,
+        )
+    except ModuleNotFoundError as e:
+        _ast_config["nmodl_support"] = False
+        _ast_config["exception"] = e
 from typing import Union, Optional, Any
 
 # aliases to avoid repeatedly doing multiple hash-table lookups
@@ -350,3 +364,61 @@ def _get_memb_flux(self, states) -> list:
             return self._memb_scales * rates
         else:
             return []
+
+    def ast(self, regions=None):
+        """Provide an AST representation of the rate.
+
+        Args:
+            regions (List[weakref.ref]): A list of weak reference `rxd.Region`
+                                         if None, all regions where the Rate
+                                         is valid are used.
+
+        Returns:
+            List[nmodl.ast]: A list of ASTs for each region in regions.
+            List[str]: A list of the species (AST state names).
+        """
+        from .species import Parameter, ParameterOnRegion, ParameterOnExtracellular
+
+        _ast_check()
+
+        if not initializer.is_initialized():
+            initializer._do_init()
+        sp = self._species()
+        if isinstance(sp, (Parameter, ParameterOnRegion, ParameterOnExtracellular)):
+            return [], []
+
+        def get_ast(region):
+            if sp and hasattr(sp, "name"):
+                name = sp.ast(region).get_node_name()
+                dx = sp.ast(region, prime=True)
+            elif sp and hasattr(sp, "_species") and sp._species():
+                name = sp.ast().get_node_name()
+                dx = sp.ast(prime=True)
+            else:
+                raise RxDException(f"Unknown species: {sp}")
+            rate = rxdmath._ensure_arithmeticed(self._original_rate).ast(region)
+            return (
+                ExpressionStatement(
+                    DiffEqExpression(
+                        (
+                            BinaryExpression(
+                                dx, BinaryOperator(BinaryOp.BOP_ASSIGN), rate
+                            )
+                        )
+                    )
+                ),
+                name,
+            )
+
+        diff = []
+        species = []
+        if regions is not None:
+            regs = [rptr() for rptr in regions if rptr()]
+        else:
+            regs = self._active_regions
+
+        for region in regs:
+            d, s = get_ast(region)
+            diff.append(d)
+            species.append(s)
+        return diff, species