diff --git a/share/lib/python/neuron/nmodl/ode.py b/share/lib/python/neuron/nmodl/ode.py index 7e1fbfb6a6..d58b2256fd 100644 --- a/share/lib/python/neuron/nmodl/ode.py +++ b/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], diff --git a/share/lib/python/neuron/rxd/__init__.py b/share/lib/python/neuron/rxd/__init__.py index b48a90a606..5faa1584b5 100644 --- a/share/lib/python/neuron/rxd/__init__.py +++ b/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 diff --git a/share/lib/python/neuron/rxd/multiCompartmentReaction.py b/share/lib/python/neuron/rxd/multiCompartmentReaction.py index 9a281b9df3..35fbffc41d 100644 --- a/share/lib/python/neuron/rxd/multiCompartmentReaction.py +++ b/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 @@ def _update_rates(self): # 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 @@ def _do_memb_scales(self, cur_map): 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 diff --git a/share/lib/python/neuron/rxd/rate.py b/share/lib/python/neuron/rxd/rate.py index 084999527a..5717490d3d 100644 --- a/share/lib/python/neuron/rxd/rate.py +++ b/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 diff --git a/share/lib/python/neuron/rxd/reaction.py b/share/lib/python/neuron/rxd/reaction.py index d38a71eb35..7c6d7d8e72 100644 --- a/share/lib/python/neuron/rxd/reaction.py +++ b/share/lib/python/neuron/rxd/reaction.py @@ -5,6 +5,24 @@ get_scheme_rate1_rate2_regions_custom_dynamics_mass_action, ) from .rxdException import RxDException +from .rxdmath import _ast_config, _ast_check + +if _ast_config["nmodl_support"]: + try: + from neuron.nmodl.ast import ( + Double, + Compartment, + ExpressionStatement, + DiffEqExpression, + BinaryExpression, + BinaryOperator, + BinaryOp, + Name, + String, + ) + except ModuleNotFoundError as e: + _ast_config["nmodl_support"] = False + _ast_config["exception"] = e from typing import Any @@ -259,3 +277,113 @@ def __repr__(self) -> str: def _do_memb_scales(self) -> None: # nothing to do since NEVER a membrane flux pass + + def ast(self, regions=None): + """Provide an AST representation of the reactions. + + Args: + regions (List[weakref.ref]): A list of weak references `rxd.Region` + if None all regions where the rate is + valid are used. + + 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 ReactionStatement + List[str]: A list of the species (AST state names) + List[nmodl.ast]: A list of Compartment if using ReactionStatement otherwise empty + """ + from .species import Parameter, ParameterOnRegion, ParameterOnExtracellular + + kinetic_block = _ast_check() + + if not initializer.is_initialized(): + initializer._do_init() + + def get_ast(region): + if kinetic_block == "off" or ( + kinetic_block == "mass_action" and self._custom_dynamics + ): + rate = self._rate_arithmeticed + frate = (-rate).ast(region) + brate = (rate).ast(region) + diff, species = [], [] + for idx, sref in enumerate(self._sources + self._dests): + sp = sref() + if isinstance( + sp, (Parameter, ParameterOnRegion, ParameterOnExtracellular) + ): + continue + rast = frate if idx < len(self._sources) else brate + 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}") + diff.append( + ExpressionStatement( + DiffEqExpression( + ( + BinaryExpression( + dx, BinaryOperator(BinaryOp.BOP_ASSIGN), rast + ) + ) + ) + ) + ) + species.append(name) + return diff, species + else: + react = self._scheme.ast(region, use_react_var=True) + # replace placeholder rates + if hasattr(self.f_rate, "ast"): + kf = self.f_rate.ast(region) + else: + kf = Double(str(self.f_rate)) + + if hasattr(self.b_rate, "ast"): + kb = self.b_rate.ast(region) + else: + kb = Double(str(self.b_rate)) + react.expression1 = kf + react.expression2 = kb + species = [] + for sref in set(self._sources + self._dests): + sp = sref() + if isinstance( + sp, (Parameter, ParameterOnRegion, ParameterOnExtracellular) + ): + continue + if sp and hasattr(sp, "name"): + name = sp.ast(region).get_node_name() + elif sp and hasattr(sp, "_species") and sp._species(): + name = sp.ast().get_node_name() + else: + raise RxDException(f"Unknown species: {sp}") + species.append(name) + return [react], species + + reactions = [] + species = [] + + if regions is None: + regs = self._react_regions + else: + regs = [rptr() for rptr in regions if rptr()] + + for region in regs: + r, s = get_ast(region) + reactions += r + species += s + return reactions, species diff --git a/share/lib/python/neuron/rxd/region.py b/share/lib/python/neuron/rxd/region.py index 943fbaac2b..96ee81eb81 100644 --- a/share/lib/python/neuron/rxd/region.py +++ b/share/lib/python/neuron/rxd/region.py @@ -15,6 +15,48 @@ import warnings import math import ctypes +from .rxdmath import _ast_config + +if _ast_config["nmodl_support"]: + try: + from neuron.nmodl.ast import ( + Name, + String, + Integer, + Double, + DerivativeBlock, + StatementBlock, + KineticBlock, + StateBlock, + AssignedDefinition, + Program, + NeuronBlock, + Suffix, + ExpressionStatement, + DiffEqExpression, + ParenExpression, + BinaryExpression, + BinaryOperator, + BinaryOp, + AstNodeType, + LocalListStatement, + 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 + from neuron.nmodl.ode import optimize_odes + except ModuleNotFoundError as e: + _ast_config["nmodl_support"] = False + _ast_config["exception"] = e _all_regions = [] _region_count = 0 @@ -76,6 +118,8 @@ class _c_region: "_region_ids", "_voltage_dependent", "_vptrs", + "_name_to_index", + "_optimized_rates", ) def __init__(self, regions: Any) -> None: @@ -99,6 +143,8 @@ def __init__(self, regions: Any) -> None: self._ecs_params_ids = None self._voltage_dependent = False self._vptrs = None + self._name_to_index = None + self._optimized_rates = None for rptr in self._regions: r = rptr() self._overlap = h.SectionList( @@ -284,6 +330,344 @@ def _initalize(self): self._ecs_initalize() self._initialized = True + def ast(self): + """Return AST for the set of reactions""" + if not _ast_config["nmodl_support"]: + return + from . import Rate, rxd, MultiCompartmentReaction + from .species import ( + SpeciesOnRegion, + Parameter, + ParameterOnRegion, + SpeciesOnExtracellular, + ParameterOnExtracellular, + ) + + # if not self._initialized: + # self._initalize() + + def parse_kinetic_block(reactions, blocks): + rate_stmts = [] + local_consts = {} + lookup = AstLookupVisitor() + kinetic_blocks = [ + KineticBlock(Name(String("reactions")), [], StatementBlock(reactions)) + ] + tmp_ast = Program(blocks + kinetic_blocks) + dsl.symtab.SymtabVisitor().visit_program(tmp_ast) + dsl.visitor.KineticBlockVisitor().visit_program(tmp_ast) + stmts = lookup.lookup(tmp_ast, AstNodeType.DERIVATIVE_BLOCK) + for stmt in stmts: + for sb in stmt.statement_block.statements: + if isinstance(sb, LocalListStatement): + for nnode in sb.variables: + local_consts[nnode.get_node_name()] = None + elif hasattr(sb, "expression"): + ex = sb.expression + if isinstance(ex, BinaryExpression): + local_consts[ex.lhs.get_node_name()] = dsl.to_nmodl(ex.rhs) + + else: + rate_stmts.append(sb) + return rate_stmts, local_consts + + # Build mapping: AST flat name -> ("species"|"params", local_id, region_id) + self._name_to_index = {} + for sptr in self._react_species: + 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) + + # ECS species/params name mapping + if self.num_ecs_species + self.num_ecs_params > 0: + from .species import _all_species as _asp + + # Build reverse lookup: _ExtracellularSpecies id -> (parent Species, ecs region) + _ecs_to_parent = {} + for sp_ref in _asp: + sp = sp_ref() + if sp is not None and hasattr(sp, "_extracellular_instances"): + for ecs_region, ecs_obj in sp._extracellular_instances.items(): + _ecs_to_parent[id(ecs_obj)] = (sp, ecs_region) + + for sptr in self._ecs_react_species: + ecs_inst = sptr() + ecs_sid = self._ecs_species_ids[ecs_inst._grid_id] + if id(ecs_inst) in _ecs_to_parent: + parent_sp, ecs_reg = _ecs_to_parent[id(ecs_inst)] + flat_name = parent_sp[ecs_reg].ast().get_node_name() + self._name_to_index[flat_name] = ("species_3d", ecs_sid) + + for sptr in self._ecs_react_params: + ecs_inst = sptr() + ecs_sid = self._ecs_params_ids[ecs_inst._grid_id] + if id(ecs_inst) in _ecs_to_parent: + parent_sp, ecs_reg = _ecs_to_parent[id(ecs_inst)] + flat_name = parent_sp[ecs_reg].ast().get_node_name() + self._name_to_index[flat_name] = ("params_3d", ecs_sid) + + species = [] + blocks = [] + reactions = [] + rates = [] + multicompartmentReactions = [] + states = [] + # 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): + # 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: + 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 = 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" + or ( + _ast_config["kinetic_block"] == "mass_action" + and rptr()._custom_dynamics + ) + ): + rates += rast + else: + reactions += rast + + for name in set(species): + states.append( + AssignedDefinition( + Name(String(name)), None, None, None, None, None, None + ) + ) + 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 + 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 r, mid, flx in zip(rast, mult, flux): + multicompartmentReactions.append((r, 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() + 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[] + 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( + _int_or_double(charge), + BinaryOperator(BinaryOp.BOP_MULTIPLICATION), + rhs, + ) + _add_flux(f"_flux_{sid}_{frid}_", frhs) + + # multiply by a constant to scale units + rhs = BinaryExpression( + mult, + BinaryOperator(BinaryOp.BOP_MULTIPLICATION), + rhs, + ) + if var_name in merged: + merged[var_name] = ( + binexpr.lhs, + BinaryExpression( + merged[var_name][1], + BinaryOperator(BinaryOp.BOP_ADDITION), + rhs, + ), + ) + else: + merged[var_name] = (binexpr.lhs, rhs) + + # Collect RHS strings, constants, and function calls + # and use CSE can eliminate common subexpressions + all_rhs_strings = [] + function_calls = set() + group_stmts_list = [] + + stmts = [] + for _, (lhs, rhs) in merged.items(): + stmts.append( + ExpressionStatement( + DiffEqExpression( + BinaryExpression( + lhs, + BinaryOperator(BinaryOp.BOP_ASSIGN), + rhs, + ) + ) + ) + ) + all_rhs_strings.append(dsl.to_nmodl(rhs)) + for ref in lookup.lookup(rhs, AstNodeType.VAR_NAME): + ref_name = ref.get_node_name() + if ref_name not in merged: + constants.add(ref_name) + for fc in lookup.lookup(rhs, AstNodeType.FUNCTION_CALL): + function_calls.add(fc.get_node_name()) + + rhs_ids = [] + for vname in merged: + if "_flux_" in vname: + _, _, sid, rid, _ = vname.split("_") + rhs_ids.append(f"flux[{sid}][{rid}]") + else: + spids = self._name_to_index[vname] + if len(spids) == 2: + _, sid = spids + rhs_ids.append(f"rhs_3d[{sid}]") + else: + _, sid, rid = spids + rhs_ids.append(f"rhs[{sid}][{rid}]") + # Run CSE + simplification globally across all groups + code, tmp_vars = optimize_odes( + all_rhs_strings, + list(merged.keys()), + list(constants), + local_consts, + function_calls, + rhs_ids=rhs_ids, + ) + + self._optimized_rates = { + "var_names": list(merged.keys()), + "code": code, + "tmp_vars": tmp_vars, + "merged": merged, + } + + blocks.append( + DerivativeBlock( + Name(String("rates")), + StatementBlock(stmts), + ) + ) + + return Program( + [ + NeuronBlock( + StatementBlock( + [Suffix(Name(String("SUFFIX")), Name(String("rxd")))] + ) + ) + ] + + blocks + ) + class Extracellular: """Declare an extracellular region diff --git a/share/lib/python/neuron/rxd/rxd.py b/share/lib/python/neuron/rxd/rxd.py index f508ffcf05..e950a99e28 100644 --- a/share/lib/python/neuron/rxd/rxd.py +++ b/share/lib/python/neuron/rxd/rxd.py @@ -568,6 +568,98 @@ def _cxx_compile(formula): return reaction +def _compile_reactions_from_ast(creg): + """Build and compile the reaction callback using AST-optimized rates. + + Uses creg._optimized_rates (CSE'd C code) and creg._name_to_index + (AST name -> array index mapping) populated by creg.ast(). + + Returns (compiled_fn, mc_mult_count, mc_mult_list) or None if the AST + path is not available. + """ + from .rxdmath import _ast_config + + if not _ast_config.get("nmodl_support"): + return None + + if not initializer.is_initialized(): + initializer._do_init() + creg.ast() + + has_rates = hasattr(creg, "_optimized_rates") and creg._optimized_rates + has_mcr = hasattr(creg, "_mcr_rates") and creg._mcr_rates + if not has_rates and not has_mcr: + return None + + name_to_index = creg._name_to_index + + # Sort names longest-first to avoid partial matches + sorted_names = sorted(name_to_index.keys(), key=len, reverse=True) + + def substitute_names(line): + """Replace AST variable names with array access expressions. + And _mult_id with mult[id] for multicompartment reactions. + ICS entries (3-tuple): species[sid][rid] or params[sid][rid] + ECS entries (2-tuple): species_3d[ecs_sid] or params_3d[ecs_sid] + """ + result = re.sub(r"_mult_(\d+)", r"mult[\1]", line) + for name in sorted_names: + entry = name_to_index[name] + if len(entry) == 3: + kind, sid, rid = entry + replacement = f"{kind}[{sid}][{rid}]" + else: + kind, sid = entry + replacement = f"{kind}[{sid}]" + result = re.sub( + r"\b" + re.escape(name) + r"\b", + replacement, + result, + ) + return result + + # Build mc_mult_list for MultiCompartmentReactions in this c_region + from . import multiCompartmentReaction + + mc_mult_count = 0 + mc_mult_list = [] + for rptr in creg._react_regions: + r = rptr() + if isinstance(r, multiCompartmentReaction.MultiCompartmentReaction): + mc_mult_count += len(r._sources) + len(r._dests) + mc_mult_list.extend(r._mult.flatten()) + + fxn_string = _c_headers + fxn_string += ( + "void reaction(double** species, double** params, double** rhs, " + "double* mult, double* species_3d, double* params_3d, " + "double* rhs_3d, double** flux, double v)\n{" + ) + + tmps = creg._optimized_rates["tmp_vars"] + if tmps: + fxn_string += "\n\tdouble " + ", ".join(tmps) + ";" + + # create optimized code function string + flux_lines = [] + for line in creg._optimized_rates["code"]: + # Replace symbolic variable names with array access + c_line = substitute_names(line) + if c_line.startswith("flux["): + flux_lines.append(c_line) + else: + fxn_string += f"\n\t{c_line};" + + # group of membrane fluxes + if flux_lines: + fxn_string += "\n\tif (flux)\n\t{" + for line in flux_lines: + fxn_string += f"\n\t\t{line};" + fxn_string += "\n\t}" + fxn_string += "\n}\n}\n" + return _cxx_compile(fxn_string), mc_mult_count, mc_mult_list + + _h_ptrvector = h.PtrVector _h_vector = h.Vector @@ -1089,7 +1181,7 @@ def _get_node_indices(species, region, sec3d, x3d, sec1d, x1d): ): indices3d.append(_point_indices[region][point]) vols3d.append(surf[point][0] if point in surf else region.dx ** 3) - # print f'found node {node._index} with coordinates ({node.x3d:g}, {node.y3d:g}, {node.z3d:g})' + # print 'found node %d with coordinates (%g, %g, %g)' % (node._index, node.x3d, node.y3d, node.z3d) # discard duplicates... # TODO: really, need to figure out all the 3d nodes connecting to a given 1d endpoint, then unique that # print f'3d matrix indices: {indices3d!r}' @@ -1123,6 +1215,55 @@ def _get_node_indices(species, region, sec3d, x3d, sec1d, x1d): return index_1d, indices3d, vol1d, vols3d +def ast(): + """ + Generates an Abstract Syntax Tree (AST) representation of the rxd model. + Reactions are grouped by the sections they have in common. This function + iterates over groups of reactions and returns their AST representations. + + ### **Behavior Based on `rxd._ast_config["kinetic_block"]` Setting**: + - `"off"`: Default beahviour. Converts all reactions into `DERIVATIVE` + blocks (`nmodl.ast.DerivativeBlock`). + - `"on"`: Converts reactions into `KINETIC` blocks (`nmodl.ast.KineticBlock`) and + rate equations into `DERIVATIVE` blocks. + - `"mass_action"`: Uses `KINETIC` blocks for **mass-action reactions** while + keeping **of AST nodes (`nmodl.ast.Program`)** representing the rxd model. + + ### **Example Use Case**: + ```python + from neuron import h, rxd + from neuron.units import nM, mV + from nmodl.ast import view + + # create a simple rxd model -- calcium buffering + # Where -- soma + soma = h.Section(name="soma") + cyt = rxd.Region([soma], name="cyt") + + # Who -- calcium and a buffer + ca = rxd.Species(cyt, name="ca", charge=2, initial=60 * nM) + buf = rxd.Species(cyt, name="buf", initial=10 * nM) + cabuf = rxd.Species(cyt, name="cabuf", initial=0) + + # What -- buffering mass-action reaction + buffering = rxd.Reaction(ca + buf, cabuf, 0.02, 0.01) + + # View the AST + view(rxd.ast()[0]) + + Returns: + List[nmodl.ast.AstNode]: A list of AST nodes representing transformed reactions. + """ + + if not initializer.is_initialized(): + _init() + grouped_reactions = [] + for cr in region._c_region_lookup.values(): + if cr[0] not in grouped_reactions: + grouped_reactions.append(cr[0]) + return [cr.ast() for cr in grouped_reactions] + + def _compile_reactions(): # clear all previous reactions (intracellular & extracellular) and the # supporting indexes @@ -1343,6 +1484,27 @@ def localize_index(creg, rate): if not creg._react_regions: continue creg._initalize() + + # Try AST-optimized compile + ast_result = _compile_reactions_from_ast(creg) + if ast_result is not None: + ast_compiled, mc_mult_count, mc_mult_list = ast_result + register_rate( + creg.num_species, + creg.num_params, + creg.num_regions, + creg.num_segments, + creg.get_state_index(), + creg.num_ecs_species, + creg.num_ecs_params, + creg.get_ecs_species_ids(), + creg.get_ecs_index(), + mc_mult_count, + numpy.array(mc_mult_list, dtype=ctypes.c_double), + _list_to_pyobject_array(creg._vptrs), + ast_compiled, + ) + continue mc_mult_count = 0 mc_mult_list = [] species_ids_used = numpy.zeros((creg.num_species, creg.num_regions), bool) diff --git a/share/lib/python/neuron/rxd/rxdmath.py b/share/lib/python/neuron/rxd/rxdmath.py index 618412cc29..10fb92e5d3 100644 --- a/share/lib/python/neuron/rxd/rxdmath.py +++ b/share/lib/python/neuron/rxd/rxdmath.py @@ -4,6 +4,87 @@ from . import initializer from typing import Union, Any, Callable, Optional +# _ast_config used for two flags +# nmodl_support default to True -- will be set to False if nmodl is not +# installed of install without python bindings + +# "kinetic_block" determines how to handle reactions +# set to "off" for ast to only use DERIVATIVE blocks +# set to "on" for reactions to be KINETIC blocks +# set to "mass_action" for mass action reaction to be KINETIC blocks while +# non-mass action reactions and rates are DERIVATIVE blocks + +_ast_config = {"nmodl_support": True, "kinetic_block": "off"} + + +if _ast_config["nmodl_support"]: + try: + from neuron.nmodl.ast import ( + Name, + String, + ReactionStatement, + VarName, + Integer, + BinaryOperator, + BinaryOp, + BinaryExpression, + UnaryOperator, + UnaryOp, + UnaryExpression, + ReactionOperator, + ReactVarName, + FunctionCall, + Double, + LocalVar, + ParenExpression, + ) + + OpPrecedence = { + BinaryOp.BOP_ADDITION: 1, + BinaryOp.BOP_SUBTRACTION: 1, + BinaryOp.BOP_MULTIPLICATION: 2, + BinaryOp.BOP_DIVISION: 2, + BinaryOp.BOP_POWER: 3, + } + + def needBrackets(parent, child): + if OpPrecedence[child] < OpPrecedence[parent]: + return True + if ( + parent in (BinaryOp.BOP_DIVISION, BinaryOp.BOP_SUBTRACTION) + and OpPrecedence[child] == OpPrecedence[parent] + ): + return True + return False + + def ParenBinaryExpression(lhs, op, rhs): + """Add parenthesis to a BinaryExpression if needed""" + + if lhs.is_binary_expression(): + if needBrackets(op, lhs.op.value): + lhs = ParenExpression(lhs) + + if rhs.is_binary_expression(): + if needBrackets(op, rhs.op.value): + rhs = ParenExpression(rhs) + return BinaryExpression(lhs, BinaryOperator(op), rhs) + + except ModuleNotFoundError as e: + _ast_config["nmodl_support"] = False + _ast_config["exception"] = e + + +def _ast_check(): + """Raise an exception if AST is not supported, otherwise return the "kinetic_block" configuration.""" + if not _ast_config["nmodl_support"]: + if "exception" in _ast_config: + raise _ast_config["exception"] + else: + raise RxDException( + 'NMODL AST are disabled set rxd._ast_config["nmodl_support"] to True' + ) + return _ast_config["kinetic_block"] + def _vectorized(f: Callable, objs: Any) -> Any: if hasattr(objs, "__len__"): @@ -201,6 +282,15 @@ def _voltage_dependent(self): except AttributeError: return False + def ast(self, region=None, use_react_var=False): + if _ast_config["nmodl_support"]: + if hasattr(self._obj, "ast"): + obj = self._obj.ast(region) + else: + obj = Name(String(self._obj)) + fun = Name(String(self._fname)) + return FunctionCall(fun, [obj]) + class _Function2: def __init__(self, obj1: Any, obj2: Any, f: Callable, fname: str) -> None: @@ -248,6 +338,20 @@ def _voltage_dependent(self): pass return False + def ast(self, region=None, use_react_var=False): + if _ast_config["nmodl_support"]: + if hasattr(self._obj1, "ast"): + obj1 = self._obj1.ast(region) + else: + obj1 = Name(String(self._obj1)) + if hasattr(self._obj2, "ast"): + obj2 = self._obj2.ast(region) + else: + obj2 = Name(String(self._obj2)) + + fun = Name(String(self._fname)) + return FunctionCall(fun, [obj1, obj2]) + # wrappers for the functions in module math from python 2.7 def acos(obj): @@ -530,6 +634,18 @@ def _involved_species(self, the_dict): self._a._involved_species(the_dict) self._b._involved_species(the_dict) + def ast(self, region=None, use_react_var=False): + if _ast_config["nmodl_support"]: + if hasattr(self._a, "ast"): + lhs = self._a.ast(region) + else: + lhs = Name(String(self._a)) + if hasattr(self._b, "ast"): + rhs = self._b.ast(region) + else: + rhs = Name(String(self._b)) + return ParenBinaryExpression(lhs, BinaryOp.BOP_MULTIPLICATION, rhs) + class _Quotient: def __init__(self, a, b): @@ -539,6 +655,18 @@ def __init__(self, a, b): def __repr__(self): return f"({self._a!r})/({self._b!r})" + def ast(self, region=None, use_react_var=False): + if _ast_config["nmodl_support"]: + if hasattr(self._a, "ast"): + lhs = self._a.ast(region) + else: + lhs = Name(String(self._a)) + if hasattr(self._b, "ast"): + rhs = self._b.ast(region) + else: + rhs = Name(String(self._b)) + return ParenBinaryExpression(lhs, BinaryOp.BOP_DIVISION, rhs) + # Change any Species to _ExtracellularSpecies so _semi_compile gives the # _grid_id and not the species _id def _ensure_extracellular(self, extracellular=None, intracellular3d=None): @@ -584,6 +712,18 @@ def _involved_species(self, the_dict): self._a._involved_species(the_dict) self._b._involved_species(the_dict) + def ast(self, region=None, use_react_var=False): + if _ast_config["nmodl_support"]: + if hasattr(self._a, "ast"): + lhs = self._a.ast(region) + else: + lhs = Name(String(self._a)) + if hasattr(self._b, "ast"): + rhs = self._b.ast(region) + else: + rhs = Name(String(self._b)) + return ParenBinaryExpression(lhs, BinaryOp.BOP_DIVISION, rhs) + class _Reaction: def __init__(self, lhs, rhs, direction): @@ -607,6 +747,23 @@ def _voltage_dependent(self): pass return False + def ast(self, region=None, use_react_var=False): + if _ast_config["nmodl_support"]: + lhs = self._lhs.ast(region, use_react_var) + rhs = self._rhs.ast(region, use_react_var) + # TODO: Placeholder rates should be replaced by in rxd.Reaction + if region: + rid = region._id + rname = region.name if region.name else "" + rint = Integer(rid, Name(String(rname))) + else: + rint = Integer(0, Name(String("unassigned"))) + + # should be replaced by rxd.Reaction with the actual values + kf = Double("0.0") + kb = Double("0.0") + return ReactionStatement(lhs, ReactionOperator(), rhs, kf, kb) + class _Arithmeticed: def __init__(self, item, valid_reaction_term=True): @@ -730,6 +887,63 @@ def __repr__(self): result = "0" return result + def ast(self, region=None, use_react_var=False): + if _ast_config["nmodl_support"]: + from . import species + + nodes = None + for item, count in zip( + list(self._items.keys()), list(self._items.values()) + ): + if hasattr(item, "ast"): + item_node = item.ast(region) + elif isinstance(item, int): + item_node = Integer(item, Name(String(str(item)))) + elif isinstance(item, float): + item_node = Double(str(item)) + else: + item_node = LocalVar(Name(String(item))) + + if count == 1: + if use_react_var and item_node.is_var_name(): + term_node = ReactVarName( + Integer(1, Name(String("1"))), item_node + ) + else: + term_node = item_node + elif count == -1: + if len(self._items) == 1: + term_node = UnaryExpression( + UnaryOperator(UnaryOp.UOP_NEGATION), item_node + ) + else: + term_node = item_node + else: + x = count if len(self._items) == 1 else abs(count) + if use_react_var and item_node.is_var_name(): + term_node = ReactVarName( + Integer(x, Name(String(str(x)))), item_node + ) + else: + term_node = ParenBinaryExpression( + Integer(x, Name(String(str(x)))), + BinaryOp.BOP_MULTIPLICATION, + item_node, + ) + + if nodes is None: + nodes = term_node + else: + if count < 0: + nodes = ParenBinaryExpression( + nodes, BinaryOp.BOP_SUBTRACTION, term_node + ) + else: + nodes = ParenBinaryExpression( + nodes, BinaryOp.BOP_ADDITION, term_node + ) + return nodes + @property def _voltage_dependent(self): for item in self._items: @@ -880,6 +1094,10 @@ def __repr__(self): def _voltage_dependent(self): return True + def ast(self, region=None, use_react_var=False): + if _ast_config["nmodl_support"]: + return VarName(Name(String("v")), None, None) + def __init__(self): super(Vm, self).__init__(Vm._Vm(), valid_reaction_term=True) diff --git a/share/lib/python/neuron/rxd/species.py b/share/lib/python/neuron/rxd/species.py index b649de6898..48d50765f6 100644 --- a/share/lib/python/neuron/rxd/species.py +++ b/share/lib/python/neuron/rxd/species.py @@ -10,6 +10,21 @@ from .rxdException import RxDException from . import initializer from collections.abc import Callable +from .rxdmath import _ast_config + +if _ast_config["nmodl_support"]: + try: + from neuron.nmodl.ast import ( + PrimeName, + ReactVarName, + VarName, + Name, + String, + Integer, + ) + except ModuleNotFoundError as e: + _ast_config["nmodl_support"] = False + _ast_config["exception"] = e from typing import Any, Optional @@ -333,6 +348,21 @@ def d(self, value): _diffs[self._indices1d()] = value rxd._setup_matrices() + def ast(self, region=None, prime=False): + if region is not None: + regions = region if hasattr(region, "__len__") else [region] + for r in regions: + if r in self._regions: + return self[r].ast(prime=prime) + elif r in self._extracellular_instances: + return self[r].ast(prime=prime) + if len(self._regions) == 1 and len(self._extracellular_instances) == 0: + return self[self._regions[0]].ast(prime=prime) + if len(self._regions) == 0 and len(self._extracellular_instances) == 1: + for ecs in self._extracellular_instances: + return self[ecs].ast(prime=prime) + raise RxDException(f"Invalid Region: {region}") + class SpeciesOnExtracellular(_SpeciesMathable): def __init__(self, species, extracellular): @@ -481,6 +511,69 @@ def d(self, value: float) -> None: def defined_on_region(self, r: Any) -> bool: return r == self._extracellular() + def ast(self, region=None, prime=False): + """ + Generate an AST node representing the species variable with region-specific information, + optionally as a derivative (primed) variable. + The name will be `Specie Name@Extracellular[Grid ID]` + + This method constructs an AST node using classes from `nmodl.ast` that encapsulates the + species’ identity along with its associated region. The resulting node is either a `VarName` + (for a normal species variable) or a `PrimeName` (for its time derivative) depending on + the `prime` flag. + + + Parameters: + region (optional): A region argument for compatibility (not used -- the region will + always be the one the species is defined on + prime (bool): If True, returns a `PrimeName` node (representing a derivative); + otherwise returns a `VarName` node. + + Returns: + An AST node (from `nmodl.ast`): + - A `PrimeName` node if `prime` is True. + - A `VarName` node otherwise. + + """ + if not _ast_config["nmodl_support"]: + if "exception" in _ast_config: + raise _ast_config["exception"] + else: + raise RxDException( + 'NMODL AST are disabled set rxd._ast_config["nmodl_support"] to True' + ) + if not initializer.is_initialized(): + initializer._do_init() + + rint = Integer( + self._id, Name(String(str(self._extracellular()._region._short_repr()))) + ) + name = ( + self._species().name + if self._species().name is not None + else f"{self._species().__class__.__name__}_{self._species()._id}" + ) + + # strip charactered used to give the region and region id + name = name.replace("@", "").replace("[", "").replace("]", "") + + # make a SymPy-compatible name -- with the index so it can also be used in Prime + region_name = str(self._extracellular()._region._short_repr()) + flat_name = f"{name}_{region_name}_{rint.eval()}" + + # Unlike VarName , PrimeName does not support an index -- so the index has been added to the name + if prime: + return VarName( + PrimeName(String(flat_name), Integer(1, None)), + None, + None, + ) + return VarName( + Name(String(flat_name)), + None, + None, + ) + class SpeciesOnRegion(_SpeciesMathable): def __init__(self, species: Any, region: Any) -> None: @@ -626,6 +719,71 @@ def instance3d(self): def _id(self): return self._species()._id + def ast(self, region=None, prime=False): + """ + Generate an AST node representing the species variable with region-specific information, + optionally as a derivative (primed) variable. + The name will be `Specie Name@Region Name[Region Id]` + + This method constructs an AST node using classes from `nmodl.ast` that encapsulates the + species’ identity along with its associated region. The resulting node is either a `VarName` + (for a normal species variable) or a `PrimeName` (for its time derivative) depending on + the `prime` flag. + + + Parameters: + region (optional): A region argument for compatibility (not used -- the region will + always be the one the species is defined on + prime (bool): If True, returns a `PrimeName` node (representing a derivative); + otherwise returns a `VarName` node. + + Returns: + An AST node (from `nmodl.ast`): + - A `PrimeName` node if `prime` is True. + - A `VarName` node otherwise. + + """ + if not _ast_config["nmodl_support"]: + if "exception" in _ast_config: + raise _ast_config["exception"] + else: + raise RxDException( + 'NMODL AST are disabled set rxd._ast_config["nmodl_support"] to True' + ) + if not initializer.is_initialized(): + initializer._do_init() + + rint = Integer(int(self._region()._id), Name(String(str(self._region().name)))) + name = ( + self._species().name + if self._species().name is not None + else f"{self._species().__class__.__name__}_{self._species()._id}" + ) + + # strip characters used to give the region and region id + name = name.replace("@", "").replace("[", "").replace("]", "") + + # Build a SymPy-compatible name + region_name = str(self._region().name) + flat_name = f"{name}_{region_name}_{rint.eval()}" + + # avoid using 'v' -- used for rxdmath.v: + # if name == 'v': + # name = f"{self._species().__class__.__name__}_v" + + # PrimeName take a String argument, so unlike VarName they cannot be indexed -- index added to the name. + if prime: + return VarName( + PrimeName(String(flat_name), Integer(1, None)), + None, + None, + ) + return VarName( + Name(String(flat_name)), + None, + None, + ) + def _xyz(seg): """Return the (x, y, z) coordinate of the center of the segment.""" diff --git a/src/nmodl/language/templates/pybind/pyast.cpp b/src/nmodl/language/templates/pybind/pyast.cpp index 8149dd6a97..e753e10eac 100644 --- a/src/nmodl/language/templates/pybind/pyast.cpp +++ b/src/nmodl/language/templates/pybind/pyast.cpp @@ -58,6 +58,11 @@ void init_ast_module(py::module& m) { .value("BOP_EXACT_EQUAL", BinaryOp::BOP_EXACT_EQUAL) .export_values(); + py::enum_(m_ast, "UnaryOp") + .value("UOP_NEGATION", UnaryOp::UOP_NEGATION) + .value("UOP_NOT", UnaryOp::UOP_NOT) + .export_values(); + py::enum_(m_ast, "AstNodeType", docstring::ast_nodetype_enum()) // clang-format off {% for node in nodes %} diff --git a/src/nrnpython/rxd.cpp b/src/nrnpython/rxd.cpp index 8a1b4772e0..b3c38aea71 100644 --- a/src/nrnpython/rxd.cpp +++ b/src/nrnpython/rxd.cpp @@ -673,6 +673,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]); @@ -729,7 +730,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++]; } } } diff --git a/test/rxd/conftest.py b/test/rxd/conftest.py index 2a5230c647..535e7e70a3 100644 --- a/test/rxd/conftest.py +++ b/test/rxd/conftest.py @@ -86,6 +86,7 @@ def neuron_nosave_instance(neuron_import): rxd.species._has_3d = False rxd.rxd._zero_volume_indices = numpy.ndarray(0, dtype=ctypes.c_long) rxd.set_solve_type(dimension=1) + rxd._ast_config = {"nmodl_support": True, "kinetic_block": "off"} @pytest.fixture diff --git a/test/rxd/test_ast.py b/test/rxd/test_ast.py new file mode 100644 index 0000000000..154b000b05 --- /dev/null +++ b/test/rxd/test_ast.py @@ -0,0 +1,379 @@ +import pytest +import json + +try: + from neuron.nmodl import to_json + + skip = False +except: + skip = True +from testutils import compare_data, tol + + +@pytest.fixture +def setup_section(neuron_instance): + """Setup a NEURON section for reactions""" + h, rxd, data, save_path = neuron_instance + sec = h.Section(name="soma") + yield (neuron_instance, sec) + + +@pytest.mark.skipif(skip, reason="nmodl not installed") +def test_species_ast(setup_section): + """Test that Species correctly generates an AST""" + (h, rxd, data, save_path), sec = setup_section + cyt = rxd.Region([sec], name="cyt", nrn_region="i") + ca = rxd.Species(cyt, name="ca", charge=2) + # Convert to AST and then to JSON + node = json.loads(to_json(ca.ast(cyt))) + ast = {"VarName": [{"Name": [{"String": [{"name": "ca_cyt_0"}]}]}]} + assert node == ast + + +@pytest.mark.skipif(skip, reason="nmodl not installed") +def test_rate_ast(setup_section): + """Test that Rate equations correctly generate an AST""" + (h, rxd, data, save_path), sec = setup_section + + cyt = rxd.Region([sec], name="cyt", nrn_region="i") + ca = rxd.Species(cyt, name="ca", charge=2) + rate = rxd.Rate(ca, -0.1 * ca) + + # create AST and convert to JSON + nodes, species = rate.ast() + node = json.loads(to_json(nodes[0])) + ast = { + "ExpressionStatement": [ + { + "DiffEqExpression": [ + { + "BinaryExpression": [ + { + "VarName": [ + { + "PrimeName": [ + {"String": [{"name": "ca_cyt_0"}]}, + {"Integer": [{"name": "1"}]}, + ] + } + ] + }, + {"BinaryOperator": [{"name": "="}]}, + { + "BinaryExpression": [ + { + "VarName": [ + { + "Name": [ + {"String": [{"name": "ca_cyt_0"}]} + ] + } + ] + }, + {"BinaryOperator": [{"name": "*"}]}, + {"Double": [{"name": "-0.1"}]}, + ] + }, + ] + } + ] + } + ] + } + assert ast == node + assert ca[cyt].ast().get_node_name() in species + + +@pytest.mark.skipif(skip, reason="nmodl not installed") +def test_reaction_ast(setup_section): + """Test that Reaction correctly generates an AST""" + (h, rxd, data, save_path), sec = setup_section + cyt = rxd.Region([sec], name="cyt", nrn_region="i") + ca = rxd.Species(cyt, name="ca", charge=2) + buf = rxd.Parameter(cyt, name="buf") + cabuf = rxd.Species(cyt, name="cabuf") + + reaction = rxd.Reaction(ca + buf, cabuf, 0.1, 0.05) + + rxd.rxdmath._ast_config["kinetic_block"] = "mass_action" + h.finitialize(-70) + react, species = reaction.ast() + node = json.loads(to_json(react[0])) + ast = { + "ReactionStatement": [ + { + "BinaryExpression": [ + { + "ReactVarName": [ + {"Integer": [{"Name": [{"String": [{"name": "1"}]}]}]}, + { + "VarName": [ + {"Name": [{"String": [{"name": "ca_cyt_0"}]}]} + ] + }, + ] + }, + {"BinaryOperator": [{"name": "+"}]}, + { + "ReactVarName": [ + {"Integer": [{"Name": [{"String": [{"name": "1"}]}]}]}, + { + "VarName": [ + {"Name": [{"String": [{"name": "buf_cyt_0"}]}]} + ] + }, + ] + }, + ] + }, + {"ReactionOperator": [{"name": "<->"}]}, + { + "ReactVarName": [ + {"Integer": [{"Name": [{"String": [{"name": "1"}]}]}]}, + {"VarName": [{"Name": [{"String": [{"name": "cabuf_cyt_0"}]}]}]}, + ] + }, + {"Double": [{"name": "0.1"}]}, + {"Double": [{"name": "0.05"}]}, + ] + } + + assert node == ast + + myspecies = [ca[cyt].ast(), cabuf.ast(cyt)] + for sp in myspecies: + assert sp.get_node_name() in species + + # should not include parameters + assert buf[cyt].ast().get_node_name() not in species + + rxd.rxdmath._ast_config["kinetic_block"] = "on" + react, species = reaction.ast() + node = json.loads(to_json(react[0])) + assert node == ast + + rxd.rxdmath._ast_config[ + "kinetic_block" + ] = "off" # default with correct volume scaling + + rates, species = reaction.ast() + node = json.loads(to_json(rates[0])) + ast = { + "ExpressionStatement": [ + { + "DiffEqExpression": [ + { + "BinaryExpression": [ + { + "VarName": [ + { + "PrimeName": [ + {"String": [{"name": "ca_cyt_0"}]}, + {"Integer": [{"name": "1"}]}, + ] + } + ] + }, + {"BinaryOperator": [{"name": "="}]}, + { + "FunctionCall": [ + {"Name": [{"String": [{"name": "-"}]}]}, + { + "BinaryExpression": [ + { + "BinaryExpression": [ + { + "BinaryExpression": [ + { + "Double": [ + {"name": "0.1"} + ] + }, + { + "BinaryOperator": [ + {"name": "*"} + ] + }, + { + "VarName": [ + { + "Name": [ + { + "String": [ + { + "name": "ca_cyt_0" + } + ] + } + ] + } + ] + }, + ] + }, + {"BinaryOperator": [{"name": "*"}]}, + { + "VarName": [ + { + "Name": [ + { + "String": [ + { + "name": "buf_cyt_0" + } + ] + } + ] + } + ] + }, + ] + }, + {"BinaryOperator": [{"name": "-"}]}, + { + "BinaryExpression": [ + {"Double": [{"name": "0.05"}]}, + {"BinaryOperator": [{"name": "*"}]}, + { + "VarName": [ + { + "Name": [ + { + "String": [ + { + "name": "cabuf_cyt_0" + } + ] + } + ] + } + ] + }, + ] + }, + ] + }, + ] + }, + ] + } + ] + } + ] + } + + assert node == ast + + +@pytest.mark.skipif(skip, reason="nmodl not installed") +def test_multicompartment_reaction_ast(setup_section): + """Test that MultiCompartmentReaction generates an AST""" + (h, rxd, data, save_path), sec = setup_section + cyt = rxd.Region([sec], name="cyt", nrn_region="i") + mem = rxd.Region([sec], name="mem", geometry=rxd.ScalableBorder(1)) + + er = rxd.Region([sec], name="er") + ca = rxd.Species([cyt, er], name="ca", charge=2) + + reaction = rxd.MultiCompartmentReaction(ca[cyt], ca[er], 0.1, 0.05, membrane=mem) + + rxd.rxdmath._ast_config["kinetic_block"] = "mass_action" + react, species = reaction.ast() + node = json.loads(to_json(react)) + + ast = { + "ReactionStatement": [ + { + "ReactVarName": [ + {"Integer": [{"Name": [{"String": [{"name": "1"}]}]}]}, + {"VarName": [{"Name": [{"String": [{"name": "ca_cyt_0"}]}]}]}, + ] + }, + {"ReactionOperator": [{"name": "<->"}]}, + { + "ReactVarName": [ + {"Integer": [{"Name": [{"String": [{"name": "1"}]}]}]}, + {"VarName": [{"Name": [{"String": [{"name": "ca_er_2"}]}]}]}, + ] + }, + {"Double": [{"name": "0.1"}]}, + {"Double": [{"name": "0.05"}]}, + ] + } + + assert node == ast + for sp in [ca[cyt], ca[er]]: + assert sp.ast().get_node_name() in species + + rxd.rxdmath._ast_config["kinetic_block"] = "off" # defualt + rates, species = reaction.ast() + node = json.loads(to_json(rates[0])) + ast = { + "ExpressionStatement": [ + { + "DiffEqExpression": [ + { + "BinaryExpression": [ + { + "VarName": [ + { + "PrimeName": [ + {"String": [{"name": "ca_cyt_0"}]}, + {"Integer": [{"name": "1"}]}, + ] + } + ] + }, + {"BinaryOperator": [{"name": "="}]}, + { + "BinaryExpression": [ + { + "BinaryExpression": [ + { + "VarName": [ + { + "Name": [ + { + "String": [ + {"name": "ca_cyt_0"} + ] + } + ] + } + ] + }, + {"BinaryOperator": [{"name": "*"}]}, + {"Double": [{"name": "0.1"}]}, + ] + }, + {"BinaryOperator": [{"name": "-"}]}, + { + "BinaryExpression": [ + { + "VarName": [ + { + "Name": [ + { + "String": [ + {"name": "ca_er_2"} + ] + } + ] + } + ] + }, + {"BinaryOperator": [{"name": "*"}]}, + {"Double": [{"name": "0.05"}]}, + ] + }, + ] + }, + ] + } + ] + } + ] + } + assert node == ast + for sp in [ca[cyt], ca[er]]: + assert sp.ast().get_node_name() in species