probability-flow 0.1.0__py3-none-any.whl

probability_flow/__init__.py

@@ -0,0 +1,41 @@
+"""probability_flow: a from-scratch, modular discrete Bayesian-network library.
+
+The design is in `docs/SPEC.md`; settled choices in `docs/DECISIONS.md` and
+milestones in `docs/ROADMAP.md`. The public API currently lives in
+`probability_flow.core` and is re-exported here for convenience:
+
+    from probability_flow import Node, ExactSolver
+"""
+from importlib.metadata import PackageNotFoundError, version
+
+from .core import (
+    CPD,
+    BayesianNetwork,
+    CompiledCPD,
+    ExactSolver,
+    IndependentEvidenceCPD,
+    LoopySolver,
+    Node,
+    NoisyAndCPD,
+    NoisyOrCPD,
+    TabularCPD,
+)
+
+try:
+    __version__ = version("probability-flow")
+except PackageNotFoundError:  # running from a source checkout, not installed
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "__version__",
+    "Node",
+    "BayesianNetwork",
+    "CompiledCPD",
+    "ExactSolver",
+    "LoopySolver",
+    "CPD",
+    "TabularCPD",
+    "IndependentEvidenceCPD",
+    "NoisyOrCPD",
+    "NoisyAndCPD",
+]

probability_flow/aspic/__init__.py

@@ -0,0 +1,21 @@
+"""probability_flow.aspic: the ASPIC argument-compilation layer.
+
+The first of several planned domain wrappers (legal, medical, AI-safety) over the
+pure-BN `core`. Build an argument out of premises, conclusions, and attacks, then
+`compile()` the target to an ordinary `BayesianNetwork`. See `docs/aspic.md`.
+
+    from probability_flow.aspic import Premise, Axiom, Conclusion
+"""
+from .argument import ArgumentWarning, Axiom, Conclusion, Premise
+from .generate import (
+    ArgumentGenerator,
+    DifficultyTargets,
+    StructuralParams,
+    generate,
+)
+from .handle import Argument
+
+__all__ = [
+    "Premise", "Axiom", "Conclusion", "ArgumentWarning", "Argument",
+    "ArgumentGenerator", "StructuralParams", "DifficultyTargets", "generate",
+]

probability_flow/aspic/argument.py

@@ -0,0 +1,182 @@
+"""The ASPIC authoring layer: free constructors over the core BN.
+
+`Premise`, `Axiom`, and `Conclusion` are core `Node`s (so a compiled argument is
+queried by passing the object straight to a solver, exactly as in core). They are
+one underlying node type; the distinction is intention plus light validation. The
+argumentative structure is *declared* with fluent methods on the downstream node
+and only *lowered* to core CPDs by `compile()` (see `compile.py`). Keeping the
+declaration separate from the core CPD makes `compile()` a pure function of the
+graph, so it is repeatable.
+
+Every argumentative edge is a method on its downstream node: `support` / `rebut` /
+`strict` add an upstream argument to a conclusion, `undermine` attacks a premise,
+`undercut` attacks an inference. The whole graph is therefore reachable by walking
+inputs from the root target, the traversal `compile()` performs. See
+`docs/aspic.md`.
+"""
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+from ..core import Node
+
+if TYPE_CHECKING:
+    from ..core import BayesianNetwork
+    from .handle import Argument
+
+
+class ArgumentWarning(UserWarning):
+    """Soft validation: a leaf conclusion, rebutting a premise, a non-disjoint
+    undercut. These do not stop compilation (they match core's permissiveness),
+    they flag a likely modelling slip."""
+
+
+class _Claim(Node):
+    """A proposition in the argument: a core `Node` plus declared, not-yet-lowered
+    argumentative edges. Not constructed directly; use `Premise` / `Axiom` /
+    `Conclusion`."""
+
+    role = "claim"
+
+    def __init__(self, name: str, prior: float = 0.5):
+        super().__init__(name, prior=prior)
+        self._edges: list[tuple["Node", float, str]] = []  # (src, lr, kind)
+        self._strict: list["Node"] = []                    # strict sources
+        self._undercuts: list[tuple["Node", "Node"]] = []  # (attacked source, undercutter)
+        self._no_undermine = False
+        # opaque, serialized but uninterpreted by the library (see docs/aspic.md):
+        self.desc: str | None = None              # a longer description
+        self.revealable_to_judge = None           # benchmark metadata, semantics TBD
+
+    def _require_new_edge(self, src: "Node") -> None:
+        """At most one argumentative edge joins a pair of claims. This is what lets
+        an undercut be addressed by its endpoints and keeps serialized edge ids
+        (`source->target`) unique."""
+        if any(s is src for s, _lr, _kind in self._edges) or any(
+            s is src for s in self._strict
+        ):
+            raise ValueError(
+                f"{self.name!r} already has an edge from {src.name!r}; at most one "
+                "argumentative edge may join a pair of claims"
+            )
+
+    # --- defeasible and strict support (methods on the downstream conclusion) --
+
+    def support(self, src: "Node", lr: float) -> "Node":
+        """Add a defeasible argument *for* this conclusion (`lr > 1`). Returns
+        `src`, so an inline source can be built further upstream, as with core
+        `add_input`."""
+        if not lr > 1:
+            raise ValueError(
+                f"support lr must be > 1 (got {lr}); use rebut for an argument against"
+            )
+        self._require_new_edge(src)
+        self._edges.append((src, float(lr), "support"))
+        return src
+
+    def rebut(self, src: "Node", lr: float) -> "Node":
+        """Add a defeasible argument *against* this conclusion (`0 < lr < 1`)."""
+        if not 0 < lr < 1:
+            raise ValueError(
+                f"rebut lr must be in (0, 1) (got {lr}); use support for an argument for"
+            )
+        if self.role == "premise":
+            warnings.warn(
+                f"rebutting {self.name!r}, a premise: a rebutted node is really a "
+                "conclusion. Either model it as one, or use undermine to attack a premise.",
+                ArgumentWarning,
+                stacklevel=2,
+            )
+        self._require_new_edge(src)
+        self._edges.append((src, float(lr), "rebut"))
+        return src
+
+    def strict(self, src: "Node") -> "Node":
+        """Add a strict argument: `src` true forces this conclusion (a rebut cannot
+        pull it down). Lowered by parent-divorcing."""
+        self._require_new_edge(src)
+        self._strict.append(src)
+        return src
+
+    # --- attacks (also methods on the downstream node) ------------------------
+
+    def undermine(self, by: "Node", lr: float) -> "Node":
+        """Attack this premise with `by` (`0 < lr < 1`). Mechanically a rebut into
+        the attacked node, named distinctly because it is the ASPIC-correct verb
+        for attacking a premise. An axiom cannot be undermined."""
+        if self._no_undermine:
+            raise ValueError(
+                f"cannot undermine {self.name!r}: an axiom has no defeasible premise to attack"
+            )
+        if not 0 < lr < 1:
+            raise ValueError(f"undermine lr must be in (0, 1) (got {lr})")
+        self._require_new_edge(by)
+        self._edges.append((by, float(lr), "undermine"))
+        return by
+
+    def undercut(self, source: "Node", by: "Node") -> "Node":
+        """Attack the `source -> self` inference with `by`: where `source`
+        supports (or strictly entails) this conclusion, `by` asserts the inference
+        does not apply. Full neutralization (`tau = 1`) is hard-coded. Returns the
+        undercutter `by`."""
+        self._undercuts.append((source, by))
+        return by
+
+    def compile(self) -> "BayesianNetwork":
+        """Lower this argument (as the target) to a core `BayesianNetwork`."""
+        from .compile import compile_argument
+
+        return compile_argument(self)
+
+    def assemble(self) -> "Argument":
+        """Bundle this argument (as the target) into an `Argument` handle: the home
+        for serialization (`to_json` / `save`), the compiled network (`.bn`), and
+        the argument-level metrics. Authoring is unchanged; this just hands back an
+        object to serialize and measure from."""
+        from .handle import Argument
+
+        return Argument(self)
+
+    @property
+    def bn(self) -> "BayesianNetwork":
+        """The compiled network for this argument (a fresh compile each access)."""
+        return self.compile()
+
+    def render(self, **kwargs):
+        """Draw the argument at the argument level (premises, conclusions, and
+        attacks), hiding the lowered splice machinery. See `aspic.visualization`.
+        For the full compiled network instead, use `self.bn.render()`."""
+        from .visualization import render_argument
+
+        return render_argument(self, **kwargs)
+
+
+class Premise(_Claim):
+    """An asserted proposition: a leaf with a credence and no incoming argument.
+    Underminable. `prior` is its credence."""
+
+    role = "premise"
+
+    def __init__(self, name: str, prior: float = 0.5):
+        super().__init__(name, prior=prior)
+
+
+class Axiom(Premise):
+    """A premise held with certainty (`prior = 1`) that cannot be undermined."""
+
+    role = "axiom"
+
+    def __init__(self, name: str):
+        super().__init__(name, prior=1.0)
+        self._no_undermine = True
+
+
+class Conclusion(_Claim):
+    """A proposition argued for: it carries a prior (default `0.5`) and accumulates
+    support, rebut, and strict arguments."""
+
+    role = "conclusion"
+
+    def __init__(self, name: str, prior: float = 0.5):
+        super().__init__(name, prior=prior)

probability_flow/aspic/calibrate.py

@@ -0,0 +1,178 @@
+"""JAX calibration of a built argument (optional — needs the `[jax]` extra).
+
+The generator's structure choices (depth, branching, edge kinds) are discrete and
+non-differentiable, so the core stays plain numpy. But for a *fixed* structure the
+continuous parameters — node priors and edge likelihood ratios — are smooth, and a
+generated graph is a polytree, so its root posterior is the topo-ordered forward
+pass of the same message math. This module expresses that forward pass in JAX, so:
+
+- `sensitivities(arg)` returns `d p_root / d parameter` for every prior and LR (via
+  `jax.grad`) — a per-parameter importance signal; and
+- `calibrate_posterior(arg, target)` solves the root's prior and edge LRs for a
+  target posterior by projected gradient descent and writes them back. This is the
+  multi-parameter upgrade of the dependency-free bisection in `generate` (which
+  tunes the root prior alone); it keeps every other subtree's marginals fixed.
+
+Defined for the CPD shapes the ASPIC compiler emits on a polytree
+(`IndependentEvidence`, the strict noisy-OR, the undercut AND-NOT splice). A
+non-polytree (shared / non-disjoint structure) is out of scope — the forward
+marginal pass is only exact on a tree.
+"""
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .handle import Argument
+
+
+def _require_jax():
+    try:
+        import jax
+        import jax.numpy as jnp
+    except ImportError as exc:  # pragma: no cover - exercised only without the extra
+        raise ImportError(
+            "calibrate needs JAX; install the optional extra: "
+            "pip install 'probability-flow[jax]'"
+        ) from exc
+    return jax, jnp
+
+
+def _logit(p: float) -> float:
+    p = min(1 - 1e-6, max(1e-6, p))
+    return math.log(p / (1 - p))
+
+
+def _is_and_not_splice(cpd) -> bool:
+    from ..core import TabularCPD
+
+    if not isinstance(cpd, TabularCPD) or len(cpd.inputs) != 2:
+        return False
+    import numpy as np
+
+    p1 = np.exp(cpd.log_table[..., 1])
+    return np.allclose(p1, np.array([[0.0, 0.0], [1.0, 0.0]]))
+
+
+def _combos(k: int):
+    import numpy as np
+
+    return np.array([[(c >> i) & 1 for i in range(k)] for c in range(2**k)], dtype=float)
+
+
+def _forward(arg: "Argument"):
+    """Build the polytree forward pass. Returns `(p_root, theta0, meta)` where
+    `p_root(theta)` is a JAX function of the flat parameter vector, `theta0` are the
+    current values (logit priors, log LRs), and `meta` labels each entry."""
+    jax, jnp = _require_jax()
+    from ..core import IndependentEvidenceCPD, NoisyOrCPD
+
+    bn = arg.bn
+    nodes = list(bn.nodes)                         # topological (inputs before node)
+    cpd_of = {n: bn.compiled_cpd(n).cpd for n in nodes}
+    ins_of = {n: tuple(bn.compiled_cpd(n).inputs) for n in nodes}
+
+    prior_idx, lr_idx, theta0, meta = {}, {}, [], []
+    for n in nodes:
+        cpd = cpd_of[n]
+        if isinstance(cpd, IndependentEvidenceCPD):
+            prior_idx[n] = len(theta0)
+            theta0.append(_logit(n.prior))
+            meta.append(("prior", n.name))
+            for inp, lr in zip(cpd.inputs, cpd.lrs):
+                lr_idx[(n, inp)] = len(theta0)
+                theta0.append(math.log(lr))
+                meta.append(("lr", (inp.name, n.name)))
+    combos = {n: jnp.asarray(_combos(len(ins_of[n]))) for n in nodes
+              if isinstance(cpd_of[n], IndependentEvidenceCPD) and ins_of[n]}
+
+    def p_root(theta):
+        m = {}
+        for n in nodes:
+            cpd, ins = cpd_of[n], ins_of[n]
+            if isinstance(cpd, IndependentEvidenceCPD):
+                b = theta[prior_idx[n]]
+                if not ins:
+                    m[n] = jax.nn.sigmoid(b)
+                else:
+                    w = jnp.stack([theta[lr_idx[(n, i)]] for i in ins])
+                    pin = jnp.stack([m[i] for i in ins])
+                    C = combos[n]
+                    wts = jnp.prod(jnp.where(C == 1, pin, 1 - pin), axis=1)
+                    m[n] = jnp.sum(jax.nn.sigmoid(b + C @ w) * wts)
+            elif isinstance(cpd, NoisyOrCPD):
+                a = jnp.asarray(cpd.activations)
+                pin = jnp.stack([m[i] for i in ins])
+                m[n] = 1 - (1 - cpd.leak) * jnp.prod(1 - a * pin)
+            elif _is_and_not_splice(cpd):
+                e, u = ins
+                m[n] = m[e] * (1 - m[u])
+            else:
+                raise NotImplementedError(
+                    f"calibrate does not handle {type(cpd).__name__} on {n.name!r}"
+                )
+        return m[arg.target]
+
+    return p_root, jnp.asarray(theta0), meta
+
+
+def sensitivities(arg: "Argument") -> dict:
+    """`d p_root / d parameter` at the current values, for every prior and LR.
+
+    Keys are `("prior", node_name)` and `("lr", (source_name, target_name))`; values
+    are gradients in the natural unconstrained space (logit prior, log LR). A large
+    magnitude marks a parameter the verdict is sensitive to — a manipulability /
+    importance signal."""
+    jax, _ = _require_jax()
+    p_root, theta0, meta = _forward(arg)
+    grad = jax.grad(p_root)(theta0)
+    return {key: float(g) for key, g in zip(meta, grad)}
+
+
+def calibrate_posterior(arg: "Argument", target: float, *,
+                        steps: int = 500, step_size: float = 0.3) -> float:
+    """Solve the root's prior and edge LRs so its posterior hits `target`, by
+    projected gradient descent (priors kept in `[0.02, 0.98]`, LRs kept to their
+    original sign with `|log lr| <= log 50`), then write the result back onto the
+    argument. Every other subtree's marginals are held fixed. Returns the achieved
+    posterior. Use the bisection in `generate` for the dependency-free single-knob
+    version."""
+    jax, jnp = _require_jax()
+    from ..core import IndependentEvidenceCPD, LoopySolver
+
+    bn = arg.bn
+    root = arg.target
+    cpd = bn.compiled_cpd(root).cpd
+    if not isinstance(cpd, IndependentEvidenceCPD):
+        raise NotImplementedError("calibrate_posterior targets a defeasible root")
+    ins = list(cpd.inputs)
+    solver = LoopySolver(bn)
+    pin = jnp.asarray([solver.prob(i, 1) for i in ins])         # fixed input marginals
+    k = len(ins)
+    C = jnp.asarray(_combos(k)) if k else None
+    w0 = jnp.asarray([math.log(lr) for lr in cpd.lrs])
+    signs = jnp.sign(w0)
+    lo_b, hi_b = _logit(0.02), _logit(0.98)
+    lo_w, hi_w = math.log(1.01), math.log(50.0)
+
+    def predict(theta):
+        b, w = theta[0], theta[1:]
+        if k == 0:
+            return jax.nn.sigmoid(b)
+        wts = jnp.prod(jnp.where(C == 1, pin, 1 - pin), axis=1)
+        return jnp.sum(jax.nn.sigmoid(b + C @ w) * wts)
+
+    grad = jax.grad(lambda th: (predict(th) - target) ** 2)
+    theta = jnp.concatenate([jnp.asarray([_logit(root.prior)]), w0])
+    for _ in range(steps):
+        theta = theta - step_size * grad(theta)
+        b = jnp.clip(theta[0], lo_b, hi_b)
+        w = signs * jnp.clip(jnp.abs(theta[1:]), lo_w, hi_w)
+        theta = jnp.concatenate([jnp.asarray([b]), w])
+
+    root.prior = float(jax.nn.sigmoid(theta[0]))
+    new_lrs = [float(math.exp(x)) for x in theta[1:]]
+    for j, (src, _lr, kind) in enumerate(root._edges):           # aligned with cpd.inputs
+        root._edges[j] = (src, new_lrs[j], kind)
+    return float(predict(theta))

probability_flow/aspic/compile.py

@@ -0,0 +1,175 @@
+"""Lowering an ASPIC argument to a core `BayesianNetwork`. See `docs/aspic.md`.
+
+`compile_argument(target)` collects everything reachable from the target, runs
+validation (raising on sign / axiom errors, warning on soft slips), lowers each
+claim to core CPDs, then defers to the core compiler. The lowerings:
+
+- premise / conclusion with only defeasible edges -> `IndependentEvidenceCPD`
+  (a premise is just the zero-support case; an undermined premise gains an
+  `lr < 1` input).
+- a conclusion with strict edges -> parent-divorcing: a hidden `D` carries the
+  defeasible part (or just the prior), and the conclusion becomes a leak-0
+  `NoisyOr` over its strict inputs and `D`, each at activation 1.
+- an undercut on the `source -> C` edge -> a spliced `X = source AND NOT U`
+  (`tau = 1`) replacing `source` on that edge, with multiple undercutters first
+  combined by an upstream leak-0 `NoisyOr` into one `U`.
+"""
+from __future__ import annotations
+
+import warnings
+
+from ..core import BayesianNetwork, IndependentEvidenceCPD, Node, TabularCPD
+from .argument import ArgumentWarning, _Claim
+
+
+def compile_argument(target: "_Claim") -> "BayesianNetwork":
+    claims = _reachable(target)
+    _validate(claims)
+    for c in claims:
+        if isinstance(c, _Claim):
+            _lower(c)
+    return BayesianNetwork.from_nodes([target])
+
+
+# --- traversal -------------------------------------------------------------
+
+def _upstream(c: "Node") -> list["Node"]:
+    """The declared upstream sources of a claim: edge and strict sources plus any
+    undercutters. Plain (non-claim) nodes are treated as leaves."""
+    out: list["Node"] = []
+    out.extend(src for src, _lr, _kind in getattr(c, "_edges", []))
+    out.extend(getattr(c, "_strict", []))
+    out.extend(by for _source, by in getattr(c, "_undercuts", []))
+    return out
+
+
+def _reachable(target: "Node") -> list["Node"]:
+    """Every node reachable from the target through declared argument edges, in a
+    deterministic, dependency-respecting order (sources before the nodes using
+    them is not required here; the core compiler re-orders)."""
+    seen: dict[int, "Node"] = {}
+    order: list["Node"] = []
+    stack = [target]
+    while stack:
+        c = stack.pop()
+        if id(c) in seen:
+            continue
+        seen[id(c)] = c
+        order.append(c)
+        stack.extend(_upstream(c))
+    return order
+
+
+def _leaf_premises(c: "Node") -> set["Node"]:
+    """The leaf premises (no incoming argument) upstream of `c`, for the undercut
+    disjointness check."""
+    out: set["Node"] = set()
+    seen: set[int] = set()
+    stack = [c]
+    while stack:
+        n = stack.pop()
+        if id(n) in seen:
+            continue
+        seen.add(id(n))
+        ups = _upstream(n)
+        if not ups:
+            out.add(n)
+        else:
+            stack.extend(ups)
+    return out
+
+
+# --- validation ------------------------------------------------------------
+
+def _validate(claims: list["Node"]) -> None:
+    for c in claims:
+        if not isinstance(c, _Claim):
+            continue
+
+        if c.role == "conclusion" and not c._edges and not c._strict:
+            warnings.warn(
+                f"{c.name!r} is a conclusion with no incoming argument; a leaf "
+                "should be a Premise.",
+                ArgumentWarning,
+                stacklevel=2,
+            )
+
+        edge_sources = {id(src) for src, _lr, _kind in c._edges}
+        strict_sources = {id(s) for s in c._strict}
+        for source, by in c._undercuts:
+            if id(source) not in edge_sources and id(source) not in strict_sources:
+                raise ValueError(
+                    f"undercut on {c.name!r}: {source.name!r} is not a source of any "
+                    "edge into it; undercut attacks an existing inference"
+                )
+            shared = _leaf_premises(source) & _leaf_premises(by)
+            if shared:
+                names = ", ".join(sorted(p.name for p in shared))
+                warnings.warn(
+                    f"undercut on the {source.name!r} -> {c.name!r} edge shares "
+                    f"premise(s) {{{names}}} with its undercutter {by.name!r}; the "
+                    "splice is exact only for disjoint branches.",
+                    ArgumentWarning,
+                    stacklevel=2,
+                )
+
+
+# --- lowering --------------------------------------------------------------
+
+def _and_not_cpd(x: "Node", e: "Node", u: "Node") -> "TabularCPD":
+    """`X = 1` iff `E = 1` and `U = 0`; the `tau = 1` undercut splice."""
+    # table[e][u] = P(X | e, u); only (E=1, U=0) yields X=1.
+    table = [
+        [[1.0, 0.0], [1.0, 0.0]],   # E=0: X=0 regardless of U
+        [[0.0, 1.0], [1.0, 0.0]],   # E=1: U=0 -> X=1, U=1 -> X=0
+    ]
+    return TabularCPD(x, [e, u], table=table)
+
+
+def _splice_for(c: "_Claim") -> dict[int, "Node"]:
+    """Build the undercut splice node for each attacked source on `c`, keyed by the
+    source's id. Multiple undercutters on one edge are OR-combined into one `U`."""
+    grouped: dict[int, tuple["Node", list["Node"]]] = {}
+    for source, u in c._undercuts:
+        grouped.setdefault(id(source), (source, []))[1].append(u)
+
+    splice: dict[int, "Node"] = {}
+    for sid, (source, undercutters) in grouped.items():
+        if len(undercutters) == 1:
+            u_node = undercutters[0]
+        else:
+            u_node = Node(f"{c.name}/U[{source.name}]")
+            u_node.noisy_or(leak=0.0)
+            for u in undercutters:
+                u_node.add_input(u, activation=1.0)
+        x = Node(f"{source.name} (spliced)")
+        x.set_cpd(_and_not_cpd(x, source, u_node))
+        splice[sid] = x
+    return splice
+
+
+def _lower(c: "_Claim") -> None:
+    """Set `c`'s core CPD (and any hidden helper nodes) from its declared edges.
+    Idempotent: the CPD is rebuilt from scratch, so re-compiling is safe."""
+    c._cpd = IndependentEvidenceCPD(c)  # reset for a clean, repeatable lowering
+
+    splice = _splice_for(c)
+
+    def resolve(src: "Node") -> "Node":
+        return splice.get(id(src), src)
+
+    defeasible = [(resolve(src), lr) for src, lr, _kind in c._edges]
+    strict = [resolve(src) for src in c._strict]
+
+    if strict:
+        # parent-divorcing: D holds the defeasible part (or just the prior).
+        d = Node(f"{c.name}/defeasible", prior=c.prior)
+        for src, lr in defeasible:
+            d.add_input(src, lr=lr)
+        c.noisy_or(leak=0.0)
+        for src in strict:
+            c.add_input(src, activation=1.0)
+        c.add_input(d, activation=1.0)
+    else:
+        for src, lr in defeasible:
+            c.add_input(src, lr=lr)