fix: gradient-fit a piecewise model with a count likelihood

pybnf/bngsim_model/net_model.py

@@ -1914,18 +1914,50 @@ def _result_to_data(self, result, print_functions=False):
         return data
 
     @staticmethod
-    def _extract_output_sensitivities(result, print_functions):
+    def _differentiable_expression_names(result):
+        """The expressions on ``result`` that bngsim can hand back an output sensitivity for.
+
+        bngsim differentiates a global function symbolically (lanl/bngsim#198) and
+        **refuses**, per function, any body carrying a non-differentiable construct --
+        an ``if()`` conditional, a comparison, ``min``/``max``/``abs``/``floor``, a table
+        function. It records the per-function verdict on the Result, and
+        ``output_sensitivities`` raises ``ValueError`` if a refused function is among the
+        requested selectors. ``has_sensitivities_expressions`` is only "some expression
+        block exists", not "every expression is differentiable", so it cannot stand in.
+
+        Asking for a refused selector fails the whole simulation, and a *scored* column
+        is almost never one of these functions -- so filter them out. This matters for
+        any piecewise model: an epidemic model whose rates switch at ``if(t >= tau)``
+        has every one of its functions refused, and before this filter every simulation
+        of such a model died with ``ValueError`` on the gradient path, making the fit's
+        objective ``inf`` everywhere. Should a scored column genuinely be a refused
+        expression, its absence surfaces later as a clean "no sensitivity column"
+        gradient error rather than a dead simulation.
+
+        Empty support map (an older bngsim, or a Result loaded from disk) ⇒ no verdicts
+        to filter on ⇒ every expression is kept, exactly as before.
+        """
+        names = list(getattr(result, 'expression_names', None) or [])
+        support = getattr(result, '_expression_sens_support', None) or {}
+        if not support:
+            return names
+        return [name for name in names if support.get(name) is None]
+
+    @classmethod
+    def _extract_output_sensitivities(cls, result, print_functions):
         """Read the native-space ∂g/∂θ tensor off a sensitivity-bearing Result.
 
         Selectors mirror the Data columns -- ``observable:<name>`` for every
-        observable, plus ``expression:<name>`` for each expression when
-        ``print_functions`` is on and the backend computed expression
-        sensitivities. The ``parameter`` axis is read whenever sensitivity params
-        were requested; the ``ic`` axis whenever IC species were.
+        observable, plus ``expression:<name>`` for each *differentiable* expression
+        (:meth:`_differentiable_expression_names`) when ``print_functions`` is on and
+        the backend computed expression sensitivities. The ``parameter`` axis is read
+        whenever sensitivity params were requested; the ``ic`` axis whenever IC species
+        were.
         """
         selectors = ['observable:%s' % name for name in result.observable_names]
         if print_functions and getattr(result, 'has_sensitivities_expressions', False):
-            selectors += ['expression:%s' % name for name in result.expression_names]
+            selectors += ['expression:%s' % name
+                          for name in cls._differentiable_expression_names(result)]
         param_names = list(result.sensitivity_params)
         ic_species = list(result.sensitivity_ic_species)
         d_param = None
@@ -1969,7 +2001,8 @@ def _extract_ss_output_sensitivities(self, ss_result, print_functions):
         """
         selectors = ['observable:%s' % name for name in ss_result.observable_names]
         if print_functions and getattr(ss_result, 'expression_names', None):
-            selectors += ['expression:%s' % name for name in ss_result.expression_names]
+            selectors += ['expression:%s' % name
+                          for name in self._differentiable_expression_names(ss_result)]
         param_names = list(ss_result.sensitivity_params)
         d_param = None
         if param_names:

pybnf/noise/negative_binomial.py

@@ -7,6 +7,16 @@
 from .base import NoiseModel
 from .location import MEDIAN
 
+#: The clip :meth:`NegBinomial.data_fit` puts on ``prob = r/(r+mean)`` to keep a
+#: degenerate point's ``-logpmf`` finite.
+_PROB_CLIP = 1e-10
+
+#: The mean-side image of that clip: ``prob <= 1 - _PROB_CLIP`` iff
+#: ``mean >= r * _PROB_CLIP/(1 - _PROB_CLIP)``. Every derivative that divides by the mean
+#: floors it here (:meth:`NegBinomial._mean_for_slope`) so a zero prediction -- which the
+#: value path scores finitely -- yields a finite slope rather than ``nan``/``-inf``.
+_MEAN_FLOOR_REL = _PROB_CLIP / (1 - _PROB_CLIP)
+
 
 def _d_betainc_d_b(a, b, x, rel=1e-6):
     """``d/db`` of the regularized incomplete beta ``betainc(a, b, x) = I_x(a, b)`` w.r.t. its
@@ -108,11 +118,34 @@ def _mean(self, prediction, noise):
             return _mean_for_median(prediction, noise)
         return prediction
 
+    def _mean_for_slope(self, prediction, noise):
+        """:meth:`_mean`, floored away from zero for the expressions that divide by it.
+
+        A MEAN-centered prediction of **exactly** zero is not pathological -- it is what
+        any model whose output is gated off over part of the fit window predicts there (an
+        epidemic model before its start time ``t0``, a stimulus-driven readout before the
+        stimulus). :meth:`data_fit` scores such a point finitely because it clips
+        ``prob = r/(r+mean)`` at ``1 - _PROB_CLIP``; every *derivative* below instead
+        divides by the mean (``r (mean - obs) / (mean (r + mean))``, ``r / (mean (r+mean))``),
+        which at ``mean == 0`` is ``0/0 -> nan`` for an observed zero and ``-inf`` otherwise.
+        That ``nan`` propagates into the gradient and out to the optimizer, which then tries
+        to assign a free parameter ``nan``.
+
+        So mirror the value path's clip on the mean side: ``prob <= 1 - _PROB_CLIP`` is
+        exactly ``mean >= r _PROB_CLIP / (1 - _PROB_CLIP)``, i.e. this floor. The slope
+        is then the finite value belonging to the objective actually being scored -- large
+        (order ``obs / (r _PROB_CLIP)``) where the model predicts zero against a positive
+        count, which is correct: the objective there really is nearly flat-then-cliff, and a
+        large finite push toward a positive prediction is the right search direction.
+        MEDIAN centering never reaches the floor (``_mean_for_median`` returns a strictly
+        positive mean for any prediction), so it is a no-op there."""
+        return max(self._mean(prediction, noise), noise * _MEAN_FLOOR_REL)
+
     def data_fit(self, prediction, observation, noise, extra=None):
         if observation < 0:
             return 0
         mean = self._mean(prediction, noise)
-        prob = np.clip(noise / (noise + mean), 1e-10, 1 - 1e-10)
+        prob = np.clip(noise / (noise + mean), _PROB_CLIP, 1 - _PROB_CLIP)
         assert isinstance(noise, float)
         # log of the negative-binomial PMF P(observation | r=noise, prob)
         # == scipy.stats.nbinom.logpmf(observation, noise, prob).
@@ -145,12 +178,14 @@ def d_data_fit_d_prediction(self, prediction, observation, noise, extra=None):
         A negative observation contributes nothing (the count-domain guard, mirroring
         :meth:`data_fit`). A prediction clamped to the count floor (``pred <= 0``, where ``target``
         floors at 0 and the median stops moving) has slope 0 -- a kink at ``pred == 0`` where PyBNF
-        takes the floor-side subgradient, like the Laplace kink (#454). The gradient uses the
-        un-clipped analytic form (``data_fit`` clips ``prob`` only at pathological extremes)."""
+        takes the floor-side subgradient, like the Laplace kink (#454). The mean-slope divides by
+        the mean, so it reads it through :meth:`_mean_for_slope` -- the value path's ``prob`` clip
+        expressed on the mean -- which keeps a MEAN-centered zero prediction finite instead of
+        ``nan``; away from that floor the form is the un-clipped analytic one."""
         if observation < 0:
             return 0.0
         r = noise
-        mean = self._mean(prediction, r)
+        mean = self._mean_for_slope(prediction, r)
         d_fit_d_mean = r * (mean - observation) / (mean * (r + mean))
         if self.location is not MEDIAN:
             return d_fit_d_mean   # MEAN: d mean/d pred = 1
@@ -191,11 +226,15 @@ def d_nll_d_noise_params(self, prediction, observation, noise, extra=None):
           at a clamped prediction. (Validated FD-first: the partial-only form is off 5-15%, sometimes
           far more, on the median.)
 
-        A negative observation contributes nothing (the count-domain guard)."""
+        A negative observation contributes nothing (the count-domain guard). The mean is read
+        through :meth:`_mean_for_slope` for consistency with :meth:`d_data_fit_d_prediction` --
+        it matters only for the MEDIAN coupling term, which divides by the mean; the MEAN
+        partial above is finite at a zero prediction either way, and the floor shifts it by
+        ``O(_PROB_CLIP)``. MEDIAN centering never reaches the floor."""
         if observation < 0:
             return {'dispersion': 0.0}
         r = noise
-        mean = self._mean(prediction, r)
+        mean = self._mean_for_slope(prediction, r)
         prob = r / (r + mean)
         d_fit_d_r = (digamma(r) - digamma(observation + r) - np.log(prob) - 1.0
                      + (r + observation) / (r + mean))
@@ -222,7 +261,16 @@ def location_fisher(self, prediction, observation, noise, extra=None):
           ``d mean/d pred`` is the non-elementary implicit derivative; its location Fisher is out of
           scope for this cut -- refused, pointing at ``fit_type = lbfgs`` (which fits it via the
           scalar data-fit gradient). A negative observation contributes no curvature (the
-          count-domain guard, mirroring :meth:`data_fit`)."""
+          count-domain guard, mirroring :meth:`data_fit`).
+
+        A **zero mean** deliberately reports **no curvature** rather than the
+        :meth:`_mean_for_slope` floor the *gradient* uses. The two want opposite things: the
+        true ``I_mean`` diverges as ``1/mean``, so flooring it would put a ``~1/(r
+        _PROB_CLIP)`` weight on every gated-off row (an epidemic model's whole pre-``t0``
+        stretch) and let those rows dominate the Fisher matrix, crushing the trust-region
+        step. Zero says "this row carries no information about the mean", which is the stable
+        reading and leaves the gradient -- which does floor -- to supply the descent
+        direction."""
         if self.location is MEDIAN:
             from ..gradient.errors import GradientNotSupported
             raise GradientNotSupported(

tests/test_bngsim_output_sensitivities.py

@@ -458,3 +458,61 @@ def test_scored_continuation_dose_response_refuses_on_gradient_path():
     with pytest.raises(PybnfError) as exc:
         model.execute('/tmp', 'dr_cont', 120)
     assert 'reset_conc' in str(exc.value).lower() or 'seed' in str(exc.value).lower()
+
+
+# --------------------------------------- non-differentiable expressions ----
+
+class _StubResult:
+    """The three attributes the selector builder reads off a bngsim ``Result``."""
+
+    has_sensitivities_expressions = True
+    observable_names = ['Atot']
+    sensitivity_params = ['k']
+    sensitivity_ic_species = []
+
+    def __init__(self, expression_names, support):
+        self.expression_names = expression_names
+        self._expression_sens_support = support
+
+    def output_sensitivities(self, selectors, axis='parameter'):
+        for sel in selectors:
+            name = sel.split(':', 1)[1]
+            reason = self._expression_sens_support.get(name)
+            if sel.startswith('expression:') and reason is not None:
+                raise ValueError(
+                    "output_sensitivities: expression '%s' has no output sensitivity -- %s"
+                    % (name, reason))
+        return np.zeros((3, len(selectors), 1))
+
+
+IF_REASON = 'uses unsupported construct: if() conditional'
+
+
+def test_non_differentiable_expressions_are_left_out_of_the_selector_request():
+    """bngsim refuses an output sensitivity for any function whose body carries an ``if()``
+    (or a comparison / min / max / floor / table function), and raises if such a selector is
+    requested -- which failed the whole simulation. Every function of a piecewise model (an
+    epidemic model switching rates at ``if(t >= tau)``) is refused, so before this filter the
+    gradient path scored ``inf`` everywhere on such a model. Only the differentiable ones are
+    requested; the observable the fit actually scores is unaffected."""
+    result = _StubResult(['smooth_f', 'switch_f'],
+                         {'smooth_f': None, 'switch_f': IF_REASON})
+    sens = bngsim_model.BngsimModel._extract_output_sensitivities(result, True)
+    assert sens.selectors == ['observable:Atot', 'expression:smooth_f']
+
+
+def test_every_expression_refused_leaves_only_the_observables():
+    """The Mallela/Lin COVID case: all 14 functions are ``if()`` chains, so the request
+    degenerates to the observables -- which is all a fit scoring a Molecules observable needs."""
+    names = ['v_rate', 'Ytheta', 'Lambdatau', 'Ptau']
+    result = _StubResult(names, {n: IF_REASON for n in names})
+    sens = bngsim_model.BngsimModel._extract_output_sensitivities(result, True)
+    assert sens.selectors == ['observable:Atot']
+
+
+def test_missing_support_map_keeps_every_expression():
+    """An older bngsim (or a Result read back from disk) records no per-function verdict;
+    with nothing to filter on the request is unchanged from before the filter existed."""
+    result = _StubResult(['f1', 'f2'], {})
+    sens = bngsim_model.BngsimModel._extract_output_sensitivities(result, True)
+    assert sens.selectors == ['observable:Atot', 'expression:f1', 'expression:f2']

tests/test_gradient_assembly.py

@@ -1295,6 +1295,39 @@ def test_family_prediction_derivative_matches_finite_difference(family, extra):
     np.testing.assert_allclose(ana, num, rtol=1e-5)
 
 
+@pytest.mark.parametrize('obs', [0.0, 1.0, 4212.0], ids=['obs_zero', 'obs_one', 'obs_big'])
+def test_neg_bin_mean_gradient_finite_at_a_zero_prediction(obs):
+    """A MEAN-centered prediction of **exactly** zero has a finite data fit and must have a
+    finite slope. The closed form ``r (mean - obs)/(mean (r + mean))`` is ``0/0`` there (``nan``
+    for an observed zero, ``-inf`` otherwise), and a ``nan`` gradient reaches the optimizer as a
+    ``nan`` parameter proposal. It is not a pathological point: any model gated off over part of
+    the fit window predicts exactly zero there (an epidemic model before its start time), and
+    ``data_fit`` scores it finitely via its ``prob`` clip -- so the slope mirrors that clip."""
+    family, r = NegBinomial(location=MEAN), 3.0
+    assert np.isfinite(family.data_fit(0.0, obs, r))
+    slope = family.d_data_fit_d_prediction(0.0, obs, r)
+    assert np.isfinite(slope)
+    # An observed positive count at a zero prediction pushes the prediction UP (slope < 0);
+    # an observed zero is already at its optimum from below (slope >= 0).
+    assert slope < 0 if obs > 0 else slope >= 0
+    assert np.isfinite(family.d_nll_d_noise_params(0.0, obs, r)['dispersion'])
+
+
+def test_neg_bin_mean_zero_prediction_slope_matches_the_scored_objective():
+    """The zero-prediction slope is not an arbitrary sentinel: it is the un-floored closed form
+    at the very mean the value path's ``prob`` clip already pretends to score, so value and
+    gradient stay consistent and neither jumps across the floor."""
+    family, r, obs = NegBinomial(location=MEAN), 3.0, 7.0
+    floor = r * 1e-10 / (1 - 1e-10)     # prob == 1 - 1e-10, data_fit's upper clip
+    np.testing.assert_allclose(family.data_fit(0.0, obs, r), family.data_fit(floor, obs, r),
+                               rtol=1e-12)
+    closed_form = r * (floor - obs) / (floor * (r + floor))
+    np.testing.assert_allclose(family.d_data_fit_d_prediction(0.0, obs, r), closed_form,
+                               rtol=1e-12)
+    np.testing.assert_allclose(family.d_data_fit_d_prediction(floor, obs, r), closed_form,
+                               rtol=1e-12)
+
+
 # ===================== exact square-root-loss least-squares residual (layer G follow-up, #459) ===
 
 @pytest.mark.parametrize('family, extra', [