From ecc5465f4690d4ec4ca8be30d1eb32cd0d54f598 Mon Sep 17 00:00:00 2001 From: Bill Hlavacek Date: Sat, 25 Jul 2026 16:44:38 -0600 Subject: [PATCH] fix: gradient-fit a piecewise model with a count likelihood Two independent defects made `fit_type = lbfgs` (or any gradient fit) score `inf` or `nan` everywhere on a model built from `if()`-gated rate laws with a negative-binomial count likelihood -- the shape of every COVID-19 compartmental job in the corpus (Lin 2021, Mallela 2024). 1. Over-requested output-sensitivity selectors. `_extract_output_sensitivities` asked bngsim for an `expression:` sensitivity for EVERY global function whenever `print_functions` was on. bngsim differentiates function bodies symbolically and refuses, per function, any body carrying a non-differentiable construct -- `if()`, a comparison, min/max/abs/floor, a table function (lanl/bngsim#198) -- raising `ValueError` if such a selector is requested. `has_sensitivities_expressions` only says "an expression block exists", so it could not stand in for "this expression is differentiable". A piecewise model has every function refused, so every simulation died and the objective was `inf` at every point. Filter the request to the differentiable expressions (`_differentiable_expression_names`); a scored column is essentially never one of these, and if it were, its absence now surfaces as a clean gradient error rather than a dead simulation. No support map (older bngsim, a Result loaded from disk) => unchanged behavior. 2. Divide-by-zero in the negative-binomial derivatives. `d(data_fit)/d(mean) = r (mean - obs) / (mean (r + mean))` is `0/0 -> nan` (observed zero) or `-inf` at a MEAN-centered prediction of exactly 0. That 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 t0 -- and `data_fit` scores it finitely by clipping `prob = r/(r+mean)` at `1 - 1e-10`. The `nan` propagated out to the optimizer, which then tried to assign a free parameter `nan`. Mirror the value path's clip on the mean side (`_mean_for_slope`) so every expression that divides by the mean gets the finite slope belonging to the objective actually being scored. MEDIAN centering never reaches the floor. Verified end to end on the Mallela et al. 2024 NYC job (42 species, 88 reactions, 14 `if()` functions, 649 daily counts, neg_bin with a fixed dispersion): before, every start returned `inf`; after, L-BFGS-B converges to nll 5235.80, improving on the published MAP's 5285.16 in the same box. --- pybnf/bngsim_model/net_model.py | 47 ++++++++++++++--- pybnf/noise/negative_binomial.py | 62 ++++++++++++++++++++--- tests/test_bngsim_output_sensitivities.py | 58 +++++++++++++++++++++ tests/test_gradient_assembly.py | 33 ++++++++++++ 4 files changed, 186 insertions(+), 14 deletions(-) diff --git a/pybnf/bngsim_model/net_model.py b/pybnf/bngsim_model/net_model.py index 6f8d7db6..c0dee25c 100644 --- a/pybnf/bngsim_model/net_model.py +++ b/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:`` for every - observable, plus ``expression:`` 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:`` 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: diff --git a/pybnf/noise/negative_binomial.py b/pybnf/noise/negative_binomial.py index 0311bcf3..4d9273c3 100644 --- a/pybnf/noise/negative_binomial.py +++ b/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( diff --git a/tests/test_bngsim_output_sensitivities.py b/tests/test_bngsim_output_sensitivities.py index 75b1f63b..ec01133a 100644 --- a/tests/test_bngsim_output_sensitivities.py +++ b/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'] diff --git a/tests/test_gradient_assembly.py b/tests/test_gradient_assembly.py index 2d1b4bd9..204774f3 100644 --- a/tests/test_gradient_assembly.py +++ b/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', [