diff --git a/documentation/source/api/python.rst b/documentation/source/api/python.rst index c13cfb2..ac92021 100644 --- a/documentation/source/api/python.rst +++ b/documentation/source/api/python.rst @@ -15,6 +15,7 @@ Python API Reference python/optimization.rst python/fitting.rst python/fitspec.rst + python/mcmc.rst Extras diff --git a/documentation/source/api/python/fitspec.rst b/documentation/source/api/python/fitspec.rst index 1f105b5..1368ea6 100644 --- a/documentation/source/api/python/fitspec.rst +++ b/documentation/source/api/python/fitspec.rst @@ -18,6 +18,7 @@ the result. hjcfit template -o my-fit.toml a specification to edit hjcfit check my-fit.toml say what it would do; run nothing hjcfit fit my-fit.toml -o out.json run it, and keep the result + hjcfit sample my-fit.toml -o post sample the posterior of the rates ``examples/fit_template.ipynb`` is the same thing as a notebook, over ``examples/CH82.toml``. @@ -185,6 +186,11 @@ The specification :members: from_dict, validate .. autoclass:: SearchSpec + +.. autoclass:: LikelihoodSpec + :members: solver + +.. autoclass:: MCMCSpec :members: from_dict, validate .. autoexception:: SpecError @@ -215,6 +221,14 @@ Carrying it out .. autofunction:: write_result +.. autofunction:: sample + +.. autoclass:: SampleOutcome + +.. autofunction:: samples_as_dict + +.. autofunction:: write_samples + .. autoexception:: RunnerError The command @@ -230,5 +244,5 @@ environment on Windows. Exit status is 0 on success, 1 on anything the command can explain -- one line beginning with ``hjcfit:`` -- and 2 from ``fit`` when the search did not -converge. The estimates are still printed in that case; it is the status that +converge, or from ``sample`` when the chains disagree (split-R-hat above 1.1). The estimates are still printed in that case; it is the status that says so, because that is what a script reads. diff --git a/documentation/source/api/python/fitting.rst b/documentation/source/api/python/fitting.rst index 5adbfd5..1962f1e 100644 --- a/documentation/source/api/python/fitting.rst +++ b/documentation/source/api/python/fitting.rst @@ -102,6 +102,12 @@ The fitter .. autodata:: FAILURE_COST +.. autodata:: SOLVER_OPTIONS + +``HJCFitter(..., solver=dict(nmax=2, xtol=1e-12))`` passes them to every +record's likelihood. Options left out keep the defaults of +:py:class:`~HJCFIT.likelihood.Log10Likelihood`. + The result """""""""" diff --git a/documentation/source/api/python/mcmc.rst b/documentation/source/api/python/mcmc.rst new file mode 100644 index 0000000..d8950e5 --- /dev/null +++ b/documentation/source/api/python/mcmc.rst @@ -0,0 +1,195 @@ +.. _python_mcmc_api: + +Bayesian inference +------------------ + +.. currentmodule:: HJCFIT.likelihood.mcmc + +:ref:`python_fitting_api` finds the rate constants that maximise the HJC +likelihood. This module treats them as random variables instead, and describes +their posterior distribution: the likelihood of the records multiplied by a +prior. The approach follows Epstein, Calderhead, Girolami & Sivilotti (2016) +*Biophys J* 111:333-348, who sampled that posterior by Markov chain Monte +Carlo using the same exact missed-events likelihood, and it uses their two +samplers. + +Like the fitter, it needs a mechanism from SCALCS (``pip install +hjcfit[fitting]``), and it is duck-typed on one, so it does not import scalcs +itself. + +The posterior +""""""""""""" + +.. code-block:: python + + import HJCFIT + from HJCFIT.likelihood.fitting import Record + from HJCFIT.likelihood.mcmc import LogPosterior + from scalcs.samples import samples + + bursts = HJCFIT.read_idealized_bursts("CH82", tau=1e-4, tcrit=4e-3) + record = Record(conc=100e-9, tres=1e-4, tcrit=4e-3, + groups=tuple(tuple(b) for b in bursts)) + + mechanism = samples.CH82() + posterior = LogPosterior(mechanism, [record]) + + posterior(mechanism.theta()) # ln prior + ln likelihood + +.. autoclass:: LogPosterior + :members: log_prior, log_likelihood, __call__, k, nevals, nfailures + +Two things differ from fitting on purpose: + +* **Nothing is clipped.** :py:class:`~HJCFIT.likelihood.fitting.HJCFitter` + resets a rate that leaves its limits, which suits a search. For a + posterior, a point outside the prior has zero density. Moving it to the + boundary instead would pile probability there. +* **A likelihood that cannot be computed gives** ``-inf``, **not a penalty.** + A sampler rejects the move. :py:attr:`LogPosterior.nfailures` counts these, + so a mechanism that the likelihood struggles with is visible. + +Priors +"""""" + +.. autoclass:: UniformPrior + :members: from_mechanism, logpdf, contains, sample, k + +.. autoclass:: LogUniformPrior + +The default prior is :py:meth:`UniformPrior.from_mechanism`: flat between the +limits each free rate already carries. Those are the limits a +maximum-likelihood fit resets against, so the mode of the posterior is the +maximum-likelihood estimate. They are also the prior of Epstein et al. +(2016): scalcs' default limits are U(1e-2, 1e6) s\ :sup:`-1`, and +U(1e-2, 1e10) M\ :sup:`-1` s\ :sup:`-1` for association rates. + +Solver settings matter +"""""""""""""""""""""" + +Pass ``solver=`` to reproduce a published value. On that paper's three AChR +records, its root-finding settings (``nmax=2``, tolerances 1e-12) give a +natural log-likelihood 0.38 higher than the defaults at the same rates. See +:py:data:`~HJCFIT.likelihood.fitting.SOLVER_OPTIONS`. + +Sampling +"""""""" + +Two stages, as in that paper. A Metropolis-within-Gibbs pilot walks the +logarithm of one rate at a time and finds the mode. An adaptive Metropolis +chain, started at that mode, learns the posterior covariance and samples +along correlations that the pilot crosses only slowly: + +.. code-block:: python + + from HJCFIT.likelihood.mcmc import ( + adaptive_sample, effective_sample_size, mwg_sample) + + pilot = mwg_sample(posterior, mechanism.theta(), n=10000, burnin=5000, + rng=1) + start, _ = pilot.mode() + chain = adaptive_sample(posterior, start, n=100000, burnin=50000, rng=2) + + draws = chain.kept() # (50000, k) rates, burn-in dropped + effective_sample_size(draws) # one value per rate + +Step sizes are tuned only during burn-in, so the samples that are kept come +from a chain with a fixed kernel. The adaptive covariance keeps learning +afterwards, but with a weight that falls as :math:`1/n`, which is what keeps +the chain ergodic (Haario et al. 2001). + +A good first check of any chain is to plot ``chain.log_posterior`` against +the iteration number, and ``chain.proposals`` against ``chain.samples`` for a +correlated pair. Proposals falling well outside the cloud of accepted samples +mean the proposal does not fit the posterior (Epstein et al. 2016, Fig. 3). + +.. autofunction:: mwg_sample + +.. autofunction:: adaptive_sample + +.. autoclass:: Chain + :members: n, kept, acceptance_rate, mode, save, load + +How much a chain is worth +""""""""""""""""""""""""" + +Samples from a Markov chain are correlated, so :math:`N` of them carry less +information than :math:`N` independent draws. The effective sample size is how +many independent draws they are worth. + +.. autofunction:: effective_sample_size + +.. autofunction:: significant_lags + +.. autofunction:: autocorrelation + +The Gaussian approximation +"""""""""""""""""""""""""" + +Maximum-likelihood standard errors come from the curvature of the log +likelihood at its maximum. The estimate is taken to be normally distributed +about the maximum, with covariance equal to the inverse of the negative +Hessian. :py:func:`gaussian_approximation` computes that distribution at any +mode: a fit's estimate, or the best sample of a chain. Overlaid on the +marginals a sampler finds, it shows where standard errors would misstate the +uncertainty. For the AChR data in Epstein et al. (2016, Fig. 7), that happens +for the monoliganded opening and shutting rates of site B. + +.. code-block:: python + + from HJCFIT.likelihood.mcmc import gaussian_approximation + + mode, _ = chain.mode() + approx = gaussian_approximation(posterior, mode) + approx.sd, approx.correlation # standard errors and correlations + approx.relative_error # how well the Hessian is determined + +.. autofunction:: gaussian_approximation + +.. autoclass:: GaussianApproximation + :members: sd, correlation, relative_error, marginal_pdf, sample + +.. autofunction:: hessian + +A Hessian needs a record that determines every rate. On the CH82 sample +record, a single concentration fitted with eight free rates, one direction is +so flat that the likelihood's own numerical noise decides its curvature. +``relative_error`` is then large, and at small steps the point may not even +appear to be a maximum. That is a statement about identifiability, not about +the differentiation: fitting several concentrations at once removes it. + +From a specification +"""""""""""""""""""" + +Everything above is available from a fit specification (see +:ref:`python_fitspec_api`), with no Python written. Add an ``[mcmc]`` section +and, to reproduce a published value, a ``[likelihood]`` section with the +root-finding settings it was computed with:: + + [likelihood] + nmax = 2 + xtol = 1e-12 + rtol = 1e-12 + + [mcmc] + n = 20000 # iterations per chain + burnin = 5000 # of which tuning, not kept + chains = 4 # in parallel processes + +and run it:: + + hjcfit sample my-fit.toml -o post + +By default this fits first and starts every chain at the maximum-likelihood +estimate. The inverse Hessian there shapes the first proposals, at the optimal +scale. ``start = "guess"`` runs the paper's own pilot sampler from the initial +guess instead. Chains after the first start from draws of the Gaussian +approximation, so that chains which agree have had a chance not to. The +command prints, for each rate, the posterior mean, standard deviation, 95% +interval, effective sample size and split-:math:`\hat R`, and writes +``post.json`` (all of that, the specification and the provenance) and one +``post_chain.npz`` per chain (:py:meth:`Chain.load` reads them). It exits +with status 2 when :math:`\hat R` exceeds 1.1: the chains plainly disagree, so +run them longer. + +.. autofunction:: potential_scale_reduction diff --git a/likelihood/python/CMakeLists.txt b/likelihood/python/CMakeLists.txt index 7486a61..487ad2c 100644 --- a/likelihood/python/CMakeLists.txt +++ b/likelihood/python/CMakeLists.txt @@ -70,7 +70,8 @@ if(pythonBindings) # Named rather than globbed, because CMake will not re-run on a new file # appearing and a glob would make the wheel depend on when cmake last ran. set(PYTHON_MODULES __init__.py random.py _methods.py optimization.py - simplex_hjc.py fitting.py fitspec.py runner.py cli.py) + simplex_hjc.py fitting.py fitspec.py runner.py cli.py + mcmc.py) install(FILES ${PYTHON_MODULES} ${CMAKE_CURRENT_BINARY_DIR}/likelihood.py diff --git a/likelihood/python/cli.py b/likelihood/python/cli.py index 13865c6..b3aa375 100644 --- a/likelihood/python/cli.py +++ b/likelihood/python/cli.py @@ -22,6 +22,7 @@ hjcfit template -o ch82.toml a specification to edit hjcfit check ch82.toml say what it would do; run nothing hjcfit fit ch82.toml -o out.json run it, and keep the result + hjcfit sample ch82.toml -o post sample the posterior ([mcmc]) Also reachable as ``python -m HJCFIT.likelihood.cli`` when the console script is not on the path, which is the usual state of affairs inside a @@ -80,6 +81,20 @@ def _parser(): fit.add_argument('-q', '--quiet', action='store_true', help='print the result and nothing before it') + sample = sub.add_parser( + 'sample', help='sample the posterior of the rates, as [mcmc] says') + _add_spec_argument(sample) + sample.add_argument('-o', '--output', metavar='PREFIX', default=None, + help='write PREFIX.json (summary, provenance) and ' + 'PREFIX_chain.npz (the chains)') + sample.add_argument('--chains', type=int, default=None, + help='number of chains, overriding [mcmc] chains') + sample.add_argument('--processes', type=int, default=None, + help='worker processes (default: one per chain, up ' + 'to the number of CPUs; 1 runs them in turn)') + sample.add_argument('-q', '--quiet', action='store_true', + help='print the result and nothing before it') + show = sub.add_parser( 'show', help='read a specification and write it back, normalised') _add_spec_argument(show) @@ -162,7 +177,8 @@ def _check(args): # only thing here that proves the records, the mechanism and the dead time # are mutually possible. A guess at which the likelihood cannot be # computed is a guess no search can start from. - fitter = HJCFitter(mec, records, log_params=spec.search.log_params) + fitter = HJCFitter(mec, records, log_params=spec.search.log_params, + solver=spec.likelihood.solver() or None) print('') print('log10L at the guess: {0:.4f}'.format(fitter.log10_likelihood())) print('nothing was fitted; run: hjcfit fit {0}'.format(args.spec)) @@ -191,6 +207,36 @@ def _fit(args): return 0 if outcome.result.success else 2 +def _sample(args): + from dataclasses import replace + + from .runner import sample, write_samples + + spec = _load(args.spec) + if args.chains is not None: + if args.chains < 1: + raise SpecError('--chains: must be at least 1, not {0}' + .format(args.chains)) + spec = replace(spec, mcmc=replace(spec.mcmc, chains=args.chains)) + if not args.quiet: + print(spec) + if not spec.mcmc.as_dict(): + print(' sampling: {0} (the defaults; no [mcmc] section)'.format( + spec.mcmc)) + print('') + outcome = sample(spec, processes=args.processes, verbose=not args.quiet) + if not args.quiet: + print('') + print(outcome) + if args.output: + paths = write_samples(args.output, outcome) + print('') + print('wrote {0}'.format(", ".join(paths))) + # As for fit: numbers are printed either way, and the status says whether + # to trust them. Chains that plainly disagree have not converged. + return 0 if not outcome.max_rhat > 1.1 else 2 + + def _show(args): sys.stdout.write(_load(args.spec).to_toml()) return 0 @@ -202,7 +248,8 @@ def main(argv=None): :param argv: Arguments, without the program name. None takes ``sys.argv``. :returns: 0 on success, 1 on anything this module can explain, and 2 - from ``fit`` when the search did not converge. + from ``fit`` when the search did not converge or from ``sample`` + when the chains disagree (split-R-hat above 1.1). """ parser = _parser() args = parser.parse_args(argv) @@ -215,7 +262,7 @@ def main(argv=None): return 0 handlers = {'template': _template, 'check': _check, 'fit': _fit, - 'show': _show} + 'sample': _sample, 'show': _show} try: return handlers[args.command](args) except SpecError as error: diff --git a/likelihood/python/fitspec.py b/likelihood/python/fitspec.py index 5569cae..38be130 100644 --- a/likelihood/python/fitspec.py +++ b/likelihood/python/fitspec.py @@ -68,8 +68,8 @@ one place, and says which names exist when one does not. """ __docformat__ = "restructuredtext en" -__all__ = ['DataSpec', 'MechanismSpec', 'SearchSpec', 'FitSpec', - 'SpecError', 'TEMPLATE', 'load_toml_bytes'] +__all__ = ['DataSpec', 'MechanismSpec', 'SearchSpec', 'LikelihoodSpec', + 'MCMCSpec', 'FitSpec', 'SpecError', 'TEMPLATE', 'load_toml_bytes'] from dataclasses import dataclass, field @@ -88,6 +88,12 @@ class SpecError(ValueError): #: Accepted values of :py:attr:`SearchSpec.method`. METHODS = ('simplex', 'scipy') +#: The root-finding options of the likelihood. The same tuple as +#: ``HJCFIT.likelihood.fitting.SOLVER_OPTIONS``, repeated because this module +#: imports nothing but the standard library; a test asserts the two agree. +SOLVER_KEYS = ('nmax', 'xtol', 'rtol', 'itermax', 'lower_bound', + 'upper_bound') + def _number(value, where): if isinstance(value, bool) or not isinstance(value, (int, float)): @@ -474,6 +480,190 @@ def as_dict(self): return d +@dataclass(frozen=True) +class LikelihoodSpec: + """ Root-finding settings for every record's likelihood. + + The names are those of + :py:data:`~HJCFIT.likelihood.fitting.SOLVER_OPTIONS`; anything left + out keeps :py:class:`~HJCFIT.likelihood.Log10Likelihood`'s default. + They are not cosmetic. On the three AChR records of Epstein et al. + (2016), ``nmax = 2`` with tolerances of 1e-12 -- that paper's settings + -- move the natural log-likelihood by 0.38 against the defaults at the + same rates, so a published value is reproduced only with the settings + it was computed with. + + Applies to ``hjcfit fit`` and ``hjcfit sample`` alike. + """ + + nmax: int = None + xtol: float = None + rtol: float = None + itermax: int = None + lower_bound: float = None + upper_bound: float = None + + KEYS = SOLVER_KEYS + + @classmethod + def from_dict(cls, d, where='likelihood'): + _unexpected(d, cls.KEYS, where) + values = {} + for name in ('nmax', 'itermax'): + if d.get(name) is not None: + value = d[name] + if isinstance(value, bool) or not isinstance(value, int): + raise SpecError("{0}.{1}: expected an integer, got {2!r}" + .format(where, name, value)) + values[name] = value + for name in ('xtol', 'rtol', 'lower_bound', 'upper_bound'): + if d.get(name) is not None: + values[name] = _number(d[name], "{0}.{1}".format(where, name)) + return cls(**values).validate(where) + + def validate(self, where='likelihood'): + if self.nmax is not None and self.nmax < 1: + raise SpecError("{0}.nmax: must be at least 1, not {1}" + .format(where, self.nmax)) + if self.itermax is not None and self.itermax < 1: + raise SpecError("{0}.itermax: must be at least 1, not {1}" + .format(where, self.itermax)) + for name in ('xtol', 'rtol'): + value = getattr(self, name) + if value is not None and value <= 0.0: + raise SpecError("{0}.{1}: must be positive, not {2}" + .format(where, name, value)) + if (self.lower_bound is not None and self.upper_bound is not None + and not self.lower_bound < self.upper_bound): + raise SpecError("{0}: lower_bound must be below upper_bound" + .format(where)) + return self + + def solver(self): + """ The settings given, as the ``solver`` argument of the fitter. """ + return {name: getattr(self, name) for name in self.KEYS + if getattr(self, name) is not None} + + def as_dict(self): + return self.solver() + + +@dataclass(frozen=True) +class MCMCSpec: + """ How to sample the posterior, for ``hjcfit sample``. + + The method is that of Epstein, Calderhead, Girolami & Sivilotti (2016); + see :ref:`python_mcmc_api`. + + :param start: + ``"fit"`` (the default) runs the ``[search]`` first and starts every + chain at the maximum-likelihood estimate, with the Hessian there as + the first proposal covariance. ``"guess"`` is the paper's own + scheme: a Metropolis-within-Gibbs pilot from the initial guess finds + the mode instead. With the paper's flat prior the two points + coincide; the fit is usually the cheaper way there. + :param sampler: + ``"adaptive"``, adaptive Metropolis, or ``"mwg"``, + Metropolis-within-Gibbs. + :param n: Iterations per chain, burn-in included. + :param burnin: + Iterations during which step sizes are tuned; not kept. + :param pilot_n: + ``start = "guess"`` only: iterations of the pilot, half of them + burn-in. + :param covariance: + ``"hessian"``: shape the adaptive sampler's first proposals with the + inverse Hessian at the start, at the optimal scale 2.38/sqrt(k). + Where the Hessian cannot be used -- a point that is not a maximum, a + direction no record determines -- the run says so and falls back to + ``"identity"``, a small isotropic step, which is the paper's. + :param mixture: ``"sum"`` (the paper's) or ``"choice"``; see + :py:func:`~HJCFIT.likelihood.mcmc.adaptive_sample`. + :param log_space: + Walk the logarithms of the rates. Unset, each sampler keeps its own + default: logarithms for ``"mwg"``, the rates themselves for + ``"adaptive"``. + :param prior: + ``"uniform"`` (the paper's) or ``"loguniform"``, in both cases + between each free rate's limits. + :param chains: + Independent chains, run in parallel processes. Chains after the + first start from draws of the Gaussian approximation at the start + point, when there is one, so that their agreement means something. + :param seed: Seeds chain *i* with ``seed + i``. + """ + + start: str = 'fit' + sampler: str = 'adaptive' + n: int = 20000 + burnin: int = 5000 + pilot_n: int = 3000 + covariance: str = 'hessian' + mixture: str = 'sum' + log_space: bool = None + prior: str = 'uniform' + chains: int = 1 + seed: int = 1 + + KEYS = ('start', 'sampler', 'n', 'burnin', 'pilot_n', 'covariance', + 'mixture', 'log_space', 'prior', 'chains', 'seed') + CHOICES = {'start': ('fit', 'guess'), 'sampler': ('adaptive', 'mwg'), + 'covariance': ('hessian', 'identity'), + 'mixture': ('sum', 'choice'), + 'prior': ('uniform', 'loguniform')} + + @classmethod + def from_dict(cls, d, where='mcmc'): + _unexpected(d, cls.KEYS, where) + values = {} + for name in cls.CHOICES: + if name in d: + values[name] = str(d[name]) + for name in ('n', 'burnin', 'pilot_n', 'chains', 'seed'): + if name in d: + value = d[name] + if isinstance(value, bool) or not isinstance(value, int): + raise SpecError("{0}.{1}: expected an integer, got {2!r}" + .format(where, name, value)) + values[name] = value + if 'log_space' in d: + if not isinstance(d['log_space'], bool): + raise SpecError("{0}.log_space: expected true or false, got " + "{1!r}".format(where, d['log_space'])) + values['log_space'] = d['log_space'] + return cls(**values).validate(where) + + def validate(self, where='mcmc'): + for name, allowed in self.CHOICES.items(): + if getattr(self, name) not in allowed: + raise SpecError("{0}.{1}: must be one of {2}, not {3!r}".format( + where, name, " or ".join(allowed), getattr(self, name))) + for name in ('n', 'pilot_n', 'chains'): + if getattr(self, name) < 1: + raise SpecError("{0}.{1}: must be at least 1, not {2}".format( + where, name, getattr(self, name))) + if not 0 <= self.burnin < self.n: + raise SpecError( + "{0}.burnin: must be at least 0 and below n ({1}), not {2}; " + "a chain that is all burn-in keeps nothing".format( + where, self.n, self.burnin)) + return self + + def as_dict(self): + default = MCMCSpec() + return {name: getattr(self, name) for name in self.KEYS + if getattr(self, name) != getattr(default, name) + and getattr(self, name) is not None} + + def __str__(self): + start = ("from the maximum-likelihood fit" if self.start == 'fit' + else "from a {0}-iteration pilot at the guess".format( + self.pilot_n)) + return ("{0} x {1} {2} iterations ({3} burn-in), {4}, {5} prior" + .format(self.chains, self.n, self.sampler, self.burnin, start, + self.prior)) + + @dataclass(frozen=True) class FitSpec: """ A whole fit, as data. @@ -484,6 +674,12 @@ class FitSpec: shared and the concentrations are not. :param mechanism: A :py:class:`MechanismSpec`. :param search: A :py:class:`SearchSpec`. + :param likelihood: + A :py:class:`LikelihoodSpec`: root-finding settings, for fitting + and sampling alike. + :param mcmc: + A :py:class:`MCMCSpec`, read by ``hjcfit sample`` and ignored by + ``hjcfit fit``. :param title: One line for the person reading the result. Read one with :py:meth:`from_toml`, write one with @@ -495,8 +691,10 @@ class FitSpec: mechanism: MechanismSpec = field(default_factory=MechanismSpec) search: SearchSpec = field(default_factory=SearchSpec) title: str = '' + likelihood: LikelihoodSpec = field(default_factory=LikelihoodSpec) + mcmc: MCMCSpec = field(default_factory=MCMCSpec) - KEYS = ('data', 'mechanism', 'search', 'title') + KEYS = ('data', 'mechanism', 'search', 'title', 'likelihood', 'mcmc') # -- reading ----------------------------------------------------------- @@ -525,6 +723,8 @@ def from_dict(cls, d, where='spec'): mechanism=MechanismSpec.from_dict(d.get('mechanism') or {}), search=SearchSpec.from_dict(d.get('search') or {}), title=str(d.get('title', '')), + likelihood=LikelihoodSpec.from_dict(d.get('likelihood') or {}), + mcmc=MCMCSpec.from_dict(d.get('mcmc') or {}), ).validate(where) @classmethod @@ -550,6 +750,8 @@ def validate(self, where='spec'): record.validate("data[{0}]".format(i)) self.mechanism.validate() self.search.validate() + self.likelihood.validate() + self.mcmc.validate() # Two records at the same concentration are not an error -- two patches # at one concentration is an ordinary experiment, and the reproduction # of Colquhoun, Hatton & Hawkes (2003) fits three concentrations of @@ -572,6 +774,10 @@ def as_dict(self): d['data'] = [r.as_dict() for r in self.data] d['mechanism'] = self.mechanism.as_dict() d['search'] = self.search.as_dict() + for name in ('likelihood', 'mcmc'): + table = getattr(self, name).as_dict() + if table: + d[name] = table return d def to_toml(self): @@ -589,8 +795,10 @@ def to_toml(self): lines.append('[[data]]') lines.extend(_toml_table(record.as_dict())) lines.append('') - for name in ('mechanism', 'search'): + for name in ('mechanism', 'search', 'likelihood', 'mcmc'): table = getattr(self, name).as_dict() + if not table and name in ('likelihood', 'mcmc'): + continue lines.append('[{0}]'.format(name)) lines.extend(_toml_table(table, name)) lines.append('') @@ -631,6 +839,12 @@ def __str__(self): self.search.method, 'logarithms' if self.search.log_params else 'rate constants', self.search.maxfev)) + solver = self.likelihood.solver() + if solver: + lines.append(' likelihood: {0}'.format(", ".join( + '{0} = {1:g}'.format(k, v) for k, v in solver.items()))) + if self.mcmc.as_dict(): + lines.append(' sampling: {0}'.format(self.mcmc)) return "\n".join(lines) @@ -773,4 +987,22 @@ def _toml_table(table, prefix=''): method = "simplex" # HJCFIT's own; "scipy" is Nelder-Mead with restarts log_params = true # search the logarithms: faster, and cannot go negative maxfev = 20000 + +# Root-finding settings of the likelihood, for fit and sample alike. Left out, +# each keeps its default. A published value is reproduced only with the +# settings it was computed with. +# [likelihood] +# nmax = 2 +# xtol = 1e-12 +# rtol = 1e-12 + +# For `hjcfit sample`: the posterior distribution of the free rates, by Markov +# chain Monte Carlo (Epstein et al. 2016). `hjcfit fit` ignores this section. +# [mcmc] +# start = "fit" # start at the maximum-likelihood fit; "guess" runs the +# # paper's pilot sampler from the guess instead +# n = 20000 # iterations per chain +# burnin = 5000 # of which tuning, not kept +# chains = 4 # run in parallel +# prior = "uniform" # between each rate's limits; or "loguniform" ''' diff --git a/likelihood/python/fitting.py b/likelihood/python/fitting.py index 3a54a9d..557cf5c 100644 --- a/likelihood/python/fitting.py +++ b/likelihood/python/fitting.py @@ -47,7 +47,7 @@ """ __docformat__ = "restructuredtext en" __all__ = ['Record', 'FitResult', 'HJCFitter', 'FAILURE_COST', - 'trim_to_openings'] + 'SOLVER_OPTIONS', 'trim_to_openings'] import time from dataclasses import dataclass, field @@ -61,6 +61,14 @@ #: rather than the perturbation HJCFIT itself used. FAILURE_COST = 1.0e10 +#: The root-finding options :py:class:`Log10Likelihood` accepts beyond the +#: record itself, and so what ``solver`` may hold. They are not cosmetic: on +#: the three AChR records of Epstein et al. (2016), ``nmax=2`` with tolerances +#: of 1e-12 -- that paper's settings -- gives a natural log-likelihood 0.38 +#: higher than the defaults (``nmax=3``, 1e-10) at the same rates. +SOLVER_OPTIONS = ('nmax', 'xtol', 'rtol', 'itermax', 'lower_bound', + 'upper_bound') + def trim_to_openings(intervals, amplitudes): """ Trim an alternating record so it starts and ends with an opening. @@ -225,14 +233,24 @@ class HJCFitter: Keep the best vertex at each iteration. Cheap. :param bool store_evaluations: Keep every point evaluated. Not cheap across many fits. + :param dict solver: + Root-finding options for every record's likelihood, keyed by the + names in :py:data:`SOLVER_OPTIONS`. Options left out keep + :py:class:`Log10Likelihood`'s defaults. To reproduce a published + value, pass the settings it was computed with. """ def __init__(self, mec, records, log_params=True, - store_path=False, store_evaluations=False): + store_path=False, store_evaluations=False, solver=None): self.mec = mec self.records = [r.check() for r in records] if not self.records: raise ValueError("no records to fit") + self.solver = dict(solver or {}) + unknown = sorted(set(self.solver) - set(SOLVER_OPTIONS)) + if unknown: + raise ValueError("unknown solver option(s) {0}; expected some of " + "{1}".format(unknown, list(SOLVER_OPTIONS))) self.log_params = log_params self.store_path = store_path self.store_evaluations = store_evaluations @@ -265,7 +283,7 @@ def __init__(self, mec, records, log_params=True, self.likelihoods = [ Log10Likelihood(r.as_lists(), nopen=mec.kA, tau=r.tres, - tcritical=r.tcrit) + tcritical=r.tcrit, **self.solver) for r in self.records ] self.nevals = 0 diff --git a/likelihood/python/mcmc.py b/likelihood/python/mcmc.py new file mode 100644 index 0000000..56e7b03 --- /dev/null +++ b/likelihood/python/mcmc.py @@ -0,0 +1,996 @@ +######################## +# HJCFIT computes missed-events likelihood as described in +# Hawkes, Jalali and Colquhoun (1990, 1992) +# +# Copyright (C) 2026 University College London +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +######################### + +""" Bayesian inference over a mechanism. + + :py:class:`~HJCFIT.likelihood.fitting.HJCFitter` looks for the rate + constants that maximise the HJC likelihood. This module treats them as + random variables instead, and describes their posterior distribution: the + likelihood of the records multiplied by a prior. It follows Epstein, + Calderhead, Girolami & Sivilotti (2016) *Biophys J* 111:333-348, who + sampled that posterior by Markov chain Monte Carlo with the same exact + missed-events likelihood. + + Like the fitter, **it does not import scalcs.** A mechanism is duck-typed: + ``theta()``, ``theta_unsqueeze()``, ``Rates``, ``kA``, ``set_eff()`` and + ``get_free_parameter_names()``, which + :py:class:`scalcs.mechanism.Mechanism` provides. + + Two things are deliberately different from fitting: + + * **Nothing is clipped.** The fitter resets a rate that leaves its limits + before the likelihood sees it, which suits a search. Here, a point + outside the prior has zero posterior density and is rejected. Clipping + it to the boundary would instead pile probability onto the boundary. + * **Failure is a value, not a penalty.** Where the likelihood cannot be + computed, the log posterior is ``-inf``. A sampler rejects the move, and + :py:attr:`LogPosterior.nfailures` counts how often that happened. + + Everything is in natural logarithms, because that is what probability + densities combine in. + + The samplers are the two stages of that paper: + + * :py:func:`mwg_sample`, Metropolis-within-Gibbs: a random walk on the + logarithm of one rate at a time, with each step size tuned during burn-in. + Robust from a poor start; it finds the mode, but mixes slowly along + correlated directions. + * :py:func:`adaptive_sample`, adaptive Metropolis (Haario, Saksman & + Tamminen 2001; Roberts & Rosenthal 2009): a block random walk whose + covariance is learned from the chain's own history. Started at the mode + the first stage found, it follows correlations such as the + :math:`\\alpha_2`-:math:`\\beta_2` ridge. + + Both return a :py:class:`Chain`. :py:func:`effective_sample_size` + measures what a chain is worth. + + The samplers take any callable returning a log density, not only a + :py:class:`LogPosterior`, which is how they are tested on targets whose + answers are known. +""" +__docformat__ = "restructuredtext en" +__all__ = ['UniformPrior', 'LogUniformPrior', 'LogPosterior', 'Chain', + 'mwg_sample', 'adaptive_sample', 'autocorrelation', + 'significant_lags', 'effective_sample_size', 'hessian', + 'GaussianApproximation', 'gaussian_approximation', + 'potential_scale_reduction'] + +import json +import time +from dataclasses import dataclass, field + +import numpy as np + + +def _free_rate_limits(mec): + """ (names, lower, upper) of the free rates, from their limits. """ + names, lower, upper = [], [], [] + missing = [] + for rate in mec.Rates: + if not rate.is_free: + continue + names.append(rate.name) + if not rate.limits: + missing.append(rate.name) + continue + lo, hi = rate.limits[0] + lower.append(float(lo)) + upper.append(float(hi)) + if missing: + raise ValueError("free rate(s) {0} have no limits, so a prior cannot " + "be taken from the mechanism; give the bounds " + "explicitly".format(missing)) + return tuple(names), np.array(lower), np.array(upper) + + +class UniformPrior: + """ Independent uniform distributions on the free rate constants. + + This is the prior of Epstein et al. (2016). It is flat between the + limits a maximum-likelihood fit would also impose, so the mode of the + posterior is the maximum-likelihood estimate. The paper uses + U(1e-2, 1e6) s\\ :sup:`-1` for every rate and U(1e-2, 1e10) + M\\ :sup:`-1` s\\ :sup:`-1` for association rates, which are the limits + scalcs gives its rates by default. + + :param lower: Lower bound of each free rate. + :param upper: Upper bound of each free rate. + :param names: Names of the free rates, in the same order. Optional. + """ + + def __init__(self, lower, upper, names=None): + self.lower = np.atleast_1d(np.asarray(lower, dtype=float)) + self.upper = np.atleast_1d(np.asarray(upper, dtype=float)) + if self.lower.shape != self.upper.shape or self.lower.ndim != 1: + raise ValueError("lower and upper must be matching 1-d arrays") + if not np.all(np.isfinite(self.lower) & np.isfinite(self.upper)): + raise ValueError("bounds must be finite for a proper prior") + if np.any(self.lower >= self.upper): + raise ValueError("every lower bound must be below its upper bound") + self.names = None if names is None else tuple(names) + if self.names is not None and len(self.names) != self.k: + raise ValueError("{0} names for {1} bounds" + .format(len(self.names), self.k)) + self._check_bounds() + + def _check_bounds(self): + pass + + @classmethod + def from_mechanism(cls, mec): + """ The prior bounded by each free rate's own limits. + + :raises ValueError: if a free rate has no limits. + """ + names, lower, upper = _free_rate_limits(mec) + return cls(lower, upper, names=names) + + @property + def k(self): + """ Number of parameters. """ + return len(self.lower) + + def _rates(self, rates): + rates = np.asarray(rates, dtype=float) + if rates.shape != (self.k,): + raise ValueError("expected {0} rates, got shape {1}" + .format(self.k, rates.shape)) + return rates + + def contains(self, rates): + """ Whether *rates* lie inside the bounds, boundaries included. """ + rates = self._rates(rates) + return bool(np.all((rates >= self.lower) & (rates <= self.upper))) + + def logpdf(self, rates): + """ Log density at *rates*: ``-inf`` outside the bounds. """ + if not self.contains(rates): + return -np.inf + return float(-np.sum(np.log(self.upper - self.lower))) + + def sample(self, rng, size=None): + """ Draws from the prior. + + :param rng: A :py:class:`numpy.random.Generator`. + :param size: Number of draws; None for a single vector. + """ + shape = (self.k,) if size is None else (size, self.k) + return rng.uniform(self.lower, self.upper, size=shape) + + +class LogUniformPrior(UniformPrior): + """ Independent log-uniform distributions on the free rate constants. + + Uniform in the logarithm of each rate between the bounds: every decade + is equally likely a priori. This is the scale-free alternative to + :py:class:`UniformPrior`, which puts almost all of its mass in the top + decade. The bounds must be positive. + """ + + def _check_bounds(self): + if np.any(self.lower <= 0): + raise ValueError("a log-uniform prior needs positive lower bounds") + + def logpdf(self, rates): + if not self.contains(rates): + return -np.inf + rates = self._rates(rates) + return float(-np.sum(np.log(rates)) + - np.sum(np.log(np.log(self.upper / self.lower)))) + + def sample(self, rng, size=None): + shape = (self.k,) if size is None else (size, self.k) + return np.exp(rng.uniform(np.log(self.lower), np.log(self.upper), + size=shape)) + + +class LogPosterior: + """ The log posterior density of a mechanism's free rate constants. + + Calling it with a vector of free rates returns + :math:`\\ln p(\\theta) + \\ln L(\\theta)`. The prior is evaluated + first, and the likelihood is not computed at all where the prior is + zero. + + :param mec: + A mechanism carrying its constraints. It is **modified**: every + evaluation puts its rates on the mechanism, as fitting does. + :param records: + A sequence of :py:class:`~HJCFIT.likelihood.fitting.Record`, fitted + simultaneously, each at its own concentration. + :param prior: + Anything with ``logpdf(rates)`` and ``k``. Defaults to + :py:meth:`UniformPrior.from_mechanism`. + :param dict solver: + Root-finding options for the likelihood; see + :py:data:`~HJCFIT.likelihood.fitting.SOLVER_OPTIONS`. + """ + + def __init__(self, mec, records, prior=None, solver=None): + from .fitting import HJCFitter + + # The fitter builds and holds one likelihood per record, and sums + # them at the mechanism's current rates. Its clipping lives in + # HJCFitter._apply, which is only reached through a parameter vector; + # evaluating it with none, below, is the unclipped path. + self.fitter = HJCFitter(mec, records, solver=solver) + self.mec = mec + self.names = tuple(mec.get_free_parameter_names()) + self.prior = UniformPrior.from_mechanism(mec) if prior is None else prior + if self.prior.k != len(self.names): + raise ValueError("the prior has {0} parameters but the mechanism " + "has {1} free rates".format(self.prior.k, + len(self.names))) + prior_names = getattr(self.prior, "names", None) + if prior_names is not None and tuple(prior_names) != self.names: + raise ValueError("the prior's parameters {0} are not the " + "mechanism's free rates {1}, in that order" + .format(list(prior_names), list(self.names))) + #: Likelihood evaluations requested, failures included. + self.nevals = 0 + #: Likelihood evaluations that could not be computed. + self.nfailures = 0 + + @property + def k(self): + """ Number of free parameters. """ + return len(self.names) + + @property + def records(self): + return self.fitter.records + + def _rates(self, rates): + rates = np.asarray(rates, dtype=float) + if rates.shape != (self.k,): + raise ValueError("expected {0} rates, got shape {1}" + .format(self.k, rates.shape)) + return rates + + def log_prior(self, rates): + """ :math:`\\ln p(\\theta)`. """ + return float(self.prior.logpdf(self._rates(rates))) + + def log_likelihood(self, rates): + """ :math:`\\ln L(\\theta)` summed over the records; ``-inf`` if it + cannot be computed. Rates are used as given, never clipped. """ + rates = self._rates(rates) + self.nevals += 1 + try: + if not np.all(np.isfinite(rates)): + raise ArithmeticError("rates are not finite") + self.mec.theta_unsqueeze(rates) + return float(self.fitter.ln_likelihood()) + except (ArithmeticError, ValueError, RuntimeError, FloatingPointError): + self.nfailures += 1 + return -np.inf + + def __call__(self, rates): + """ :math:`\\ln p(\\theta) + \\ln L(\\theta)`. """ + log_prior = self.log_prior(rates) + if log_prior == -np.inf: + return -np.inf + return log_prior + self.log_likelihood(rates) + + +# -------------------------------------------------------------------------- +# Chains +# -------------------------------------------------------------------------- + +@dataclass +class Chain: + """ What a sampler produced. + + :param samples: ``(n, k)`` parameter values, one row per iteration, as + rates even when the sampler walked their logarithms. + :param log_posterior: ``(n,)`` log density at each row of *samples*. + :param proposals: ``(n, k)`` the points proposed, accepted or not, as + rates. Plotting them against *samples* shows how well a proposal + fits the posterior (Epstein et al. 2016, Fig. 3). + :param accepted: Whether each proposal was accepted: ``(n,)`` for a + block sampler, ``(n, k)`` for one parameter at a time. + :param scales: The step-size multipliers in use at each iteration, + ``(n,)`` or ``(n, k)`` like *accepted*. + :param burnin: Iterations during which the sampler tuned itself. They + are not draws from the posterior; :py:meth:`kept` drops them. + :param names: The parameters' names, if the target had them. + :param sampler: Which sampler produced the chain. + :param settings: The sampler's settings, for provenance. + :param seconds: Wall-clock time of the run. + :param nevals: Calls to the log density. + :param nfailures: Of those, how many the target counted as failures. + :param proposal_covariance: The proposal covariance at the end of an + adaptive run, in the space it walked. + """ + + samples: np.ndarray + log_posterior: np.ndarray + proposals: np.ndarray + accepted: np.ndarray + scales: np.ndarray + burnin: int + names: tuple = None + sampler: str = "" + settings: dict = field(default_factory=dict) + seconds: float = 0.0 + nevals: int = 0 + nfailures: int = 0 + proposal_covariance: np.ndarray = None + + @property + def n(self): + """ Iterations, burn-in included. """ + return len(self.samples) + + def kept(self): + """ The samples after burn-in: the draws from the posterior. """ + return self.samples[self.burnin:] + + def acceptance_rate(self, after_burnin=True): + """ Fraction of proposals accepted; per parameter for one-at-a-time + sampling. """ + accepted = self.accepted[self.burnin:] if after_burnin else self.accepted + return accepted.mean(axis=0) + + def mode(self): + """ The sample with the highest log posterior, and that value. This is + where Epstein et al. (2016) start the adaptive sampler. """ + best = int(np.argmax(self.log_posterior)) + return self.samples[best].copy(), float(self.log_posterior[best]) + + def save(self, path): + """ Write the chain to a NumPy ``.npz`` file. """ + arrays = dict(samples=self.samples, log_posterior=self.log_posterior, + proposals=self.proposals, accepted=self.accepted, + scales=self.scales) + if self.proposal_covariance is not None: + arrays['proposal_covariance'] = self.proposal_covariance + meta = dict(burnin=self.burnin, names=self.names, + sampler=self.sampler, settings=self.settings, + seconds=self.seconds, nevals=self.nevals, + nfailures=self.nfailures) + np.savez_compressed(path, meta=np.array(json.dumps(meta)), **arrays) + + @classmethod + def load(cls, path): + """ Read a chain written by :py:meth:`save`. """ + with np.load(path, allow_pickle=False) as data: + meta = json.loads(str(data['meta'])) + names = meta.pop('names') + return cls(samples=data['samples'], + log_posterior=data['log_posterior'], + proposals=data['proposals'], + accepted=data['accepted'], scales=data['scales'], + proposal_covariance=(data['proposal_covariance'] + if 'proposal_covariance' in data + else None), + names=None if names is None else tuple(names), + **meta) + + +# -------------------------------------------------------------------------- +# Samplers +# -------------------------------------------------------------------------- + +class _Target: + """ The log density a sampler walks, in the space it walks it. + + Walking the logarithm of the rates changes the density by the + Jacobian of the transformation: a density :math:`\\pi(\\theta)` over + rates is :math:`\\pi(e^y)\\prod_i e^{y_i}` over :math:`y = \\ln\\theta`, + so :math:`\\sum_i y_i` is added. Leaving it out samples a different + distribution -- one that favours small rates. + """ + + def __init__(self, log_density, log_space): + self.log_density = log_density + self.log_space = log_space + self.nevals = 0 + + def to_state(self, rates): + rates = np.asarray(rates, dtype=float) + if self.log_space: + if np.any(rates <= 0): + raise ValueError("a sampler over log rates needs positive " + "starting rates") + return np.log(rates) + return rates.copy() + + def to_rates(self, state): + return np.exp(state) if self.log_space else state + + def __call__(self, state): + self.nevals += 1 + value = float(self.log_density(self.to_rates(state))) + if np.isnan(value): + return -np.inf + if self.log_space and value > -np.inf: + value += float(np.sum(state)) + return value + + +def _start(log_density, x0, log_space): + target = _Target(log_density, log_space) + state = target.to_state(x0) + if state.ndim != 1: + raise ValueError("the starting point must be a 1-d vector") + log_p = target(state) + if not np.isfinite(log_p): + raise ValueError("the log density at the starting point is {0}; a " + "chain must start where the posterior is positive" + .format(log_p)) + return target, state, log_p + + +def _accept(rng, log_p_new, log_p): + """ The Metropolis test for a symmetric proposal. """ + if log_p_new == -np.inf: + return False + return bool(np.log(rng.random()) < log_p_new - log_p) + + +def _tune(scale, window, low, high, step): + """ Shrink a scale when too few proposals are accepted, grow it when too + many are. Used during burn-in only, so that the chain that is kept is + a Markov chain with a fixed kernel. """ + rate = window.mean() + if rate < low: + return scale * (1.0 - step) + if rate > high: + return scale * (1.0 + step) + return scale + + +def _check_run(n, burnin): + if not 0 <= burnin <= n: + raise ValueError("burnin must lie between 0 and n") + + +def _finish(target, log_density, sampler, settings, started, **arrays): + return Chain(names=getattr(log_density, 'names', None), sampler=sampler, + settings=settings, seconds=time.perf_counter() - started, + nevals=target.nevals, + nfailures=int(getattr(log_density, 'nfailures', 0)), + **arrays) + + +def mwg_sample(log_density, x0, n, burnin, rng=None, *, initial_scale=1.0, + tune_every=50, acceptance=(0.1, 0.5), tune_step=0.1, + log_space=True, callback=None, callback_every=1000): + """ Metropolis-within-Gibbs: a random walk on one parameter at a time. + + Each iteration visits every parameter in turn and proposes + :math:`y_i' = y_i + s_i z`, :math:`z \\sim N(0, 1)`, on the logarithm + of the rate (a multiplicative step on the rate itself), accepting it by + the Metropolis rule. This is the pilot sampler of Epstein et al. + (2016): multiplicative steps suit rates spanning several decades, and + it locates the posterior mode reliably from a poor starting point. + + During burn-in, every *tune_every* iterations, each :math:`s_i` is + multiplied by ``1 - tune_step`` if that parameter's acceptance over the + window is below ``acceptance[0]``, or by ``1 + tune_step`` if above + ``acceptance[1]``. + + :param log_density: Callable returning the log posterior at a vector + of rates, such as a :py:class:`LogPosterior`. + :param x0: Starting rates; the density there must be positive. + :param int n: Iterations, each a sweep over all parameters. + :param int burnin: Iterations during which the step sizes are tuned. + :param rng: A :py:class:`numpy.random.Generator`, or a seed. + :param initial_scale: Starting :math:`s_i`, a scalar or one per + parameter. On log rates, 1 is a step of about a factor of e. + :param bool log_space: Walk the logarithms of the rates (default) or + the rates themselves. + :param callback: Called as ``callback(iteration, chain)`` every + *callback_every* iterations with the chain so far. + :returns: A :py:class:`Chain` with per-parameter ``accepted`` and + ``scales``. + """ + _check_run(n, burnin) + rng = np.random.default_rng(rng) + settings = dict(n=n, burnin=burnin, initial_scale=initial_scale, + tune_every=tune_every, acceptance=list(acceptance), + tune_step=tune_step, log_space=log_space) + started = time.perf_counter() + target, state, log_p = _start(log_density, x0, log_space) + k = len(state) + scale = np.broadcast_to(np.asarray(initial_scale, dtype=float), (k,)).copy() + settings['initial_scale'] = scale.tolist() + low, high = acceptance + + samples = np.empty((n, k)) + proposals = np.empty((n, k)) + accepted = np.zeros((n, k), dtype=bool) + scales = np.empty((n, k)) + log_posterior = np.empty(n) + + def partial(upto): + return _finish(target, log_density, 'mwg', settings, started, + samples=target.to_rates(samples[:upto]), + log_posterior=log_posterior[:upto], + proposals=target.to_rates(proposals[:upto]), + accepted=accepted[:upto], scales=scales[:upto], + burnin=min(burnin, upto)) + + for it in range(n): + scales[it] = scale + for i in range(k): + proposal = state.copy() + proposal[i] += scale[i] * rng.standard_normal() + proposals[it, i] = proposal[i] + log_p_new = target(proposal) + if _accept(rng, log_p_new, log_p): + state, log_p = proposal, log_p_new + accepted[it, i] = True + samples[it] = state + # The log density of the rates, without the Jacobian of the walk. + log_posterior[it] = log_p - (np.sum(state) if log_space else 0.0) + + done = it + 1 + if done <= burnin and done % tune_every == 0: + window = accepted[done - tune_every:done] + for i in range(k): + scale[i] = _tune(scale[i], window[:, i], low, high, tune_step) + if callback is not None and done % callback_every == 0: + callback(done, partial(done)) + + return partial(n) + + +class _RunningCovariance: + """ Mean and covariance of a growing sample, one row at a time (Welford). + Divides by n - 1, as ``numpy.cov`` does. """ + + def __init__(self, k): + self.count = 0 + self.mean = np.zeros(k) + self._m2 = np.zeros((k, k)) + + def add(self, x): + self.count += 1 + delta = x - self.mean + self.mean += delta / self.count + self._m2 += np.outer(delta, x - self.mean) + + @property + def covariance(self): + return self._m2 / (self.count - 1) + + +def _cholesky(matrix): + """ Lower Cholesky factor, adding the smallest jitter that makes a + near-singular sample covariance usable. """ + jitter = 0.0 + size = float(np.mean(np.abs(np.diag(matrix)))) or 1.0 + for _ in range(12): + try: + return np.linalg.cholesky(matrix + jitter * np.eye(len(matrix))) + except np.linalg.LinAlgError: + jitter = size * 1e-12 if jitter == 0.0 else jitter * 10.0 + raise np.linalg.LinAlgError("proposal covariance is not positive definite") + + +def adaptive_sample(log_density, x0, n, burnin, rng=None, *, adapt_start=100, + beta=0.05, initial_step=0.1, initial_covariance=None, + mixture='sum', tune_every=50, acceptance=(0.1, 0.5), + tune_step=0.1, log_space=False, callback=None, + callback_every=1000): + """ Adaptive Metropolis: a block random walk that learns its covariance. + + For the first *adapt_start* iterations the proposal is + :math:`x + (\\epsilon/\\sqrt{k})\\,L_0 z`, where :math:`\\epsilon` is + *initial_step* and :math:`L_0 L_0^T` is *initial_covariance*, or the + identity if none is given. From then on :math:`\\hat\\Sigma`, the + covariance of every sample so far, takes over. It is scaled by + :math:`2.38^2/k`, which is optimal for a Gaussian target (Roberts & + Rosenthal 2001). It is also mixed with the small step, so that the + chain cannot lock onto a degenerate :math:`\\hat\\Sigma`: + + * ``mixture='sum'``, as in Epstein et al. (2016): + :math:`x' = x + (1-\\beta)\\,L z_1 + \\beta\\,(\\epsilon/\\sqrt{k})\\,L_0 z_2`, + where :math:`L L^T = (2.38^2/k)\\,s\\,\\hat\\Sigma`. + * ``mixture='choice'``, as in Roberts & Rosenthal (2009): the small + step with probability :math:`\\beta`, otherwise the learned one. + + Both are symmetric. During burn-in, a global multiplier :math:`s` on + :math:`\\hat\\Sigma` is tuned from the acceptance rate every + *tune_every* iterations, as in :py:func:`mwg_sample`. + + By default the walk is on the rates themselves, as in that paper; start + it at the mode :py:func:`mwg_sample` found. With ``log_space=True`` it + walks their logarithms instead. + + :param log_density: Callable returning the log posterior at a vector + of rates, such as a :py:class:`LogPosterior`. + :param x0: Starting rates; the density there must be positive. + :param int n: Iterations. + :param int burnin: Iterations during which *s* is tuned. The + covariance keeps adapting afterwards. The weight of each new sample + in it falls as :math:`1/n`, which is what keeps the chain ergodic + (Haario et al. 2001). + :param rng: A :py:class:`numpy.random.Generator`, or a seed. + :param int adapt_start: Iterations before the learned covariance is + used. + :param float beta: Weight of the small step. + :param initial_covariance: :math:`L_0 L_0^T`. The inverse Hessian at a + maximum-likelihood fit is a good choice. It shapes the small step + throughout. + :param callback: Called as ``callback(iteration, chain)`` every + *callback_every* iterations with the chain so far. + :returns: A :py:class:`Chain`; ``proposal_covariance`` is the learned + covariance at the end, scaled as it was being used. + """ + if mixture not in ('sum', 'choice'): + raise ValueError("mixture must be 'sum' or 'choice', not " + + repr(mixture)) + if not 0.0 <= beta <= 1.0: + raise ValueError("beta must lie between 0 and 1") + _check_run(n, burnin) + rng = np.random.default_rng(rng) + settings = dict(n=n, burnin=burnin, adapt_start=adapt_start, beta=beta, + initial_step=initial_step, mixture=mixture, + tune_every=tune_every, acceptance=list(acceptance), + tune_step=tune_step, log_space=log_space, + initial_covariance=initial_covariance is not None) + started = time.perf_counter() + target, state, log_p = _start(log_density, x0, log_space) + k = len(state) + if initial_covariance is None: + small = initial_step / np.sqrt(k) * np.eye(k) + else: + c0 = np.asarray(initial_covariance, dtype=float) + if c0.shape != (k, k): + raise ValueError("initial_covariance must be {0} x {0}".format(k)) + small = initial_step / np.sqrt(k) * _cholesky(c0) + optimal = 2.38 ** 2 / k + s = 1.0 + low, high = acceptance + + running = _RunningCovariance(k) + samples = np.empty((n, k)) + proposals = np.empty((n, k)) + accepted = np.zeros(n, dtype=bool) + scales = np.empty(n) + log_posterior = np.empty(n) + learned_factor = None + + def partial(upto): + return _finish(target, log_density, 'adaptive', settings, started, + samples=target.to_rates(samples[:upto]), + log_posterior=log_posterior[:upto], + proposals=target.to_rates(proposals[:upto]), + accepted=accepted[:upto], scales=scales[:upto], + burnin=min(burnin, upto), + proposal_covariance=( + None if learned_factor is None + else learned_factor @ learned_factor.T)) + + for it in range(n): + scales[it] = s + if it < adapt_start or running.count < 2: + proposal = state + small @ rng.standard_normal(k) + else: + learned_factor = _cholesky(optimal * s * running.covariance) + learned = learned_factor @ rng.standard_normal(k) + isotropic = small @ rng.standard_normal(k) + if mixture == 'sum': + proposal = state + (1.0 - beta) * learned + beta * isotropic + elif rng.random() < beta: + proposal = state + isotropic + else: + proposal = state + learned + proposals[it] = proposal + log_p_new = target(proposal) + if _accept(rng, log_p_new, log_p): + state, log_p = proposal, log_p_new + accepted[it] = True + samples[it] = state + log_posterior[it] = log_p - (np.sum(state) if log_space else 0.0) + running.add(state) + + done = it + 1 + if done <= burnin and done % tune_every == 0: + s = _tune(s, accepted[done - tune_every:done], low, high, + tune_step) + if callback is not None and done % callback_every == 0: + callback(done, partial(done)) + + return partial(n) + + +# -------------------------------------------------------------------------- +# Diagnostics +# -------------------------------------------------------------------------- + +def autocorrelation(x, max_lag): + """ Sample autocorrelation of a series at lags ``0 .. max_lag``. + + The biased estimator: autocovariances divided by the length of the + series, then by the variance. It is the usual estimator for MCMC + output, and the one MATLAB's ``autocorr`` computes. + """ + x = np.asarray(x, dtype=float) + if x.ndim != 1: + raise ValueError("autocorrelation takes a 1-d series") + n = len(x) + max_lag = min(int(max_lag), n - 1) + centred = x - x.mean() + variance = centred @ centred / n + if variance == 0.0: + raise ValueError("the series is constant") + return np.array([centred[:n - lag] @ centred[lag:] / n / variance + for lag in range(max_lag + 1)]) + + +def significant_lags(x, max_lag=200, nsd=2.0): + """ The first lag whose autocorrelation falls below + :math:`\\text{nsd}/\\sqrt{N}`, the approximate bound for a series with + no autocorrelation; *max_lag* if none does. + + Epstein et al. (2016, Table 2) report this as the number of + significant lags, and truncate the effective sample size there. + """ + x = np.asarray(x, dtype=float) + rho = autocorrelation(x, max_lag) + below = np.nonzero(rho < nsd / np.sqrt(len(x)))[0] + return int(below[0]) if len(below) else len(rho) - 1 + + +def effective_sample_size(x, max_lag=None): + """ Effective sample size by Geyer's (1992) initial monotone sequence. + + Autocorrelations are summed in adjacent pairs, + :math:`\\Gamma_m = \\rho_{2m} + \\rho_{2m+1}`. The pairs are made + non-increasing and kept while positive, and then + :math:`\\text{ESS} = N / (-1 + 2\\sum_m \\Gamma_m)`. The denominator + is floored at 1, so the estimate never exceeds *N*. This is the + estimator of Epstein et al. (2016). + + :param x: A series, or an ``(n, k)`` array of them. + :param max_lag: Autocorrelations beyond this lag are not used. None + takes each series' own :py:func:`significant_lags`. That paper + truncates every parameter at the significant lags of + :math:`\\alpha_2`; pass that number to do the same. + :returns: A float, or one per column. + """ + x = np.asarray(x, dtype=float) + if x.ndim == 2: + return np.array([effective_sample_size(col, max_lag) for col in x.T]) + lag = significant_lags(x) if max_lag is None else int(max_lag) + rho = autocorrelation(x, lag) + npairs = len(rho) // 2 + gamma = rho[0:2 * npairs:2] + rho[1:2 * npairs:2] + gamma = np.minimum.accumulate(gamma) + tau = max(1.0, -rho[0] + 2.0 * gamma[gamma > 0].sum()) + return len(x) / tau + + +# -------------------------------------------------------------------------- +# The Gaussian approximation at the mode +# -------------------------------------------------------------------------- + +def _second_differences(f, x, f0, steps): + """ Central second differences of *f* at *x* with per-coordinate + *steps*: the Hessian to O(h^2). """ + k = len(x) + + def at(offsets): + value = float(f(x + offsets)) + if not np.isfinite(value): + moved = np.nonzero(offsets)[0].tolist() + raise ValueError( + "the log density is {0} a step of {1} away from the point, " + "along parameter(s) {2}. The step leaves the prior or the " + "likelihood fails there; use a smaller rel_step, or a point " + "further from the boundary".format( + value, offsets[moved].tolist(), moved)) + return value + + h = np.zeros((k, k)) + for i in range(k): + e = np.zeros(k) + e[i] = steps[i] + h[i, i] = (at(e) - 2.0 * f0 + at(-e)) / steps[i] ** 2 + for i in range(k): + for j in range(i + 1, k): + ei = np.zeros(k) + ej = np.zeros(k) + ei[i] = steps[i] + ej[j] = steps[j] + h[i, j] = h[j, i] = (at(ei + ej) - at(ei - ej) - at(ej - ei) + + at(-ei - ej)) / (4.0 * steps[i] * steps[j]) + return h + + +def hessian(f, x, rel_step=3e-3, steps=None): + """ Hessian of a scalar function by central differences with one + Richardson extrapolation. + + Second differences with steps :math:`h` and :math:`h/2` each carry an + :math:`O(h^2)` error. :math:`(4 D_{h/2} - D_h) / 3` cancels it, + leaving :math:`O(h^4)`, and the difference between that and + :math:`D_{h/2}` estimates the error left. + + Steps are relative, :math:`h_i = \\text{rel\\_step} \\times |x_i|`, + because rate constants span eight decades: one absolute step cannot + suit a rate of 1 s\\ :sup:`-1` and one of 10\\ :sup:`8` + M\\ :sup:`-1` s\\ :sup:`-1` at once. + + The step is a trade-off. Too large, and the extrapolation has not yet + removed the curvature error. Too small, and noise in *f* is divided by + :math:`h^2`. The HJC likelihood carries noise from its root finding: + about 5e-8 in ln L with the default tolerances. On the three AChR + records of Epstein et al. (2016), steps from 3e-2 down to 1e-3 give + the same standard deviations, with the smallest error estimate near + the default. Tightening the likelihood's ``solver`` tolerances lowers + the noise floor. + + :param f: Scalar function of a vector. + :param x: Where to differentiate. + :param float rel_step: Step relative to each coordinate. + :param steps: Absolute steps, one per coordinate, instead. + :returns: ``(H, error)``, both ``k x k``. + :raises ValueError: if *f* is not finite at any point used. A + difference across a prior boundary or a failed likelihood would + otherwise give a meaningless number without complaint. + """ + x = np.asarray(x, dtype=float) + if steps is None: + steps = rel_step * np.abs(x) + if np.any(steps == 0.0): + raise ValueError("a relative step is zero where a coordinate is " + "zero; give absolute steps") + else: + steps = np.broadcast_to(np.asarray(steps, dtype=float), + x.shape).copy() + f0 = float(f(x)) + if not np.isfinite(f0): + raise ValueError("the log density at the point is {0}".format(f0)) + coarse = _second_differences(f, x, f0, steps) + fine = _second_differences(f, x, f0, steps / 2.0) + extrapolated = (4.0 * fine - coarse) / 3.0 + return extrapolated, np.abs(extrapolated - fine) + + +@dataclass +class GaussianApproximation: + """ A normal distribution matched to the posterior at its mode. + + The mean is the mode and the covariance is :math:`(-H)^{-1}`, the + inverse of the negative Hessian of the log posterior there. With a + flat prior, the mode is the maximum-likelihood estimate and this is + the asymptotic distribution of that estimate: the basis of + maximum-likelihood standard errors and correlations (Colquhoun, + Hatton & Hawkes 2003). Epstein et al. (2016, Figs 4 and 7) overlay it + on the marginals the sampler finds. Where the two differ, the + standard errors misstate the uncertainty. + + :param mode: Where the approximation was made. + :param covariance: :math:`(-H)^{-1}`. + :param hessian: :math:`H`. + :param hessian_error: Estimated error in each element of *hessian*. + :param log_density: The log density at the mode. + :param names: Parameter names, if the density had them. + """ + + mode: np.ndarray + covariance: np.ndarray + hessian: np.ndarray + hessian_error: np.ndarray + log_density: float + names: tuple = None + + @property + def relative_error(self): + """ The largest estimated error in the Hessian, relative to the scale + of its row and column: :math:`\\max_{ij} \\epsilon_{ij} / + \\sqrt{|H_{ii} H_{jj}|}`. Small (1e-4 or less) when the Hessian is + well determined. Near 1, a direction is barely curved at all, and + the noise of the likelihood decides it. A single-concentration + record often does this, because it does not identify every rate. """ + scale = np.sqrt(np.outer(np.abs(np.diag(self.hessian)), + np.abs(np.diag(self.hessian)))) + return float(np.max(self.hessian_error / scale)) + + @property + def sd(self): + """ Standard deviations: square roots of the covariance diagonal. """ + return np.sqrt(np.diag(self.covariance)) + + @property + def correlation(self): + """ The correlation matrix. """ + sd = self.sd + return self.covariance / np.outer(sd, sd) + + def _index(self, parameter): + if isinstance(parameter, str): + if self.names is None: + raise ValueError("this approximation has no parameter names") + return self.names.index(parameter) + return int(parameter) + + def marginal_pdf(self, parameter, x): + """ The normal density of one parameter, by index or name, at *x*. """ + i = self._index(parameter) + sd = self.sd[i] + z = (np.asarray(x, dtype=float) - self.mode[i]) / sd + return np.exp(-0.5 * z * z) / (sd * np.sqrt(2.0 * np.pi)) + + def sample(self, rng, size=None): + """ Draws from the approximation. """ + return np.random.default_rng(rng).multivariate_normal( + self.mode, self.covariance, size=size) + + +def gaussian_approximation(log_density, mode, rel_step=3e-3, steps=None): + """ The :py:class:`GaussianApproximation` to a posterior at its mode. + + :param log_density: Callable returning the log posterior at a vector + of rates, such as a :py:class:`LogPosterior`. + :param mode: The mode: a fit's maximum-likelihood estimate, or the + best sample of a chain (:py:meth:`Chain.mode`). The Hessian is + evaluated there, so a point that is not a maximum is refused. + :param rel_step: See :py:func:`hessian`. + :param steps: See :py:func:`hessian`. + :raises ValueError: if the negative Hessian is not positive definite, + so the point is not a maximum; or see :py:func:`hessian`. + """ + mode = np.asarray(mode, dtype=float).copy() + h, error = hessian(log_density, mode, rel_step=rel_step, steps=steps) + h = 0.5 * (h + h.T) + eigenvalues = np.linalg.eigvalsh(-h) + if np.any(eigenvalues <= 0): + raise ValueError( + "the point is not a maximum: the negative Hessian has {0} " + "non-positive eigenvalue(s) (smallest {1:.3g})".format( + int(np.sum(eigenvalues <= 0)), eigenvalues.min())) + # Invert through the Cholesky factor, so that the covariance is exactly + # symmetric and positive definite. + factor = np.linalg.cholesky(-h) + inverse = np.linalg.inv(factor) + covariance = inverse.T @ inverse + return GaussianApproximation( + mode=mode, covariance=covariance, hessian=h, hessian_error=error, + log_density=float(log_density(mode)), + names=getattr(log_density, 'names', None)) + + +def potential_scale_reduction(chains): + """ Split-:math:`\\hat R` (Gelman et al. 2013) for one parameter. + + Each chain is split in half, and the variance between the halves' + means is compared with the variance within them. Near 1 when every + chain samples the same distribution; values above about 1.01 mean + the chains have not yet agreed. Splitting makes a single chain + testable too, against a drift between its first and second halves. + + :param chains: A sequence of 1-d series, one per chain, burn-in + already removed. They are truncated to the shortest. + :returns: :math:`\\hat R`. + """ + series = [np.asarray(c, dtype=float) for c in chains] + n = min(len(s) for s in series) // 2 + if n < 2: + raise ValueError("each chain needs at least four samples") + halves = np.array([h for s in series for h in (s[:n], s[n:2 * n])]) + within = halves.var(axis=1, ddof=1).mean() + between = n * halves.mean(axis=1).var(ddof=1) + if within == 0.0: + return 1.0 if between == 0.0 else np.inf + pooled = (n - 1) / n * within + between / n + return float(np.sqrt(pooled / within)) diff --git a/likelihood/python/runner.py b/likelihood/python/runner.py index 9c0a275..4287153 100644 --- a/likelihood/python/runner.py +++ b/likelihood/python/runner.py @@ -52,7 +52,8 @@ __docformat__ = "restructuredtext en" __all__ = ['Outcome', 'build_mechanism', 'load_records', 'run', 'against_limits', 'provenance', 'result_as_dict', 'write_result', - 'RunnerError'] + 'RunnerError', 'SampleOutcome', 'sample', 'samples_as_dict', + 'write_samples'] import json import os @@ -437,7 +438,8 @@ def run(spec, verbose=False, store_path=False, x0=None): search = spec.search fitter = HJCFitter(mec, records, log_params=search.log_params, - store_path=store_path) + store_path=store_path, + solver=spec.likelihood.solver() or None) if verbose: # Not the records: Outcome.__str__ prints those, so that a run with # verbose off still says what was fitted. @@ -540,3 +542,339 @@ def write_result(path, outcome): json.dump(result_as_dict(outcome), handle, indent=2, sort_keys=False) handle.write("\n") return path + + +# --------------------------------------------------------------- sampling + +def _posterior(spec): + """ Records, mechanism and log posterior, built from the specification. + + Every chain process builds its own. The likelihood objects wrap C++ + state and cannot be pickled, and a specification can. + """ + from .mcmc import LogPosterior, LogUniformPrior, UniformPrior + + records = load_records(spec) + mec = build_mechanism(spec) + prior_class = (UniformPrior if spec.mcmc.prior == 'uniform' + else LogUniformPrior) + post = LogPosterior(mec, records, prior=prior_class.from_mechanism(mec), + solver=spec.likelihood.solver() or None) + return records, mec, post + + +def _run_chain(spec, index, start, covariance): + """ One chain of ``spec.mcmc``: the unit of work of one process. """ + from .mcmc import adaptive_sample, mwg_sample + + _, _, post = _posterior(spec) + m = spec.mcmc + start = np.asarray(start, dtype=float) + kwargs = {} + if m.log_space is not None: + kwargs['log_space'] = m.log_space + log_space = (m.log_space if m.log_space is not None + else m.sampler == 'mwg') + if covariance is not None: + covariance = np.asarray(covariance, dtype=float) + if log_space: + # The same Gaussian, to first order, on the logarithms. + covariance = covariance / np.outer(start, start) + if m.sampler == 'adaptive': + if covariance is not None: + # Proposals shaped like the posterior at the optimal scale from + # the first iteration, instead of the paper's small isotropic step + # while the covariance is learned. + kwargs.update(initial_covariance=covariance, initial_step=2.38) + return adaptive_sample(post, start, n=m.n, burnin=m.burnin, + rng=m.seed + index, mixture=m.mixture, + **kwargs) + if covariance is not None: + kwargs['initial_scale'] = 2.38 * np.sqrt(np.diag(covariance)) + return mwg_sample(post, start, n=m.n, burnin=m.burnin, + rng=m.seed + index, **kwargs) + + +def _starts(post, start, covariance, chains, seed): + """ The first chain at *start*; the others at draws from the Gaussian + approximation that the prior allows, so that chains which agree have + had a chance not to. """ + starts = [np.asarray(start, dtype=float)] + rng = np.random.default_rng(seed + 10007) + for _ in range(1, chains): + point = starts[0] + if covariance is not None: + for _ in range(100): + draw = rng.multivariate_normal(starts[0], covariance) + if np.isfinite(post(draw)): + point = draw + break + starts.append(point) + return starts + + +def _chain_summary(names, chains): + """ Pooled posterior summary per free rate, over the kept samples. """ + from .mcmc import effective_sample_size, potential_scale_reduction + + rows = [] + for i, name in enumerate(names): + series = [c.kept()[:, i] for c in chains] + pooled = np.concatenate(series) + ess = 0.0 + for s in series: + try: + ess += effective_sample_size(s) + except ValueError: # a constant series: nothing moved + pass + try: + rhat = potential_scale_reduction(series) + except ValueError: + rhat = float('nan') + q = np.percentile(pooled, [2.5, 50.0, 97.5]) + rows.append({'name': name, 'mean': float(pooled.mean()), + 'sd': float(pooled.std(ddof=1)), + 'q2.5': float(q[0]), 'median': float(q[1]), + 'q97.5': float(q[2]), 'ess': float(ess), + 'rhat': float(rhat)}) + return rows + + +@dataclass +class SampleOutcome: + """ Everything one sampling run produced. + + :param spec: The specification. + :param records: The records as they were sampled against. + :param start: Where the first chain started, as rates. + :param start_from: ``'fit'`` or ``'pilot'``. + :param fit: The :py:class:`Outcome` of the fit, for ``start = "fit"``. + :param pilot: The pilot :py:class:`~HJCFIT.likelihood.mcmc.Chain`, for + ``start = "guess"``. + :param approximation: The + :py:class:`~HJCFIT.likelihood.mcmc.GaussianApproximation` at the + start, when the Hessian could be used. + :param chains: One :py:class:`~HJCFIT.likelihood.mcmc.Chain` each. + :param summary: Per free rate: mean, sd, quantiles, ESS summed over + chains and split-R-hat across them. + :param notes: Anything the run had to decide that the reader should + know, such as falling back from the Hessian. + :param seconds: Wall-clock time of the whole run. + """ + + spec: object + records: list + names: tuple + start: np.ndarray + start_from: str + chains: list + summary: list + fit: object = None + pilot: object = None + approximation: object = None + notes: list = field(default_factory=list) + seconds: float = 0.0 + provenance: dict = field(default_factory=dict, repr=False) + + @property + def max_rhat(self): + values = [row['rhat'] for row in self.summary + if np.isfinite(row['rhat'])] + return max(values) if values else float('nan') + + def __str__(self): + lines = [str(r) for r in self.records] + lines.append("") + lines.append("{0} chain{1} of {2} from the {3}; {4:.1f} min".format( + len(self.chains), "" if len(self.chains) == 1 else "s", + self.spec.mcmc.n, "maximum-likelihood fit" + if self.start_from == 'fit' else "pilot's best sample", + self.seconds / 60.0)) + for note in self.notes: + lines.append(" note: " + note) + lines.append("") + lines.append(" {0:9s} {1:>12s} {2:>11s} {3:>12s} {4:>12s} {5:>7s} " + "{6:>6s}".format("rate", "mean", "sd", "2.5%", "97.5%", + "ESS", "R-hat")) + for row in self.summary: + lines.append(" {0:9s} {1:12.4g} {2:11.3g} {3:12.4g} {4:12.4g} " + "{5:7.0f} {6:6.3f}".format( + row['name'], row['mean'], row['sd'], row['q2.5'], + row['q97.5'], row['ess'], row['rhat'])) + failures = sum(c.nfailures for c in self.chains) + acceptance = [float(np.mean(c.acceptance_rate())) + for c in self.chains] + lines.append("") + lines.append("acceptance after burn-in: {0}; likelihood failures: {1}" + .format(", ".join("{0:.2f}".format(a) + for a in acceptance), failures)) + if self.max_rhat > 1.01: + lines.append("R-hat above 1.01: the chains have not agreed yet. " + "Run them longer before using these numbers.") + return "\n".join(lines) + + +def sample(spec, processes=None, verbose=False): + """ Sample the posterior a specification describes (``[mcmc]``). + + :param spec: A :py:class:`~HJCFIT.likelihood.fitspec.FitSpec`. + :param processes: Worker processes for the chains. None uses one per + chain, up to the number of CPUs; 1 runs them one after another in + this process. + :param verbose: Print the stages as they happen. + :returns: A :py:class:`SampleOutcome`. + + The C++ likelihood already spreads each evaluation over the cores + with OpenMP, but past two or three threads it gains little. Measured + on the AChR records of Epstein et al. (2016): 57 ms on one thread, + 36 ms on two, 29 ms on four. So parallel chains are run as separate + processes, each limited to its share of the logical CPUs. + + What that buys depends on the machine. On a 4-core, 8-thread laptop, + four CH82 chains ran 1.3-1.6 times faster in four processes than one + after another in one. Fewer threads per process (one each) was + slower there, because hyperthreading helps. More physical cores, or a + likelihood that costs more per evaluation than CH82's 1-2 ms, gain + more. + + Worker processes are *spawned*. As with any spawned process on + Windows, a script that calls this with ``processes > 1`` must do so + under ``if __name__ == '__main__':``. The ``hjcfit`` command and + notebooks need nothing. + """ + import time + from .mcmc import gaussian_approximation, mwg_sample + + started = time.perf_counter() + m = spec.mcmc + records, mec, post = _posterior(spec) + names = tuple(mec.get_free_parameter_names()) + notes = [] + fit = pilot = approximation = None + + if m.start == 'fit': + if verbose: + print("fitting, to start the chains at the maximum ...") + fit = run(spec) + start = np.asarray(fit.result.free_values, dtype=float) + start_from = 'fit' + if fit.limited: + notes.append("the fit ended with {0} on a limit, which is also " + "the prior's edge".format(", ".join( + name for name, *_ in fit.limited))) + else: + if verbose: + print("running the pilot from the guess ...") + pilot = mwg_sample(post, np.asarray(mec.theta(), dtype=float), + n=m.pilot_n, burnin=m.pilot_n // 2, + rng=m.seed - 1) + start, _ = pilot.mode() + start_from = 'pilot' + if not np.isfinite(post(start)): + raise RunnerError("the posterior is zero at the starting point") + + covariance = None + if m.covariance == 'hessian': + try: + approximation = gaussian_approximation(post, start) + covariance = approximation.covariance + if approximation.relative_error > 1e-2: + notes.append("the Hessian at the start is poorly determined " + "(relative error {0:.1g}); a direction the " + "records barely constrain".format( + approximation.relative_error)) + except ValueError as error: + notes.append("no Hessian at the start ({0}); the chains begin " + "with the small isotropic step instead".format( + str(error).split(';')[0].split('.')[0])) + + starts = _starts(post, start, covariance, m.chains, m.seed) + workers = 1 if processes == 1 else min(m.chains, processes + or os.cpu_count() or 1) + if verbose: + print("{0} chain{1} of {2} iterations in {3} process{4} ...".format( + m.chains, "" if m.chains == 1 else "s", m.n, workers, + "" if workers == 1 else "es")) + if workers == 1: + chains = [_run_chain(spec, i, s, covariance) + for i, s in enumerate(starts)] + else: + chains = _parallel_chains(spec, starts, covariance, workers) + + return SampleOutcome( + spec=spec, records=records, names=names, start=start, + start_from=start_from, chains=chains, + summary=_chain_summary(names, chains), fit=fit, pilot=pilot, + approximation=approximation, notes=notes, + seconds=time.perf_counter() - started, provenance=provenance()) + + +def _parallel_chains(spec, starts, covariance, workers): + """ Run the chains in *workers* spawned processes, each limited to its + share of the OpenMP threads. The limit is set in the environment the + processes inherit, before any of them loads the likelihood library, + and restored afterwards. """ + from concurrent.futures import ProcessPoolExecutor + import multiprocessing + + threads = max(1, (os.cpu_count() or 1) // workers) + previous = os.environ.get('OMP_NUM_THREADS') + os.environ['OMP_NUM_THREADS'] = str(threads) + try: + context = multiprocessing.get_context('spawn') + with ProcessPoolExecutor(max_workers=workers, + mp_context=context) as pool: + futures = [pool.submit(_run_chain, spec, i, s, covariance) + for i, s in enumerate(starts)] + return [f.result() for f in futures] + finally: + if previous is None: + os.environ.pop('OMP_NUM_THREADS', None) + else: + os.environ['OMP_NUM_THREADS'] = previous + + +def samples_as_dict(outcome): + """ A sampling outcome as JSON-ready data: the specification, the start, + the summary and the provenance. The chains themselves go in ``.npz`` + files beside it (:py:func:`write_samples`). """ + approx = outcome.approximation + return { + 'hjcfit_samples': 1, + 'spec': outcome.spec.as_dict(), + 'free_names': list(outcome.names), + 'start': {'from': outcome.start_from, + 'rates': [float(v) for v in outcome.start]}, + 'fit': None if outcome.fit is None else result_as_dict( + outcome.fit)['fit'], + 'hessian_sd': None if approx is None else [float(v) + for v in approx.sd], + 'summary': outcome.summary, + 'chains': [{'n': c.n, 'burnin': c.burnin, 'seconds': c.seconds, + 'acceptance': np.atleast_1d(c.acceptance_rate()).tolist(), + 'nfailures': c.nfailures} for c in outcome.chains], + 'notes': list(outcome.notes), + 'seconds': outcome.seconds, + 'provenance': outcome.provenance, + } + + +def write_samples(prefix, outcome): + """ Write ``.json`` (:py:func:`samples_as_dict`) and one + ``_chain.npz`` per chain + (:py:meth:`~HJCFIT.likelihood.mcmc.Chain.save`). + + :returns: The paths written. + """ + prefix = str(prefix) + if prefix.endswith('.json'): + prefix = prefix[:-5] + paths = [prefix + '.json'] + with open(paths[0], 'w', encoding='utf-8', newline='\n') as handle: + json.dump(samples_as_dict(outcome), handle, indent=2) + handle.write("\n") + for i, chain in enumerate(outcome.chains): + path = '{0}_chain{1}.npz'.format(prefix, i) + chain.save(path) + paths.append(path) + return paths diff --git a/tests/test_dcio_integration.py b/tests/test_dcio_integration.py index 22c46fa..307da9a 100644 --- a/tests/test_dcio_integration.py +++ b/tests/test_dcio_integration.py @@ -272,3 +272,42 @@ def test_both_searches_find_the_same_maximum(self): reached.append(result.log10_likelihood) assert reached[0] == pytest.approx(reached[1], abs=1e-2), reached + + +class TestPosteriorWithARealMechanism: + """LogPosterior against scalcs' CH82 and the shipped record. + + test_mcmc.py proves the posterior needs no mechanism library. This checks + that the prior scalcs' own rate limits give is the one intended, and that + inside it the posterior is the fitter's likelihood plus a constant. + """ + + @pytest.fixture() + def posterior(self): + from scalcs.samples import samples + + from HJCFIT.likelihood.fitting import Record + from HJCFIT.likelihood.mcmc import LogPosterior + + bursts = HJCFIT.read_idealized_bursts("CH82", tau=1e-4, tcrit=4e-3) + record = Record(conc=100e-9, tres=1e-4, tcrit=4e-3, + groups=tuple(tuple(b) for b in bursts)) + return LogPosterior(samples.CH82(), [record]) + + def test_the_default_prior_is_the_mechanism_limits(self, posterior): + prior = posterior.prior + assert prior.names == posterior.names + free = [r for r in posterior.mec.Rates if r.is_free] + np.testing.assert_array_equal(prior.lower, + [r.limits[0][0] for r in free]) + np.testing.assert_array_equal(prior.upper, + [r.limits[0][1] for r in free]) + + def test_inside_the_prior_it_is_the_golden_likelihood(self, posterior): + """The same golden value as the fitter's, at the shipped guess.""" + theta = np.asarray(posterior.mec.theta(), dtype=float) + expected = (2286.9746 * np.log(10.0) + - np.sum(np.log(posterior.prior.upper + - posterior.prior.lower))) + assert posterior(theta) == pytest.approx(expected, abs=1e-2) + assert posterior.nfailures == 0 diff --git a/tests/test_mcmc.py b/tests/test_mcmc.py new file mode 100644 index 0000000..a9f57d2 --- /dev/null +++ b/tests/test_mcmc.py @@ -0,0 +1,204 @@ +"""Bayesian layer: priors and the log posterior. + +As in test_fitting.py, nothing here imports scalcs. The posterior is +duck-typed on a mechanism, so the fake one from test_fitting.py is enough; +the real-mechanism checks live in test_dcio_integration.py. +""" + +import subprocess +import sys + +import numpy as np +import pytest + +from HJCFIT.likelihood.fitting import SOLVER_OPTIONS, HJCFitter, Record +from HJCFIT.likelihood.mcmc import LogPosterior, LogUniformPrior, UniformPrior + +from test_fitting import FakeMechanism + + +@pytest.fixture() +def record(): + """A short alternating record from the fake mechanism's rough timescales, + as in test_fitting.py.""" + rng = np.random.default_rng(4) + n = 41 + intervals = np.empty(n) + intervals[0::2] = rng.exponential(1.0 / 1000.0, len(intervals[0::2])) + intervals[1::2] = rng.exponential(1.0 / 100.0, len(intervals[1::2])) + return Record(conc=0.0, groups=(tuple(intervals),), tres=0.0) + + +LIMITS = [[1.0, 1e5]] + + +def test_mcmc_does_not_import_scalcs(): + out = subprocess.run( + [sys.executable, "-c", + "import sys, HJCFIT.likelihood.mcmc; print('scalcs' in sys.modules)"], + capture_output=True, text=True, check=True) + assert out.stdout.strip() == "False" + + +# -------------------------------------------------------------------------- +# Priors +# -------------------------------------------------------------------------- + +class TestUniformPrior: + + def test_density_inside_is_the_product_of_reciprocal_widths(self): + prior = UniformPrior([0.0, 10.0], [2.0, 30.0]) + assert prior.logpdf([1.0, 20.0]) == pytest.approx(-np.log(2 * 20)) + + def test_boundaries_are_inside(self): + prior = UniformPrior([0.0, 10.0], [2.0, 30.0]) + assert np.isfinite(prior.logpdf([0.0, 30.0])) + + def test_outside_is_minus_infinity(self): + prior = UniformPrior([0.0, 10.0], [2.0, 30.0]) + assert prior.logpdf([1.0, 31.0]) == -np.inf + assert prior.logpdf([-1e-9, 20.0]) == -np.inf + + def test_from_mechanism_reads_the_free_rates_limits(self): + prior = UniformPrior.from_mechanism(FakeMechanism(limits=LIMITS)) + assert prior.names == ("alpha", "beta") + np.testing.assert_array_equal(prior.lower, [1.0, 1.0]) + np.testing.assert_array_equal(prior.upper, [1e5, 1e5]) + + def test_from_mechanism_needs_limits(self): + with pytest.raises(ValueError, match="no limits"): + UniformPrior.from_mechanism(FakeMechanism()) + + def test_rejects_inverted_or_infinite_bounds(self): + with pytest.raises(ValueError, match="below"): + UniformPrior([1.0], [1.0]) + with pytest.raises(ValueError, match="finite"): + UniformPrior([0.0], [np.inf]) + + def test_rejects_the_wrong_number_of_rates(self): + with pytest.raises(ValueError, match="expected 2"): + UniformPrior([0.0, 0.0], [1.0, 1.0]).logpdf([0.5]) + + def test_samples_stay_inside(self): + prior = UniformPrior([0.0, 10.0], [2.0, 30.0]) + draws = prior.sample(np.random.default_rng(1), size=1000) + assert draws.shape == (1000, 2) + assert all(prior.contains(d) for d in draws) + + +class TestLogUniformPrior: + + def test_integrates_to_one(self): + prior = LogUniformPrior([1e-2], [1e6]) + x = np.logspace(-2, 6, 200001) + density = np.exp([prior.logpdf([v]) for v in x]) + assert np.trapezoid(density, x) == pytest.approx(1.0, rel=1e-4) + + def test_every_decade_equally_likely(self): + prior = LogUniformPrior([1.0], [1e4]) + draws = prior.sample(np.random.default_rng(2), size=40000)[:, 0] + counts = np.histogram(np.log10(draws), bins=4, range=(0, 4))[0] + assert np.all(np.abs(counts / 10000 - 1) < 0.05) + + def test_needs_positive_bounds(self): + with pytest.raises(ValueError, match="positive"): + LogUniformPrior([0.0], [1.0]) + + +# -------------------------------------------------------------------------- +# Solver options on the fitter +# -------------------------------------------------------------------------- + +class TestSolverOptions: + + def test_reach_every_likelihood(self, record): + fitter = HJCFitter(FakeMechanism(), [record, record], + solver=dict(nmax=2, xtol=1e-12, rtol=1e-11, + itermax=50)) + for lik in fitter.likelihoods: + assert (lik.nmax, lik.xtol, lik.rtol, lik.itermax) == ( + 2, 1e-12, 1e-11, 50) + + def test_defaults_are_left_alone(self, record): + lik = HJCFitter(FakeMechanism(), [record]).likelihoods[0] + assert (lik.nmax, lik.xtol, lik.rtol) == (3, 1e-10, 1e-10) + + def test_unknown_option_is_refused(self, record): + with pytest.raises(ValueError, match="unknown solver option"): + HJCFitter(FakeMechanism(), [record], solver=dict(tol=1e-9)) + + def test_the_names_are_what_log10likelihood_takes(self): + assert SOLVER_OPTIONS == ('nmax', 'xtol', 'rtol', 'itermax', + 'lower_bound', 'upper_bound') + + +# -------------------------------------------------------------------------- +# LogPosterior +# -------------------------------------------------------------------------- + +class TestLogPosterior: + + def test_is_prior_plus_likelihood(self, record): + mec = FakeMechanism(limits=LIMITS) + post = LogPosterior(mec, [record]) + rates = np.array([900.0, 120.0]) + fitter = HJCFitter(FakeMechanism(limits=LIMITS), [record]) + expected = fitter.ln_likelihood(np.log(rates)) - 2 * np.log(1e5 - 1.0) + assert post(rates) == pytest.approx(expected, rel=1e-12) + + def test_names_and_size_come_from_the_mechanism(self, record): + post = LogPosterior(FakeMechanism(limits=LIMITS), [record]) + assert post.names == ("alpha", "beta") and post.k == 2 + + def test_outside_the_prior_skips_the_likelihood(self, record): + post = LogPosterior(FakeMechanism(limits=LIMITS), [record]) + assert post([2e5, 100.0]) == -np.inf + assert post.nevals == 0 + + def test_the_likelihood_is_never_clipped(self, record): + """The fitter would reset alpha = 3000 to its limit, 2000. The + posterior's likelihood must see 3000 itself.""" + limits = [[1.0, 2000.0]] + post = LogPosterior(FakeMechanism(limits=limits), [record]) + unclipped = HJCFitter(FakeMechanism(), [record]) + clipped = HJCFitter(FakeMechanism(limits=limits), [record]) + rates = np.array([3000.0, 100.0]) + value = post.log_likelihood(rates) + assert value == pytest.approx( + unclipped.ln_likelihood(np.log(rates)), rel=1e-12) + assert value != pytest.approx( + clipped.ln_likelihood(np.log(rates)), rel=1e-6) + + def test_a_failure_is_minus_infinity_and_counted(self, record): + post = LogPosterior(FakeMechanism(limits=LIMITS), [record], + prior=UniformPrior([-np.float64(1e9)] * 2, + [np.float64(1e9)] * 2)) + assert post.log_likelihood([np.nan, 100.0]) == -np.inf + assert (post.nevals, post.nfailures) == (1, 1) + + def test_a_zero_rate_does_not_raise(self, record): + """A rate of zero gives a singular Q; the posterior must return + -inf or a number, never an exception.""" + post = LogPosterior(FakeMechanism(limits=[[0.0, 1e5]]), [record]) + value = post([0.0, 100.0]) + assert value == -np.inf or np.isfinite(value) + + def test_solver_options_reach_the_likelihood(self, record): + post = LogPosterior(FakeMechanism(limits=LIMITS), [record], + solver=dict(nmax=2)) + assert post.fitter.likelihoods[0].nmax == 2 + + def test_explicit_prior_must_match_the_mechanism(self, record): + with pytest.raises(ValueError, match="free rates"): + LogPosterior(FakeMechanism(limits=LIMITS), [record], + prior=UniformPrior([0.0], [1.0])) + with pytest.raises(ValueError, match="in that order"): + LogPosterior(FakeMechanism(limits=LIMITS), [record], + prior=UniformPrior([0, 0], [1, 1], + names=("beta", "alpha"))) + + def test_wrong_number_of_rates_is_an_error_not_minus_infinity( + self, record): + post = LogPosterior(FakeMechanism(limits=LIMITS), [record]) + with pytest.raises(ValueError, match="expected 2"): + post([100.0]) diff --git a/tests/test_mcmc_hessian.py b/tests/test_mcmc_hessian.py new file mode 100644 index 0000000..67645f8 --- /dev/null +++ b/tests/test_mcmc_hessian.py @@ -0,0 +1,170 @@ +"""The Hessian and the Gaussian approximation at a mode. + +Checked on functions whose second derivatives are known exactly. The +scaled case matters most: rate constants span eight decades, so one +parameter may be ~1 and another ~1e8 in the same Hessian. +""" + +import numpy as np +import pytest + +from HJCFIT.likelihood.mcmc import ( + GaussianApproximation, LogPosterior, gaussian_approximation, hessian) + + +def gaussian_log_density(mean, cov): + precision = np.linalg.inv(cov) + + def f(x): + d = np.asarray(x) - mean + return -0.5 * d @ precision @ d + return f + + +def scaled_cov(sd, rho): + sd = np.asarray(sd, dtype=float) + corr = np.full((len(sd), len(sd)), rho) + np.fill_diagonal(corr, 1.0) + return corr * np.outer(sd, sd) + + +class TestHessian: + + def test_exact_on_a_quadratic(self): + cov = scaled_cov([1.0, 2.0, 0.5], 0.6) + f = gaussian_log_density(np.array([3.0, 5.0, 7.0]), cov) + h, _ = hessian(f, [3.0, 5.0, 7.0]) + np.testing.assert_allclose(h, -np.linalg.inv(cov), rtol=1e-7) + + def test_fourth_order_on_a_curved_function(self): + """f = x0^4 + x0 x1^2 + exp(x1), whose Hessian is known exactly. + Richardson extrapolation should do much better than the plain + central difference at the same step.""" + def f(x): + return x[0] ** 4 + x[0] * x[1] ** 2 + np.exp(x[1]) + x = np.array([1.3, 0.7]) + exact = np.array([[12 * x[0] ** 2, 2 * x[1]], + [2 * x[1], 2 * x[0] + np.exp(x[1])]]) + h, err = hessian(f, x, rel_step=1e-2) + np.testing.assert_allclose(h, exact, rtol=1e-7) + assert np.all(err < 1e-2 * np.abs(exact)) + + def test_rates_eight_decades_apart(self): + """One rate ~1, one ~1e8, correlated. Relative steps handle both.""" + mean = np.array([2.0, 4e7, 5e4]) + cov = scaled_cov([0.2, 4e6, 3e3], 0.9) + h, _ = hessian(gaussian_log_density(mean, cov), mean) + np.testing.assert_allclose(h, -np.linalg.inv(cov), rtol=1e-6) + + def test_absolute_steps_where_a_coordinate_is_zero(self): + f = gaussian_log_density(np.zeros(2), np.eye(2)) + with pytest.raises(ValueError, match="absolute steps"): + hessian(f, [0.0, 0.0]) + h, _ = hessian(f, [0.0, 0.0], steps=1e-3) + np.testing.assert_allclose(h, -np.eye(2), rtol=1e-8) + + def test_a_step_across_the_boundary_is_an_error(self): + """A difference across a prior edge gives -inf, never a number.""" + def f(x): + return -np.sum(x ** 2) if np.all(x < 1.0) else -np.inf + with pytest.raises(ValueError, match="parameter"): + hessian(f, [0.9995, 0.5]) + + def test_evaluations(self): + calls = [0] + + def f(x): + calls[0] += 1 + return -np.sum(np.asarray(x) ** 2) + hessian(f, np.ones(4)) + k = 4 + assert calls[0] == 1 + 2 * (2 * k + 4 * k * (k - 1) // 2) + + +class TestGaussianApproximation: + + def test_recovers_the_covariance_of_a_gaussian(self): + mean = np.array([2.0, 4e7, 5e4]) + cov = scaled_cov([0.2, 4e6, 3e3], 0.9) + g = gaussian_approximation(gaussian_log_density(mean, cov), mean) + np.testing.assert_allclose(g.covariance, cov, rtol=1e-6) + np.testing.assert_allclose(g.sd, [0.2, 4e6, 3e3], rtol=1e-6) + np.testing.assert_allclose(g.correlation[0, 1], 0.9, rtol=1e-6) + np.testing.assert_array_equal(g.mode, mean) + assert g.log_density == pytest.approx(0.0) + + def test_covariance_is_symmetric_positive_definite(self): + cov = scaled_cov([1.0, 3.0, 0.1, 50.0], 0.95) + g = gaussian_approximation( + gaussian_log_density(np.array([1.0, 2.0, 3.0, 4.0]), cov), + [1.0, 2.0, 3.0, 4.0]) + np.testing.assert_array_equal(g.covariance, g.covariance.T) + assert np.all(np.linalg.eigvalsh(g.covariance) > 0) + + def test_refuses_a_point_that_is_not_a_maximum(self): + def saddle(x): + return -x[0] ** 2 + x[1] ** 2 + with pytest.raises(ValueError, match="not a maximum"): + gaussian_approximation(saddle, [1.0, 1.0]) + + def test_marginal_pdf_integrates_to_one(self): + g = gaussian_approximation( + gaussian_log_density(np.array([10.0]), np.array([[4.0]])), [10.0]) + x = np.linspace(-10, 30, 20001) + assert np.trapezoid(g.marginal_pdf(0, x), x) == pytest.approx(1.0) + + def test_names_come_from_the_density(self): + f = gaussian_log_density(np.array([1.0, 2.0]), np.eye(2)) + f_named = lambda x: f(x) # noqa: E731 + f_named.names = ("alpha", "beta") + g = gaussian_approximation(f_named, [1.0, 2.0]) + assert g.names == ("alpha", "beta") + assert g.marginal_pdf("beta", 2.0) == pytest.approx( + 1 / np.sqrt(2 * np.pi)) + + def test_samples_have_its_moments(self): + cov = scaled_cov([1.0, 2.0], 0.8) + g = gaussian_approximation( + gaussian_log_density(np.array([5.0, 6.0]), cov), [5.0, 6.0]) + draws = g.sample(np.random.default_rng(1), size=40000) + np.testing.assert_allclose(np.cov(draws.T), cov, rtol=0.05) + + def test_is_a_dataclass_worth_keeping(self): + g = gaussian_approximation( + gaussian_log_density(np.array([5.0]), np.eye(1)), [5.0]) + assert isinstance(g, GaussianApproximation) + assert g.hessian_error.shape == (1, 1) + + +def test_on_a_log_posterior(): + """With the fake two-state mechanism: the approximation at a fitted + maximum is a proper covariance.""" + from HJCFIT.likelihood.fitting import HJCFitter, Record + from test_fitting import FakeMechanism + + rng = np.random.default_rng(4) + intervals = np.empty(401) + intervals[0::2] = rng.exponential(1.0 / 1000.0, 201) + intervals[1::2] = rng.exponential(1.0 / 100.0, 200) + record = Record(conc=0.0, groups=(tuple(intervals),), tres=0.0) + fit = HJCFitter(FakeMechanism(), [record]).fit(search="scipy") + post = LogPosterior(FakeMechanism(limits=[[1.0, 1e5]]), [record]) + g = gaussian_approximation(post, fit.free_values) + assert g.names == ("alpha", "beta") + assert np.all(g.sd > 0) + # For an exponential rate from n intervals, the standard error is about + # rate / sqrt(n); here n = 200 openings and 200 shuttings. + np.testing.assert_allclose(g.sd / fit.free_values, 1 / np.sqrt(200), + rtol=0.25) + + +def test_relative_error_reports_how_well_the_hessian_is_determined(): + """-cosh(x - 2) has its maximum at 2 with H = -1. A huge step leaves a + visible truncation error; the default step leaves almost none.""" + def f(x): + return -np.sum(np.cosh(np.asarray(x) - 2.0)) + coarse = gaussian_approximation(f, [2.0, 2.0], rel_step=0.3) + fine = gaussian_approximation(f, [2.0, 2.0]) + assert fine.relative_error < 1e-6 + assert coarse.relative_error > 100 * fine.relative_error + np.testing.assert_allclose(fine.sd, 1.0, rtol=1e-6) diff --git a/tests/test_mcmc_samplers.py b/tests/test_mcmc_samplers.py new file mode 100644 index 0000000..b289762 --- /dev/null +++ b/tests/test_mcmc_samplers.py @@ -0,0 +1,340 @@ +"""Samplers, chains and diagnostics, on targets whose answers are known. + +Nothing here computes an HJC likelihood. The samplers take any log density, +so they are checked against distributions with known moments: + +* a correlated Gaussian, the shape of the alpha2-beta2 ridge that + motivates the adaptive sampler; +* a log-normal, which the log-space walk samples correctly only if it + applies the Jacobian of the transformation; +* a box, whose edges return -inf. + +Every run is seeded, so these are deterministic. The tolerances are several +Monte Carlo standard errors wide, so that a different seed would pass too. +""" + +import numpy as np +import pytest + +from HJCFIT.likelihood.mcmc import ( + Chain, adaptive_sample, autocorrelation, effective_sample_size, + mwg_sample, significant_lags) + + +# -------------------------------------------------------------------------- +# Targets +# -------------------------------------------------------------------------- + +class Gaussian: + """Correlated 2-d Gaussian, well away from zero so both walks apply.""" + + names = ("a", "b") + + def __init__(self, mean=(10.0, 20.0), sd=(1.0, 2.0), rho=0.9): + self.mean = np.array(mean) + sd = np.array(sd) + self.cov = np.array([[sd[0] ** 2, rho * sd[0] * sd[1]], + [rho * sd[0] * sd[1], sd[1] ** 2]]) + self.precision = np.linalg.inv(self.cov) + self.calls = 0 + + def __call__(self, x): + self.calls += 1 + d = np.asarray(x) - self.mean + return -0.5 * d @ self.precision @ d + + +class LogNormal: + """Independent log-normals: ln(theta) ~ N(mu, sigma**2). + + Written as a density over theta, 1/theta included, which is what a + posterior over rates is. Walking ln(theta) without the Jacobian would + sample N(mu - sigma**2, sigma**2) instead. + """ + + def __init__(self, mu=(0.0, 2.0), sigma=(0.5, 0.3)): + self.mu = np.array(mu) + self.sigma = np.array(sigma) + + def __call__(self, theta): + theta = np.asarray(theta) + if np.any(theta <= 0): + return -np.inf + y = np.log(theta) + return float(np.sum(-0.5 * ((y - self.mu) / self.sigma) ** 2 - y)) + + +def box(x): + x = np.asarray(x) + return 0.0 if np.all((x >= 0.0) & (x <= 1.0)) else -np.inf + + +# -------------------------------------------------------------------------- +# Metropolis-within-Gibbs +# -------------------------------------------------------------------------- + +class TestMWG: + + def test_log_space_walk_applies_the_jacobian(self): + target = LogNormal() + chain = mwg_sample(target, np.exp(target.mu), n=20000, burnin=2000, + rng=1) + y = np.log(chain.kept()) + np.testing.assert_allclose(y.mean(axis=0), target.mu, atol=0.03) + np.testing.assert_allclose(y.std(axis=0), target.sigma, rtol=0.06) + + def test_rate_space_walk_recovers_a_gaussian(self): + target = Gaussian() + chain = mwg_sample(target, target.mean, n=20000, burnin=2000, rng=2, + log_space=False) + kept = chain.kept() + np.testing.assert_allclose(kept.mean(axis=0), target.mean, atol=0.15) + np.testing.assert_allclose(np.cov(kept.T), target.cov, rtol=0.15) + + def test_tuning_shrinks_an_oversized_step(self): + target = Gaussian() + chain = mwg_sample(target, target.mean, n=3000, burnin=2000, rng=3, + log_space=False, initial_scale=50.0) + assert np.all(chain.scales[-1] < 10.0) + assert np.all(chain.scales[-1] == chain.scales[2000]) + rate = chain.acceptance_rate() + assert np.all((rate > 0.1) & (rate < 0.6)), rate + + def test_scales_are_frozen_after_burnin(self): + chain = mwg_sample(Gaussian(), (10.0, 20.0), n=600, burnin=300, + rng=4, log_space=False, initial_scale=20.0) + assert np.all(chain.scales[300:] == chain.scales[300]) + + def test_per_parameter_records(self): + chain = mwg_sample(Gaussian(), (10.0, 20.0), n=100, burnin=50, rng=5) + assert chain.accepted.shape == (100, 2) + assert chain.scales.shape == (100, 2) + assert chain.proposals.shape == (100, 2) + assert chain.acceptance_rate().shape == (2,) + + def test_one_evaluation_per_parameter_per_iteration(self): + target = Gaussian() + chain = mwg_sample(target, target.mean, n=100, burnin=0, rng=6) + assert chain.nevals == target.calls == 100 * 2 + 1 + + def test_log_posterior_is_of_the_rates_not_the_walk(self): + """The Jacobian is the walk's business; the chain reports the + target's own density at each sample.""" + target = LogNormal() + chain = mwg_sample(target, np.exp(target.mu), n=50, burnin=0, rng=7) + for x, lp in zip(chain.samples, chain.log_posterior): + assert lp == pytest.approx(target(x)) + + def test_needs_positive_rates_to_walk_logs(self): + with pytest.raises(ValueError, match="positive"): + mwg_sample(Gaussian(), (-1.0, 20.0), n=10, burnin=0) + + +# -------------------------------------------------------------------------- +# Adaptive Metropolis +# -------------------------------------------------------------------------- + +class TestAdaptive: + + @pytest.mark.parametrize("mixture", ["sum", "choice"]) + def test_recovers_a_correlated_gaussian(self, mixture): + target = Gaussian() + chain = adaptive_sample(target, target.mean, n=30000, burnin=5000, + rng=11, mixture=mixture) + kept = chain.kept() + np.testing.assert_allclose(kept.mean(axis=0), target.mean, atol=0.15) + np.testing.assert_allclose(np.cov(kept.T), target.cov, rtol=0.12) + + def test_learns_the_correlation(self): + """The point of the second stage: its proposals line up with the + posterior ridge (Epstein et al. 2016, Fig. 3 F).""" + target = Gaussian(rho=0.95) + chain = adaptive_sample(target, target.mean, n=10000, burnin=5000, + rng=12) + c = chain.proposal_covariance + assert c[0, 1] / np.sqrt(c[0, 0] * c[1, 1]) == pytest.approx( + 0.95, abs=0.03) + + def test_mixes_better_than_mwg_on_a_ridge(self): + """Table 2 of that paper, in miniature: on a correlated target the + adaptive sampler gives more effective samples per iteration.""" + target = Gaussian(rho=0.98) + mwg = mwg_sample(target, target.mean, n=10000, burnin=2000, rng=13, + log_space=False) + ada = adaptive_sample(target, target.mean, n=10000, burnin=2000, + rng=13) + ess_mwg = effective_sample_size(mwg.kept()[:, 0]) / len(mwg.kept()) + ess_ada = effective_sample_size(ada.kept()[:, 0]) / len(ada.kept()) + assert ess_ada > 2 * ess_mwg, (ess_ada, ess_mwg) + + def test_log_space_walk_applies_the_jacobian(self): + target = LogNormal() + chain = adaptive_sample(target, np.exp(target.mu), n=30000, + burnin=5000, rng=14, log_space=True) + y = np.log(chain.kept()) + np.testing.assert_allclose(y.mean(axis=0), target.mu, atol=0.03) + np.testing.assert_allclose(y.std(axis=0), target.sigma, rtol=0.06) + + def test_initial_covariance_shapes_the_first_steps(self): + """Before adapt_start the step is (initial_step / sqrt(k)) L0 z. With + L0 L0^T = diag(100, 1e-4), the steps in the first coordinate are + 10**3 times the size of those in the second.""" + chain = adaptive_sample(lambda x: 0.0, (0.0, 0.0), n=200, burnin=0, + rng=15, adapt_start=200, + initial_covariance=np.diag([100.0, 1e-4])) + steps = np.diff(chain.samples, axis=0) + ratio = steps[:, 0].std() / steps[:, 1].std() + assert ratio == pytest.approx(1e3, rel=0.2) + + def test_one_evaluation_per_iteration(self): + target = Gaussian() + chain = adaptive_sample(target, target.mean, n=300, burnin=100, + rng=16) + assert chain.nevals == target.calls == 301 + assert chain.accepted.shape == (300,) + assert chain.scales.shape == (300,) + + def test_rejects_an_unknown_mixture(self): + with pytest.raises(ValueError, match="mixture"): + adaptive_sample(Gaussian(), (10.0, 20.0), n=10, burnin=0, + mixture="both") + + +# -------------------------------------------------------------------------- +# Both samplers +# -------------------------------------------------------------------------- + +SAMPLERS = [ + pytest.param(lambda f, x0, **kw: mwg_sample(f, x0, log_space=False, **kw), + id="mwg"), + pytest.param(adaptive_sample, id="adaptive"), +] + + +@pytest.mark.parametrize("sample", SAMPLERS) +class TestBothSamplers: + + def test_a_seed_reproduces_the_chain(self, sample): + a = sample(Gaussian(), (10.0, 20.0), n=500, burnin=200, rng=21) + b = sample(Gaussian(), (10.0, 20.0), n=500, burnin=200, rng=21) + np.testing.assert_array_equal(a.samples, b.samples) + np.testing.assert_array_equal(a.proposals, b.proposals) + + def test_never_leaves_the_support(self, sample): + chain = sample(box, (0.5, 0.5), n=3000, burnin=500, rng=22) + assert np.all((chain.samples >= 0) & (chain.samples <= 1)) + assert np.any((chain.proposals < 0) | (chain.proposals > 1)), ( + "the test needs proposals outside the box to mean anything") + np.testing.assert_allclose(chain.kept().mean(axis=0), 0.5, atol=0.05) + + def test_refuses_a_start_with_zero_density(self, sample): + with pytest.raises(ValueError, match="starting point"): + sample(box, (2.0, 0.5), n=10, burnin=0) + + def test_nan_is_a_rejection(self, sample): + def target(x): + return np.nan if x[0] > 0.6 else box(x) + chain = sample(target, (0.5, 0.5), n=500, burnin=0, rng=23) + assert np.all(chain.samples[:, 0] <= 0.6) + + def test_names_come_from_the_target(self, sample): + chain = sample(Gaussian(), (10.0, 20.0), n=10, burnin=0, rng=24) + assert chain.names == ("a", "b") + + def test_burnin_cannot_exceed_the_run(self, sample): + with pytest.raises(ValueError, match="burnin"): + sample(Gaussian(), (10.0, 20.0), n=10, burnin=11) + + def test_callback_sees_the_chain_so_far(self, sample): + seen = [] + sample(Gaussian(), (10.0, 20.0), n=250, burnin=100, rng=25, + callback=lambda i, c: seen.append((i, c.n, c.burnin)), + callback_every=100) + assert seen == [(100, 100, 100), (200, 200, 100)] + + +# -------------------------------------------------------------------------- +# Chain +# -------------------------------------------------------------------------- + +class TestChain: + + @pytest.fixture() + def chain(self): + return adaptive_sample(Gaussian(), (10.0, 20.0), n=400, burnin=150, + rng=31) + + def test_kept_drops_the_burnin(self, chain): + assert chain.kept().shape == (250, 2) + np.testing.assert_array_equal(chain.kept(), chain.samples[150:]) + + def test_mode_is_the_best_sample(self, chain): + x, lp = chain.mode() + assert lp == chain.log_posterior.max() + assert Gaussian()(x) == pytest.approx(lp) + + def test_save_and_load_round_trip(self, chain, tmp_path): + path = tmp_path / "chain.npz" + chain.save(path) + back = Chain.load(path) + for name in ("samples", "log_posterior", "proposals", "accepted", + "scales", "proposal_covariance"): + np.testing.assert_array_equal(getattr(back, name), + getattr(chain, name)) + for name in ("burnin", "names", "sampler", "settings", "nevals", + "nfailures"): + assert getattr(back, name) == getattr(chain, name), name + + +# -------------------------------------------------------------------------- +# Diagnostics +# -------------------------------------------------------------------------- + +def ar1(phi, n, seed): + rng = np.random.default_rng(seed) + x = np.empty(n) + x[0] = rng.standard_normal() / np.sqrt(1 - phi ** 2) + for i in range(1, n): + x[i] = phi * x[i - 1] + rng.standard_normal() + return x + + +class TestDiagnostics: + + def test_autocorrelation_starts_at_one(self): + rho = autocorrelation(ar1(0.5, 1000, 1), 5) + assert rho[0] == 1.0 and len(rho) == 6 + + def test_autocorrelation_of_ar1_decays_geometrically(self): + rho = autocorrelation(ar1(0.8, 200000, 2), 3) + np.testing.assert_allclose(rho, 0.8 ** np.arange(4), atol=0.01) + + def test_constant_series_is_refused(self): + with pytest.raises(ValueError, match="constant"): + autocorrelation(np.ones(10), 2) + + def test_independent_draws_are_worth_their_number(self): + x = np.random.default_rng(3).standard_normal(20000) + assert effective_sample_size(x) == pytest.approx(20000, rel=0.1) + + def test_ar1_matches_theory(self): + """For AR(1), ESS / N = (1 - phi) / (1 + phi).""" + phi = 0.9 + x = ar1(phi, 100000, 4) + ess = effective_sample_size(x, max_lag=200) + assert ess / len(x) == pytest.approx((1 - phi) / (1 + phi), rel=0.15) + + def test_never_more_than_n(self): + """An anticorrelated series would otherwise claim more samples than + it has.""" + x = ar1(-0.5, 20000, 5) + assert effective_sample_size(x) <= len(x) + + def test_significant_lags_grow_with_correlation(self): + assert significant_lags(ar1(0.0, 20000, 6)) <= 3 + assert significant_lags(ar1(0.95, 20000, 7)) > 20 + + def test_one_value_per_column(self): + x = np.column_stack([ar1(0.0, 5000, 8), ar1(0.9, 5000, 9)]) + ess = effective_sample_size(x) + assert ess.shape == (2,) and ess[0] > 3 * ess[1] diff --git a/tests/test_sampling.py b/tests/test_sampling.py new file mode 100644 index 0000000..05c0303 --- /dev/null +++ b/tests/test_sampling.py @@ -0,0 +1,300 @@ +"""Sampling from a specification: the [likelihood] and [mcmc] tables, split- +R-hat, the runner's sample() and the ``hjcfit sample`` command. + +The specification tests need nothing beyond HJCFIT. The runner and command +tests need the [fitting] extra and skip without it, as in +test_runner_cli.py. They sample the CH82 sample record, which is quick, and +whose one concentration leaves one direction barely constrained. That makes +it a good test of how the run reports a poor Hessian and chains that have +not yet agreed. +""" + +import json + +import numpy as np +import pytest + +from HJCFIT.likelihood.fitspec import ( + TEMPLATE, FitSpec, LikelihoodSpec, MCMCSpec, SpecError, load_toml_bytes) +from HJCFIT.likelihood.mcmc import Chain, potential_scale_reduction + +CH82 = { + 'title': 'CH82 at 100 nM', + 'data': [{'record': 'CH82', 'conc': 1e-7, 'tres': 1e-4, 'tcrit': 4e-3}], + 'mechanism': {'sample': 'CH82', 'rates': {'beta1': 15.0, + 'beta2': 15000.0}}, +} + + +def spec_for(**changes): + return FitSpec.from_dict(dict(CH82, **changes)) + + +# -------------------------------------------------------------------------- +# [likelihood] +# -------------------------------------------------------------------------- + +class TestLikelihoodSpec: + + def test_its_names_are_the_fitters(self): + from HJCFIT.likelihood.fitspec import SOLVER_KEYS + from HJCFIT.likelihood.fitting import SOLVER_OPTIONS + assert SOLVER_KEYS == SOLVER_OPTIONS + + def test_only_what_is_given_reaches_the_solver(self): + spec = LikelihoodSpec.from_dict({'nmax': 2, 'xtol': 1e-12}) + assert spec.solver() == {'nmax': 2, 'xtol': 1e-12} + assert LikelihoodSpec().solver() == {} + + @pytest.mark.parametrize("table, message", [ + ({'nmax': 0}, "nmax"), + ({'nmax': 2.5}, "integer"), + ({'nmax': True}, "integer"), + ({'xtol': -1e-12}, "positive"), + ({'lower_bound': 0.0, 'upper_bound': -1.0}, "below"), + ({'tol': 1e-9}, "tol"), + ]) + def test_nonsense_is_refused(self, table, message): + with pytest.raises(SpecError, match=message): + LikelihoodSpec.from_dict(table) + + +# -------------------------------------------------------------------------- +# [mcmc] +# -------------------------------------------------------------------------- + +class TestMCMCSpec: + + def test_defaults(self): + m = MCMCSpec.from_dict({}) + assert (m.start, m.sampler, m.n, m.burnin, m.chains) == ( + 'fit', 'adaptive', 20000, 5000, 1) + assert m.as_dict() == {} + + def test_only_departures_are_written_back(self): + m = MCMCSpec.from_dict({'n': 5000, 'burnin': 1000, 'chains': 4}) + assert m.as_dict() == {'n': 5000, 'burnin': 1000, 'chains': 4} + + @pytest.mark.parametrize("table, message", [ + ({'start': 'mode'}, "start"), + ({'sampler': 'hmc'}, "sampler"), + ({'prior': 'flat'}, "prior"), + ({'covariance': 'fisher'}, "covariance"), + ({'n': 1000, 'burnin': 1000}, "keeps nothing"), + ({'chains': 0}, "chains"), + ({'n': 1e4}, "integer"), + ({'log_space': 1}, "true or false"), + ({'iterations': 100}, "iterations"), + ]) + def test_nonsense_is_refused(self, table, message): + with pytest.raises(SpecError, match=message): + MCMCSpec.from_dict(table) + + def test_a_whole_specification_round_trips(self): + spec = spec_for(likelihood={'nmax': 2, 'xtol': 1e-12}, + mcmc={'n': 5000, 'burnin': 1000, 'chains': 3, + 'start': 'guess', 'log_space': True}) + back = FitSpec.from_dict(load_toml_bytes(spec.to_toml().encode())) + assert back == spec + + def test_without_the_tables_nothing_is_written(self): + text = spec_for().to_toml() + assert '[likelihood]' not in text and '[mcmc]' not in text + + def test_str_says_how_it_would_sample(self): + text = str(spec_for(mcmc={'chains': 4}, + likelihood={'nmax': 2})) + assert 'sampling: 4 x 20000 adaptive' in text + assert 'likelihood: nmax = 2' in text + + def test_the_templates_example_sections_are_valid(self): + """The commented [likelihood] and [mcmc] blocks, uncommented.""" + lines, inside = [], False + for line in TEMPLATE.splitlines(): + if line in ('# [likelihood]', '# [mcmc]'): + inside = True + elif not line.strip(): + inside = False + if inside and line.startswith('# ') and ( + line.startswith('# [') or '=' in line): + line = line[2:] + lines.append(line) + text = "\n".join(lines) + assert '\n[mcmc]' in text and '\n[likelihood]' in text + spec = FitSpec.from_dict(load_toml_bytes(text.encode())) + assert spec.mcmc.chains == 4 + assert spec.likelihood.solver() == {'nmax': 2, 'xtol': 1e-12, + 'rtol': 1e-12} + + +# -------------------------------------------------------------------------- +# split-R-hat +# -------------------------------------------------------------------------- + +class TestPotentialScaleReduction: + + def test_chains_of_one_distribution_give_one(self): + rng = np.random.default_rng(0) + chains = [rng.standard_normal(4000) for _ in range(4)] + assert potential_scale_reduction(chains) == pytest.approx(1.0, + abs=0.01) + + def test_chains_that_disagree_give_more(self): + rng = np.random.default_rng(1) + chains = [rng.standard_normal(2000) + shift for shift in (0, 0, 3)] + assert potential_scale_reduction(chains) > 1.5 + + def test_a_single_chain_that_drifts_is_caught(self): + drift = np.linspace(0, 10, 4000) + np.random.default_rng(2).normal( + size=4000) + assert potential_scale_reduction([drift]) > 1.5 + + def test_too_short_is_an_error(self): + with pytest.raises(ValueError, match="four"): + potential_scale_reduction([np.ones(3)]) + + +# -------------------------------------------------------------------------- +# The runner and the command: need the [fitting] extra +# -------------------------------------------------------------------------- + +@pytest.fixture(scope='module') +def fitting_extra(): + pytest.importorskip('dcio', reason='needs the [fitting] extra') + pytest.importorskip('scalcs', reason='needs the [fitting] extra') + + +@pytest.fixture(scope='module') +def outcome(fitting_extra): + from HJCFIT.likelihood.runner import sample + spec = spec_for(mcmc={'n': 400, 'burnin': 100, 'chains': 2}) + return sample(spec, processes=1) + + +class TestSample: + + def test_chains_and_summary(self, outcome): + assert len(outcome.chains) == 2 + assert all(c.n == 400 and c.burnin == 100 for c in outcome.chains) + names = [row['name'] for row in outcome.summary] + assert names == list(outcome.names) and len(names) == 8 + for row in outcome.summary: + assert row['q2.5'] <= row['median'] <= row['q97.5'] + assert row['sd'] > 0 and row['ess'] > 0 + + def test_the_first_chain_starts_at_the_fit(self, outcome): + assert outcome.start_from == 'fit' + np.testing.assert_allclose(outcome.chains[0].samples[0], + outcome.fit.result.free_values, + rtol=0.2) + np.testing.assert_array_equal(outcome.start, + outcome.fit.result.free_values) + + def test_a_poor_hessian_is_reported_not_hidden(self, outcome): + """CH82 at one concentration has a direction the record barely + constrains; the run says so.""" + assert outcome.approximation is not None or outcome.notes + assert any('Hessian' in note for note in outcome.notes) + + def test_str_is_a_table_with_rhat(self, outcome): + text = str(outcome) + assert 'R-hat' in text and 'acceptance after burn-in' in text + + def test_written_and_read_back(self, outcome, tmp_path): + from HJCFIT.likelihood.runner import write_samples + paths = write_samples(tmp_path / 'post', outcome) + assert [p.rsplit('/', 1)[-1].rsplit('\\', 1)[-1] for p in paths] == [ + 'post.json', 'post_chain0.npz', 'post_chain1.npz'] + with open(paths[0]) as handle: + data = json.load(handle) + assert data['free_names'] == list(outcome.names) + assert data['spec']['mcmc'] == {'n': 400, 'burnin': 100, 'chains': 2} + assert data['provenance']['versions']['hjcfit'] + back = Chain.load(paths[1]) + np.testing.assert_array_equal(back.samples, outcome.chains[0].samples) + + def test_the_paper_scheme_starts_from_a_pilot(self, fitting_extra): + from HJCFIT.likelihood.runner import sample + spec = spec_for(mcmc={'start': 'guess', 'pilot_n': 60, 'n': 200, + 'burnin': 50, 'covariance': 'identity'}) + out = sample(spec, processes=1) + assert out.start_from == 'pilot' and out.fit is None + assert out.pilot.n == 60 + np.testing.assert_array_equal(out.start, out.pilot.mode()[0]) + + def test_the_mwg_sampler_can_run_the_chains(self, fitting_extra): + from HJCFIT.likelihood.runner import sample + spec = spec_for(mcmc={'sampler': 'mwg', 'n': 60, 'burnin': 20}) + out = sample(spec, processes=1) + assert out.chains[0].accepted.shape == (60, 8) + + def test_chains_in_parallel_processes(self, fitting_extra): + """Spawned processes, each rebuilding the posterior from the + specification; seeded, so the result does not depend on where a + chain ran.""" + from HJCFIT.likelihood.runner import sample + spec = spec_for(mcmc={'n': 150, 'burnin': 50, 'chains': 2}) + here = sample(spec, processes=1) + there = sample(spec, processes=2) + for a, b in zip(here.chains, there.chains): + np.testing.assert_allclose(a.samples, b.samples, rtol=1e-9) + + def test_the_likelihood_table_reaches_the_posterior(self, fitting_extra): + from HJCFIT.likelihood.runner import _posterior + _, _, post = _posterior(spec_for(likelihood={'nmax': 2, + 'xtol': 1e-12})) + lik = post.fitter.likelihoods[0] + assert (lik.nmax, lik.xtol) == (2, 1e-12) + + def test_and_the_fit(self, fitting_extra, monkeypatch): + from HJCFIT.likelihood import runner + seen = {} + original = runner.HJCFitter + + class Spy(original): + def __init__(self, *args, **kwargs): + seen.update(kwargs) + super().__init__(*args, **kwargs) + + monkeypatch.setattr(runner, 'HJCFitter', Spy) + runner.run(spec_for(likelihood={'nmax': 2}, + search={'maxfev': 50})) + assert seen['solver'] == {'nmax': 2} + + +class TestSampleCommand: + + def test_sample_writes_its_files(self, fitting_extra, tmp_path, capsys): + from HJCFIT.likelihood import cli + path = tmp_path / 'ch82.toml' + spec_for(mcmc={'n': 200, 'burnin': 50}).write_toml(str(path)) + status = cli.main(['sample', str(path), '-o', + str(tmp_path / 'post'), '--processes', '1', '-q']) + out = capsys.readouterr().out + assert status in (0, 2) # 2: chains this short may disagree + assert 'R-hat' in out + assert (tmp_path / 'post.json').exists() + assert (tmp_path / 'post_chain0.npz').exists() + + def test_chains_on_the_command_line_override_the_file( + self, fitting_extra, tmp_path, capsys): + from HJCFIT.likelihood import cli + path = tmp_path / 'ch82.toml' + spec_for(mcmc={'n': 120, 'burnin': 20}).write_toml(str(path)) + cli.main(['sample', str(path), '--chains', '2', '--processes', '1', + '-o', str(tmp_path / 'p'), '-q']) + assert (tmp_path / 'p_chain1.npz').exists() + + def test_zero_chains_is_one_line_on_stderr(self, fitting_extra, + tmp_path, capsys): + from HJCFIT.likelihood import cli + path = tmp_path / 'ch82.toml' + spec_for().write_toml(str(path)) + assert cli.main(['sample', str(path), '--chains', '0']) == 1 + assert capsys.readouterr().err.startswith('hjcfit: --chains') + + +def test_a_plain_fit_says_nothing_about_sampling(): + """Without an [mcmc] section, `hjcfit fit` and `check` should not talk + about sampling.""" + assert 'sampling' not in str(spec_for())