feat(gradient): route free params that reach the model through a condition (#511)

pybnf/algorithms/optimizers/gradient_base.py

@@ -47,7 +47,7 @@
 from .concurrent_multistart import DONE, ConcurrentMultiStartOptimizer
 from ...gradient import (
     GradientNotSupported,
-    apply_routing,
+    apply_routings,
     assemble_constraint_gradient,
     assemble_gaussian_gradient,
     route_for_model,
@@ -379,31 +379,41 @@ def _setup_gradient_path(self):
         routings -- idempotent, called once from the leaf's ``start_run`` (before the
         model scatter, so the request rides the pickle to the workers).
 
-        For each model: ``apply_routing`` the **wildtype** routing, whose request
-        lists (``sensitivity_params`` / ``sensitivity_ic``) are the union over
-        conditions -- the wildtype pins nothing, so every parameter-/IC-bound free
-        parameter is requested, a superset of any single condition's (a ``=``-pinned
-        condition only ever *drops* columns). Then build one
-        :class:`ExperimentRouting` per scored ``(model, suffix)``, carrying that
-        condition's chain-rule factors for the assembly. Raises (capability gate) if
-        the bngsim build lacks ``output_sensitivities``."""
+        For each model: build one :class:`ExperimentRouting` per scored
+        ``(model, suffix)`` (carrying that condition's chain-rule factors for the
+        assembly), then ``apply_routings`` the **union** of their sensitivity
+        requests -- plus the wildtype's. The wildtype alone is not a superset: a
+        condition can route a free parameter to a column no other experiment binds (a
+        per-condition estimated initial condition, ADR-0076), so that column is
+        reached only through the condition and must be unioned in (an extra requested
+        column is harmless; a missing one aborts the assembly). Raises (capability
+        gate) if the bngsim build lacks ``output_sensitivities``."""
         if self._routings is not None:
             return
         names = [v.name for v in self.variables]
         routings = {}
         for model in self.model_list:
-            # Union sensitivity request (wildtype) -> sets _sensitivity_request,
-            # which survives the scatter and is applied at every simulate().
-            apply_routing(model, route_for_model(model, names, condition=None))
             # Declare which of this model's outputs are scored gradient targets so
             # an incidental/unscored action (a stochastic diagnostic, a
             # carried-state pre-equilibration scan) runs sensitivity-free instead
             # of aborting the whole fit at a differentiability guard (#475). Rides
             # the scatter alongside the sensitivity request.
             model.set_scored_suffixes(self.exp_data.get(model.name, {}))
+            # Apply the UNION sensitivity request over the wildtype and every scored
+            # condition -> sets _sensitivity_request, which survives the scatter and is
+            # applied at every simulate(). The wildtype alone is NOT a superset once a
+            # condition routes a free parameter to a column no other experiment binds --
+            # a per-condition estimated initial condition (ADR-0076): its species-IC /
+            # multiplier column is reached only through that condition, so the union must
+            # include every scored routing (an extra column is harmless; a missing one aborts).
+            wildtype = route_for_model(model, names, condition=None)
+            model_routings = {}
             for suffix in self.exp_data.get(model.name, {}):
                 condition = self._condition_for_suffix(model, suffix)
-                routings[(model.name, suffix)] = route_for_model(model, names, condition)
+                model_routings[suffix] = route_for_model(model, names, condition)
+            apply_routings(model, [wildtype, *model_routings.values()])
+            for suffix, routing in model_routings.items():
+                routings[(model.name, suffix)] = routing
         self._routings = routings
 
     def _condition_for_suffix(self, model, suffix):

pybnf/bngsim_model/expressions.py

@@ -172,6 +172,38 @@ def _parse_net_species_initializers(net_lines):
     return initializers
 
 
+def _net_species_ic_seed_map(species_initializers, param_names):
+    """Map a model parameter to the species whose initial value it *bares* in a .net species
+    block, for the gradient router's per-condition estimated initial condition composition
+    (ADR-0076, #511).
+
+    A **bare** initializer ``species <- <param>`` (the whole expression is a single parameter
+    id) has ``d(IC)/d(param) = 1``, so a free parameter a condition assigns to that parameter
+    routes straight onto the species' initial-condition sensitivity axis. Such a parameter maps
+    to its species. A parameter referenced by a **non-bare** initializer expression (``2*p``,
+    ``a+b``), or that bares more than one species, maps to ``None`` -- present but non-routable
+    (its ``d(IC)/d(param)`` is not a plain ``1``), so the router refuses rather than emitting a
+    wrong column.
+    """
+    param_set = set(param_names)
+    bare = {}      # param -> set(species) it bares
+    nonbare = set()  # params a non-bare initializer references (force a refuse)
+    ident = re.compile(r'^[A-Za-z_]\w*$')
+    token = re.compile(r'[A-Za-z_]\w*')
+    for species, expr in species_initializers:
+        e = expr.strip()
+        if ident.match(e) and e in param_set:
+            bare.setdefault(e, set()).add(species)
+        else:
+            nonbare.update(t for t in token.findall(e) if t in param_set)
+    seed_map = {}
+    for param, species in bare.items():
+        seed_map[param] = next(iter(species)) if len(species) == 1 else None
+    for param in nonbare:
+        seed_map[param] = None  # a non-bare use forces a refuse, even if also bare elsewhere
+    return seed_map
+
+
 def _parse_bngl_param_block(model_lines):
     """Extract BNGL parameter definitions as ordered (name, expression) pairs."""
     params = []

pybnf/bngsim_model/net_model.py

@@ -40,6 +40,7 @@
     _eval_model_expression,
     _build_mutant_param_set,
     _parse_net_species_initializers,
+    _net_species_ic_seed_map,
 )
 from .scan import (
     _with_sim_timeout,
@@ -343,20 +344,27 @@ def has_discrete_events(self):
     def sensitivity_entity_namespace(self):
         """The bind-by-id namespaces the gradient router classifies free parameters against (#448).
 
-        Returns ``(param_ids, species_initializers)``:
+        Returns ``(param_ids, species_initializers, ic_seed_map)``:
 
         * ``param_ids`` -- the model's ``begin parameters`` namespace (the engine's
           ``param_names``), the kinetic/global ids a free parameter binds to via ``set_param``
           and thus routes to ``Simulator(sensitivity_params=)``;
         * ``species_initializers`` -- the ``(species, initial-expr)`` pairs
           (``_parse_net_species_initializers``); a free parameter that is a species' bare
           initial-value expression binds to ``Simulator(sensitivity_ic=)`` keyed by the
-          species (an IC parameter is absent from the ODE RHS, so the parameter axis is zero).
+          species (an IC parameter is absent from the ODE RHS, so the parameter axis is zero);
+        * ``ic_seed_map`` -- ``{model parameter -> species}`` for a bare species initializer
+          (``_net_species_ic_seed_map``), so the router can route a free parameter a condition
+          assigns to that parameter onto the species' IC axis (a per-condition estimated initial
+          condition, ADR-0076, #511); a non-routable seed maps to ``None``.
 
         This is the only model coupling :mod:`pybnf.gradient.routing` needs, so the routing
-        core stays backend-agnostic. No simulation -- both namespaces are known at build time.
+        core stays backend-agnostic. No simulation -- all three are known at build time.
         """
-        return list(self._engine_model.param_names), list(self._net_species_initializers)
+        param_ids = list(self._engine_model.param_names)
+        species_initializers = list(self._net_species_initializers)
+        return (param_ids, species_initializers,
+                _net_species_ic_seed_map(species_initializers, param_ids))
 
     def _sensitivity_request_kwargs(self, method):
         """Simulator kwargs requesting forward sensitivities on the gradient path.

pybnf/bngsim_sbml_model.py

@@ -212,6 +212,7 @@ def _extract_sbml_structure(self):
         self._initial_dep_names = self._compute_initial_dependency_names(doc.getModel())
         self._species_unit_factor, self._unsafe_volume = self._compute_species_unit_factors(
             doc.getModel())
+        self._ic_seed_map = self._compute_ic_seed_map(doc.getModel())
 
     @staticmethod
     def _collect_ast_names(node, out):
@@ -289,6 +290,50 @@ def _compute_initial_dependency_names(self, sbml_model):
                     worklist.append(ref)
         return dep
 
+    def _compute_ic_seed_map(self, sbml_model):
+        """Map a model parameter to the species whose initial value it *bares*.
+
+        A **bare** ``initialAssignment`` ``species = <param>`` (a single ``ci`` naming a global
+        parameter, on a concentration-based / unit-factor-1 species) has ``d(IC)/d(param) = 1``,
+        so the gradient router can route a free parameter a condition assigns to that parameter
+        (a per-condition estimated initial condition, ADR-0076, #511) straight onto the species'
+        initial-condition sensitivity axis. Such a parameter maps to its species.
+
+        A parameter that seeds a species initial value through a **non-bare** expression
+        (``2*init``, ``a+b``), or on an amount species whose value needs a non-unit
+        concentration factor, or that bares more than one species, maps to ``None`` -- present
+        but non-routable: its ``d(IC)/d(param)`` is not a plain ``1``, so the router refuses
+        (``classify_condition_target``) rather than emitting a wrong/silently-zero column.
+        """
+        param_set = set(self._global_param_names)
+        bare = {}      # param -> set(species) it bares
+        nonbare = set()  # params a species IA references non-baruely (force a refuse)
+        for i in range(sbml_model.getNumInitialAssignments()):
+            ia = sbml_model.getInitialAssignment(i)
+            symbol = ia.getSymbol()
+            if symbol not in self._species_name_set:
+                continue
+            math = ia.getMath()
+            is_bare = (
+                math is not None
+                and math.getType() == libsbml.AST_NAME
+                and math.getNumChildren() == 0
+                and math.getName() in param_set
+                and self._species_unit_factor.get(symbol, 1.0) == 1.0
+            )
+            if is_bare:
+                bare.setdefault(math.getName(), set()).add(symbol)
+            else:
+                refs = set()
+                self._collect_ast_names(math, refs)
+                nonbare.update(r for r in refs if r in param_set)
+        seed_map = {}
+        for param, species in bare.items():
+            seed_map[param] = next(iter(species)) if len(species) == 1 else None
+        for param in nonbare:
+            seed_map[param] = None  # a non-bare use forces a refuse, even if also bare elsewhere
+        return seed_map
+
     def _compute_species_unit_factors(self, sbml_model):
         """Per-species factor converting a PyBNF species value to a bngsim
         concentration, plus a flag for volumes we cannot safely convert.
@@ -909,7 +954,7 @@ def has_discrete_events(self):
     def sensitivity_entity_namespace(self):
         """The bind-by-id namespaces the gradient router classifies free parameters against (#448).
 
-        Returns ``(param_ids, species_initializers)``:
+        Returns ``(param_ids, species_initializers, ic_seed_map)``:
 
         * ``param_ids`` -- the model's global ``parameter`` ids, the kinetic ids a free
           parameter binds to via ``set_param`` and thus routes to
@@ -919,14 +964,19 @@ def sensitivity_entity_namespace(self):
           for a species sets that species' initial value (via ``set_concentration``, the
           bind-by-id convention, ADR-0034), so each species' bare initializer expression *is*
           its own name; such a free parameter routes to the initial-condition axis keyed by the
-          species (an IC parameter is absent from the ODE RHS, so its parameter axis is zero).
+          species (an IC parameter is absent from the ODE RHS, so its parameter axis is zero);
+        * ``ic_seed_map`` -- ``{model parameter -> species}`` for a bare ``initialAssignment``
+          ``species = <param>`` (:meth:`_compute_ic_seed_map`), so the router can route a free
+          parameter a condition assigns to that parameter onto the species' IC axis (a
+          per-condition estimated initial condition, ADR-0076, #511); a non-routable seed maps
+          to ``None``.
 
         This is the only model coupling :mod:`pybnf.gradient.routing` needs, so the routing core
-        stays backend-agnostic. No simulation -- both namespaces are known at load time. (A
-        species whose initial is a non-trivial *expression* of a parameter is the deferred
-        non-bare-initializer case the router documents, exactly as for the net backend.)
+        stays backend-agnostic. No simulation -- all three are known at load time.
         """
-        return list(self._global_param_names), [(s, s) for s in self._species_names]
+        return (list(self._global_param_names),
+                [(s, s) for s in self._species_names],
+                dict(self._ic_seed_map))
 
     def _sensitivity_request_kwargs(self, method):
         """Simulator kwargs requesting forward sensitivities on the gradient path.

pybnf/gradient/__init__.py

@@ -28,13 +28,16 @@
     PARAM,
     IC,
     NONE,
+    RouteContribution,
     ParamRoute,
     ExperimentRouting,
     classify_free_param,
+    classify_condition_target,
     condition_factor,
     route_experiment,
     route_for_model,
     apply_routing,
+    apply_routings,
 )
 from .assembly import (
     GradientResult,
@@ -50,13 +53,16 @@
     'PARAM',
     'IC',
     'NONE',
+    'RouteContribution',
     'ParamRoute',
     'ExperimentRouting',
     'classify_free_param',
+    'classify_condition_target',
     'condition_factor',
     'route_experiment',
     'route_for_model',
     'apply_routing',
+    'apply_routings',
     'GradientResult',
     'assemble_gaussian_gradient',
     'assemble_gradient_and_fisher_hessian',

pybnf/gradient/assembly.py

@@ -469,24 +469,25 @@ def _accumulate_experiment_fisher(objective, sim_data, exp_data, routing, index,
         _accumulate_fisher_point(objective, sim_data, exp_data, index, point, hessian)
 
 
-def _sensitivity(sens, selector, route, sim_row):
+def _sensitivity(sens, selector, contribution, sim_row, free_param):
     """The native forward sensitivity ``d(observable)/d(routed entity)`` at one time row.
 
-    Reads the parameter axis (``sensitivity_params``) for a PARAM route and the
-    initial-condition axis (``sensitivity_ic``) for an IC route, addressing the
-    entity by the route's ``key`` (parameter id, or species for an IC)."""
-    if route.target == PARAM:
+    Reads the parameter axis (``sensitivity_params``) for a PARAM contribution and the
+    initial-condition axis (``sensitivity_ic``) for an IC contribution, addressing the entity by
+    the contribution's ``key`` (parameter id, or species for an IC). ``free_param`` names the
+    routed free parameter for the diagnostic when a requested column is absent."""
+    if contribution.target == PARAM:
         axis, labels = 'parameter', sens.param_names
     else:  # IC
         axis, labels = 'ic', sens.ic_species
-    if route.key not in labels:
+    if contribution.key not in labels:
         raise GradientNotSupported(
             "Free parameter '%s' routes to %s '%s', but the simulation's "
             "sensitivity tensor has no such column (axis labels: %s). Apply the same "
             "routing to the model before running it (apply_routing)."
-            % (route.free_param, axis, route.key, ', '.join(map(str, labels)) or '(none)'))
+            % (free_param, axis, contribution.key, ', '.join(map(str, labels)) or '(none)'))
     column = sens.slice_for(selector, axis=axis)   # (n_times, n_axis)
-    return column[sim_row, labels.index(route.key)]
+    return column[sim_row, labels.index(contribution.key)]
 
 
 def _raw_sensitivity_accessor(objective, sim_data, sens, routing, index, n_param, indvar):
@@ -522,9 +523,12 @@ def tensor_sens(col_name, row):
         vec = np.zeros(n_param)
         selector = _selector_for(sens, col_name)
         for name, route in routing.routes.items():
-            if route.factor == 0.0 or route.target == NONE:
-                continue
-            vec[index[name]] = route.factor * _sensitivity(sens, selector, route, row)
+            total = 0.0
+            for c in route.contributions:
+                if c.target == NONE or c.factor == 0.0:
+                    continue
+                total += c.factor * _sensitivity(sens, selector, c, row, name)
+            vec[index[name]] = total
         return vec
 
     def raw_sens(col_name, row):
@@ -743,8 +747,11 @@ def raw_sens(model, suffix, observable, row):
         selector = _selector_for(sens, observable)
         vec = np.zeros(n_param)
         for name, route in routing.routes.items():
-            if route.factor == 0.0 or route.target == NONE:
-                continue
-            vec[index[name]] = route.factor * _sensitivity(sens, selector, route, row)
+            total = 0.0
+            for c in route.contributions:
+                if c.target == NONE or c.factor == 0.0:
+                    continue
+                total += c.factor * _sensitivity(sens, selector, c, row, name)
+            vec[index[name]] = total
         return vec
     return raw_sens