bayesmith 0.1.0__py3-none-any.whl

bayesmith/__init__.py

@@ -0,0 +1,173 @@
+"""bayesmith: a graph of operators is a Bayesian model.
+
+Deterministic operators propagate dependence; probabilistic ones contribute a
+conditional density. The graph's structure is what selects the inference
+method -- exact where a subgraph permits one, NUTS where it does not.
+"""
+
+from __future__ import annotations
+
+import importlib
+import importlib.metadata as _metadata
+from typing import Any
+
+from bayesmith.errors import (
+    BayesmithError,
+    ConvergenceError,
+    GraphError,
+    NotGaussian,
+    StructureError,
+    TraceError,
+)
+
+__all__ = [
+    # tracing
+    "trace",
+    "const",
+    "det",
+    "sample",
+    "observe",
+    "plate",
+    "NodeRef",
+    "PlateRef",
+    # graph
+    "Graph",
+    "Plate",
+    "Node",
+    "Const",
+    "Deterministic",
+    "Probabilistic",
+    # evaluation
+    "evaluate",
+    "log_joint",
+    # inference
+    "compile",
+    "Posterior",
+    "Estimate",
+    "to_numpyro",
+    "nuts",
+    "predict",
+    # exact
+    "linear_operator",
+    "check_linearity",
+    "wiener_solve",
+    "gcr_sample",
+    "condition_bound",
+    "iterative_gls",
+    "sigma_from_graph",
+    "noise_std_at",
+    "precision_at",
+    "fisher_information",
+    "parameter_covariance",
+    # diagnose
+    "identifiability",
+    "prior_sensitivity",
+    "JeffreysPrior",
+    # errors
+    "BayesmithError",
+    "GraphError",
+    "TraceError",
+    "StructureError",
+    "ConvergenceError",
+    "NotGaussian",
+]
+
+# Every public name above except the six error classes is resolved lazily,
+# on first attribute access, rather than imported here at module scope.
+# Importing eagerly would make `import bayesmith` load numpyro (hence jax) as
+# a side effect -- exactly the regression this module previously had: Python
+# always runs a package's __init__.py before any of its submodules, so even
+# `import bayesmith.errors` was dragging in the whole bridge, which broke the
+# stdlib-only contract errors.py documents for itself and which
+# test_errors_module_imports_no_heavy_dependency enforces. Only errors.py is
+# cheap and stdlib-only, so it alone is still imported eagerly above.
+#
+# name -> (owning submodule, attribute name within it)
+_LAZY_ATTRS: dict[str, tuple[str, str]] = {
+    "trace": ("bayesmith.graph.trace", "trace"),
+    "const": ("bayesmith.graph.trace", "const"),
+    "det": ("bayesmith.graph.trace", "det"),
+    "sample": ("bayesmith.graph.trace", "sample"),
+    "observe": ("bayesmith.graph.trace", "observe"),
+    "plate": ("bayesmith.graph.trace", "plate"),
+    "NodeRef": ("bayesmith.graph.trace", "NodeRef"),
+    "PlateRef": ("bayesmith.graph.trace", "PlateRef"),
+    "Graph": ("bayesmith.graph.graph", "Graph"),
+    "Plate": ("bayesmith.graph.graph", "Plate"),
+    "Node": ("bayesmith.graph.nodes", "Node"),
+    "Const": ("bayesmith.graph.nodes", "Const"),
+    "Deterministic": ("bayesmith.graph.nodes", "Deterministic"),
+    "Probabilistic": ("bayesmith.graph.nodes", "Probabilistic"),
+    "evaluate": ("bayesmith.graph.evaluate", "evaluate"),
+    "log_joint": ("bayesmith.graph.evaluate", "log_joint"),
+    "compile": ("bayesmith.dispatch.plan", "compile"),
+    "Posterior": ("bayesmith.dispatch.execute", "Posterior"),
+    "Estimate": ("bayesmith.dispatch.execute", "Estimate"),
+    "to_numpyro": ("bayesmith.bridge.numpyro_bridge", "to_numpyro"),
+    "nuts": ("bayesmith.bridge.numpyro_bridge", "nuts"),
+    "predict": ("bayesmith.bridge.numpyro_bridge", "predict"),
+    "linear_operator": ("bayesmith.exact.linearity", "linear_operator"),
+    "check_linearity": ("bayesmith.exact.linearity", "check_linearity"),
+    "wiener_solve": ("bayesmith.exact.solve", "wiener_solve"),
+    "gcr_sample": ("bayesmith.exact.solve", "gcr_sample"),
+    "condition_bound": ("bayesmith.exact.solve", "condition_bound"),
+    "iterative_gls": ("bayesmith.exact.gls", "iterative_gls"),
+    "sigma_from_graph": ("bayesmith.exact.gls", "sigma_from_graph"),
+    "noise_std_at": ("bayesmith.exact.gaussian", "noise_std_at"),
+    "precision_at": ("bayesmith.exact.gaussian", "precision_at"),
+    "fisher_information": ("bayesmith.exact.fisher", "fisher_information"),
+    "parameter_covariance": ("bayesmith.exact.fisher", "parameter_covariance"),
+    "identifiability": ("bayesmith.diagnose.identifiability", "identifiability"),
+    "prior_sensitivity": ("bayesmith.diagnose.sensitivity", "prior_sensitivity"),
+    "JeffreysPrior": ("bayesmith.diagnose.priors", "JeffreysPrior"),
+}
+
+# Subpackages reachable as `bayesmith.<name>` after a bare `import bayesmith`,
+# without eagerly importing any of them -- `bridge` in particular is what
+# pulls in numpyro, and `exact` reaches numpyro through `bridge` too (see
+# `gaussian.py`'s use of `numpyro.distributions`). `errors` is listed too for
+# __dir__'s sake even though the eager import above already binds it as a
+# real attribute, so __getattr__ is never actually consulted for it.
+#
+# `evidence` reaches jax through `sqrtinfo.py`'s module-scope import, so it is
+# listed here and NOT imported eagerly, for the same reason `exact` is not.
+# It was missing from this tuple for the whole of B11: the layer was complete,
+# dense-oracled and cross-checked, and `import bayesmith; bayesmith.evidence`
+# still raised AttributeError, so only an explicit `import bayesmith.evidence`
+# reached it. Both halves are pinned in `tests/test_public_api.py` -- that it
+# resolves, and that resolving it is what pulls jax in rather than importing
+# this package.
+_LAZY_SUBMODULES = (
+    "graph",
+    "bridge",
+    "exact",
+    "dispatch",
+    "evidence",
+    "diagnose",
+    "errors",
+)
+
+
+#: The installed distribution's version, READ from the installed metadata
+#: rather than written here. Two spellings of one version is the defect this
+#: package has spent the most effort repairing, and a release number is the
+#: worst candidate for a second copy: it changes on exactly the commit where
+#: everyone is busy doing something else.
+__version__ = _metadata.version("bayesmith")
+
+
+def __getattr__(name: str) -> Any:
+    if name in _LAZY_ATTRS:
+        module_name, attr_name = _LAZY_ATTRS[name]
+        value = getattr(importlib.import_module(module_name), attr_name)
+        globals()[name] = value  # cache: later lookups skip __getattr__
+        return value
+    if name in _LAZY_SUBMODULES:
+        module = importlib.import_module(f"{__name__}.{name}")
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_LAZY_ATTRS) | set(_LAZY_SUBMODULES))

bayesmith/bridge/__init__.py

@@ -0,0 +1 @@
+"""Bridges to external inference engines."""

bayesmith/bridge/numpyro_bridge.py

@@ -0,0 +1,205 @@
+"""Turning a graph into a NumPyro model, and running NUTS on it.
+
+This is the last row of the dispatch table: whatever structure bayesmith
+cannot solve exactly is handed to NumPyro. It is also the oracle every exact
+path is checked against, because a graph that qualifies for an exact method
+always also qualifies for NUTS.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpyro
+from numpyro.infer import MCMC, NUTS
+
+from bayesmith.graph.evaluate import apply_deterministic, apply_probabilistic
+from bayesmith.graph.graph import Graph
+from bayesmith.graph.nodes import Const, Deterministic, Probabilistic
+
+
+def to_numpyro(graph: Graph) -> Callable[[], dict[str, Any]]:
+    """Build a NumPyro model that declares the same joint distribution.
+
+    Latent and observed nodes become ``numpyro.sample`` sites carrying the
+    graph's own node names, so posterior samples come back keyed by them.
+    Deterministic nodes are recorded with ``numpyro.deterministic`` so they
+    appear in traces and predictives without contributing a density.
+    """
+
+    def model() -> dict[str, Any]:
+        env: dict[str, Any] = {}
+        for node in graph.nodes:
+            if isinstance(node, Const):
+                env[node.name] = node.value
+            elif isinstance(node, Deterministic):
+                env[node.name] = numpyro.deterministic(
+                    node.name, apply_deterministic(graph, node, env)
+                )
+            elif isinstance(node, Probabilistic):
+                distribution = apply_probabilistic(graph, node, env)
+                if node.plate:
+                    name = node.plate[0]
+                    with numpyro.plate(name, graph.plate_size(name)):
+                        env[node.name] = numpyro.sample(
+                            node.name, distribution, obs=node.observed
+                        )
+                else:
+                    env[node.name] = numpyro.sample(
+                        node.name, distribution, obs=node.observed
+                    )
+        return env
+
+    return model
+
+
+def nuts(
+    graph: Graph,
+    key: jax.Array,
+    *,
+    num_warmup: int = 1000,
+    num_samples: int = 2000,
+    num_chains: int = 1,
+    chain_method: str = "sequential",
+    progress_bar: bool = False,
+    nuts_options: Mapping[str, Any] | None = None,
+) -> dict[str, jax.Array]:
+    """Sample the posterior of ``graph`` with NUTS.
+
+    ``chain_method`` and ``nuts_options`` are here because
+    :meth:`~bayesmith.dispatch.plan.InferencePlan.sample` promises them on
+    every path, and two of its five shapes -- the graph with no exact block,
+    and the SNIS collapse -- run through this function. Until they existed
+    those two shapes silently ignored both keywords while the mixed shape,
+    which reaches ``HMCGibbs`` through
+    :func:`~bayesmith.exact.gibbs.assemble`, honoured them. The names and the
+    defaults are ``assemble``'s, so the two spellings of "run a chain" take
+    the same words.
+
+    Args:
+        graph: the model.
+        key: a PRNG key.
+        num_warmup: adaptation draws, discarded.
+        num_samples: retained draws per chain.
+        num_chains: independent chains.
+        chain_method: how ``num_chains`` are run -- ``"sequential"``,
+            ``"parallel"`` or ``"vectorized"``. All three are numpyro's own
+            and all three are legal here; ``assemble`` refuses
+            ``"vectorized"`` for a Gibbs sweep, but that refusal is about
+            ``HMCGibbs.init`` and does not apply to a bare kernel.
+        progress_bar: whether NumPyro prints progress.
+        nuts_options: keywords for the ``NUTS`` kernel itself
+            (``target_accept_prob``, ``max_tree_depth``, ``dense_mass``, ...).
+
+            **``init_strategy`` goes here, and on a narrow posterior it is
+            not a tuning knob.** NumPyro's default is ``init_to_uniform``,
+            which draws in the UNCONSTRAINED space with no knowledge of
+            where the graph's priors sit; a posterior far narrower than its
+            prior is a needle, and warmup then adapts a step size for
+            wherever it landed in the haystack. Measured on a power law
+            whose amplitude has a prior of 1e4 and a posterior width of
+            0.4, two chains of 400:
+
+            ==============================================  ======  ======
+            init                                            r_hat   ESS
+            ==============================================  ======  ======
+            default (``init_to_uniform``)                    1609     1.0
+            ``init_to_value`` at the declared values          1.006  138.6
+            ==============================================  ======  ======
+
+            The declared point does not have to be good -- only somewhere a
+            gradient can be followed. This is rheplicant's
+            ``init_to_declared`` lesson, and the lesson is what was carried:
+            the remedy already exists here as a keyword, so what was missing
+            was the sentence saying so.
+
+    Returns:
+        A mapping from latent node name to its draws.
+    """
+    mcmc = MCMC(
+        NUTS(to_numpyro(graph), **dict(nuts_options or {})),
+        num_warmup=num_warmup,
+        num_samples=num_samples,
+        num_chains=num_chains,
+        chain_method=chain_method,
+        progress_bar=progress_bar,
+    )
+    mcmc.run(key)
+    return mcmc.get_samples()
+
+
+def predict(
+    graph: Graph, samples: Mapping[str, Any], key: jax.Array | None = None
+) -> dict[str, jax.Array]:
+    """Posterior predictive: every node's value over a stack of draws.
+
+    Args:
+        graph: the model the draws came from.
+        samples: ``{latent name: (n_draws, *latent shape)}`` -- what
+            :func:`nuts` returns, unchanged.
+        key: PRNG key for any node the draws do not fix. Fixed by default,
+            so a predictive is reproducible.
+
+    Returns:
+        ``{node name: (n_draws, *node shape)}``, deterministic nodes
+        included.
+
+    Raises:
+        GraphError: if a latent is missing from ``samples``; if a stack's
+            PER-SAMPLE shape is not the latent's own; or if the stacks
+            disagree about how many draws there are.
+
+    Note:
+        **Why the per-sample shape is checked and not just the name.** This
+        is rheplicant's ``predict_from_samples`` guard, ported because the
+        failure it prevents is reachable here and is silent. Measured on a
+        length-3 latent with three draws: handing the stack in TRANSPOSED
+        returns a finite, correctly-shaped ``(3, 3)`` predictive whose every
+        entry is wrong --
+
+            correct     [[0, 2, 6], [3, 8, 15], [6, 14, 24]]
+            transposed  [[0, 6, 18], [1, 8, 21], [2, 10, 24]]
+
+        -- because NumPyro's ``Predictive`` maps over the leading axis and
+        has no independent statement of what the latent's shape is. The
+        graph does have one: each latent's ``dist_fn``. A non-square
+        transposition raises a broadcast error from three layers down that
+        names neither the site nor the axis; a square one raises nothing at
+        all.
+    """
+    from numpyro.infer import Predictive
+
+    from bayesmith.dispatch.classify import prior_environment
+    from bayesmith.errors import GraphError
+
+    declared = prior_environment(graph)
+    draws: set[int] = set()
+    for name in graph.latents:
+        if name not in samples:
+            raise GraphError(
+                f"samples is missing latent {name!r}; available: "
+                f"{sorted(samples)}. A predictive needs every latent the "
+                "graph declares, because the deterministic nodes read them."
+            )
+        expected = jnp.shape(declared[name])
+        got = jnp.shape(samples[name])
+        if got[1:] != expected:
+            raise GraphError(
+                f"samples[{name!r}] has per-sample shape {got[1:]}, but the "
+                f"latent is {expected} (its full stack is {got}). The LEADING "
+                "axis must be the draw axis. Checking only the name would let "
+                "a transposed stack broadcast into the prediction and return a "
+                "finite, correctly-shaped, wrong predictive -- silently, "
+                "whenever the draw count happens to equal the latent's size."
+            )
+        draws.add(got[0])
+    if len(draws) > 1:
+        raise GraphError(
+            f"the sample stacks disagree about the number of draws: "
+            f"{sorted(draws)}. They must all come from one run."
+        )
+    predictive = Predictive(to_numpyro(graph), posterior_samples=dict(samples))
+    return predictive(jax.random.key(0) if key is None else key)

bayesmith/diagnose/__init__.py

@@ -0,0 +1,50 @@
+"""Design-time diagnostics: what the model cannot tell apart, and what the
+priors did to it.
+
+Three questions, each answered about the model *at a point* rather than
+claimed globally -- which is why this package linearizes locally and never
+reads ``linear_in``:
+
+* :func:`identifiability` -- the rank of the joint Jacobian: which
+  combinations of latents the data is blind to, across blocks, where every
+  per-block guard is structurally silent.
+* :func:`prior_sensitivity` -- how far each latent's own prior moved the
+  mode, in posterior sigmas, from two deterministic routes that verify each
+  other.
+* :class:`JeffreysPrior` -- ``sqrt(det I)`` over a named block, evaluated
+  from the graph's own noise, with the determinant taken by ``eigvalsh``
+  plus a rank floor because ``slogdet`` and ``cholesky`` both return
+  plausible finite answers on a singular block.
+
+Ported from ``rheplicant.inference.{identifiability,sensitivity,priors}``
+(migration spec §八 step 5); the per-module cross-check records live in
+``docs/migration/``. Precision discipline: run everything, graph
+construction included, inside ``with jax.enable_x64(True):`` -- these
+verdicts live at 1e-17 of the largest singular value, and a float32 result
+is refused by name rather than silently reported.
+"""
+
+from bayesmith.diagnose.identifiability import (
+    DEFAULT_RANK_RTOL,
+    IdentifiabilityReport,
+    identifiability,
+)
+from bayesmith.diagnose.priors import JeffreysPrior
+from bayesmith.diagnose.sensitivity import (
+    CRITERION_SHIFT,
+    PriorSensitivityReport,
+    prior_sensitivity,
+)
+
+__all__ = [
+    # identifiability
+    "identifiability",
+    "IdentifiabilityReport",
+    "DEFAULT_RANK_RTOL",
+    # sensitivity
+    "prior_sensitivity",
+    "PriorSensitivityReport",
+    "CRITERION_SHIFT",
+    # joint priors
+    "JeffreysPrior",
+]