probability-flow 0.3.0__tar.gz → 0.4.1__tar.gz

{probability_flow-0.3.0 → probability_flow-0.4.1}/.gitignore

@@ -6,3 +6,6 @@ __pycache__/
 *.egg-info/
 .DS_Store
 _previews/
+
+# Local secrets
+.env

{probability_flow-0.3.0 → probability_flow-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: probability-flow
-Version: 0.3.0
+Version: 0.4.1
 Summary: A from-scratch, modular discrete Bayesian-network library.
 Project-URL: Homepage, https://github.com/scalable-oversight-benchmarks/probability-flow
 Project-URL: Repository, https://github.com/scalable-oversight-benchmarks/probability-flow
@@ -28,6 +28,7 @@ Requires-Dist: pytest; extra == 'dev'
 Requires-Dist: ruff; extra == 'dev'
 Provides-Extra: jax
 Requires-Dist: jax; extra == 'jax'
+Requires-Dist: scipy; extra == 'jax'
 Provides-Extra: viz
 Requires-Dist: matplotlib; extra == 'viz'
 Description-Content-Type: text/markdown

{probability_flow-0.3.0 → probability_flow-0.4.1}/probability_flow/aspic/__init__.py

@@ -9,13 +9,16 @@ pure-BN `core`. Build an argument out of premises, conclusions, and attacks, the
 from .argument import ArgumentWarning, Axiom, Conclusion, Premise
 from .generate import (
     ArgumentGenerator,
+    BatchParams,
     DifficultyTargets,
     StructuralParams,
     generate,
+    generate_batch,
 )
 from .handle import Argument
 
 __all__ = [
     "Premise", "Axiom", "Conclusion", "ArgumentWarning", "Argument",
     "ArgumentGenerator", "StructuralParams", "DifficultyTargets", "generate",
+    "BatchParams", "generate_batch",
 ]

{probability_flow-0.3.0 → probability_flow-0.4.1}/probability_flow/aspic/argument.py

@@ -47,6 +47,11 @@ class _Claim(Node):
         # (source_a, source_b, J): residual correlation between two defeasible
         # sources of this conclusion, lowered to a CorrelatedEvidenceCPD coupling.
         self._correlations: list[tuple["Node", "Node", "float | None"]] = []
+        # (sources, lr, leak, rule): a conjunctive support group — the sources
+        # JOINTLY support this conclusion through one noisy-AND inference (the
+        # deduction rule), lowered to a hidden NoisyAnd gate fed in as one lr.
+        self._conjunctions: list[
+            tuple[list["Node"], "float | None", float, "str | None"]] = []
         self._no_undermine = False
         # opaque, serialized but uninterpreted by the library (see docs/aspic.md):
         self.desc: str | None = None              # a longer description
@@ -56,8 +61,13 @@ class _Claim(Node):
         """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
+        in_conj = any(
+            s is src for srcs, _lr, _leak, _rule in self._conjunctions for s in srcs
+        )
+        if (
+            any(s is src for s, _lr, _kind in self._edges)
+            or any(s is src for s in self._strict)
+            or in_conj
         ):
             raise ValueError(
                 f"{self.name!r} already has an edge from {src.name!r}; at most one "
@@ -80,6 +90,51 @@ class _Claim(Node):
         self._edges.append((src, None if lr is None else float(lr), "support"))
         return src
 
+    def support_all(
+        self,
+        srcs: "list[Node]",
+        lr: "float | None" = None,
+        leak: float = 0.0,
+        rule: "str | None" = None,
+    ) -> "list[Node]":
+        """Conjunctive support: `srcs` JOINTLY support this conclusion through a
+        single noisy-AND inference — the deduction fires only when *all* the sources
+        hold. This is the faithful mapping of a deductive entailment step: N
+        antecedent siblings plus one inference rule yield one deduced conclusion,
+        the shape of every node in a MuSR-style reasoning tree. Returns `srcs`, so
+        inline sources can be built further upstream (as with `support`).
+
+        Lowered (see `compile.py`) to a hidden `NoisyAnd` gate over `srcs`, each at
+        activation 1 — the gate is the logical AND, firing iff every source holds —
+        fed into this conclusion as a single `lr` support. So all the evidential
+        strength is the one `lr > 1` (where the inference rule's re-elicited weight
+        lands), and the conclusion keeps its prior and composes with ordinary
+        `support` / `rebut` edges. `leak` is the rule's fallibility:
+        `P(fire | all sources hold) = 1 - leak`. `rule` is an optional label for the
+        inference licence (carried for audit / rendering, not interpreted). `lr` may
+        be omitted (`None`) to scaffold the group before its strength is known.
+        """
+        srcs = list(srcs)
+        if len(srcs) < 2:
+            raise ValueError(
+                "support_all needs >= 2 sources (it is a conjunction); use support "
+                "for a single argument"
+            )
+        if len({id(s) for s in srcs}) != len(srcs):
+            raise ValueError("support_all sources must be distinct")
+        if lr is not None and not lr > 1:
+            raise ValueError(
+                f"support_all lr must be > 1 (got {lr}); it is conjunctive support"
+            )
+        if not 0.0 <= leak <= 1.0:
+            raise ValueError("leak must be a probability in [0, 1]")
+        for src in srcs:
+            self._require_new_edge(src)
+        self._conjunctions.append(
+            (srcs, None if lr is None else float(lr), float(leak), rule)
+        )
+        return srcs
+
     def rebut(self, src: "Node", lr: "float | None" = None) -> "Node":
         """Add a defeasible argument *against* this conclusion (`0 < lr < 1`). `lr`
         may be omitted (`None`) to declare the edge before its strength is known."""

probability_flow-0.4.1/probability_flow/aspic/benchmark.py

@@ -0,0 +1,254 @@
+"""Benchmark helpers: generate → optimize bridge for large-batch posterior targeting.
+
+This module isolates the JAX/scipy dependency (via ``optimize``) from the pure-Python
+``generate.py``.  The primary entry point is ``generate_with_targets``, which produces
+N argument graphs each tuned to a requested root posterior.
+
+Strategy
+--------
+For each (seed, target_posterior) pair:
+
+1. Try ``generate(seed, structural, DifficultyTargets(target_posterior=t))`` — this
+   bisects the root prior to hit *t* cheaply (no JAX needed).
+2. If the bisection cannot reach *t* for the drawn structure (``RuntimeError`` from
+   the rejection loop), fall back to:
+   a. Generate an *unconstrained* graph from the same seed.
+   b. Check whether *t* is inside ``achievable_interval(arg)``.
+   c. If yes, call ``optimize(arg, target_posterior=t)`` to tune LRs.
+3. If neither path succeeds, record the failure (skip / warn / raise per
+   ``on_failure``).
+
+Both paths return an ``Argument`` with identical public interface.  The ``optimize``
+fallback is more expensive (JAX JIT + SLSQP) but reaches extreme posteriors that
+bisection cannot.
+
+Batch labeling
+--------------
+Each accepted argument is annotated with a ``_benchmark_meta`` attribute (a plain
+dict) carrying the difficulty scalars the benchmark consumer (``debate-eval``) needs:
+
+    posterior, manipulability, d_sep_count, max_depth_metric,
+    upstream_size, circuit_rank, seed, is_exact_manipulability,
+    used_optimize_fallback
+
+These come entirely from existing metrics — no new computation.
+
+JAX/scipy availability
+----------------------
+Both are required for the ``optimize`` fallback path.  When they are absent the
+module still imports cleanly; ``generate_with_targets`` will error only if an
+optimize fallback is actually needed.  Use ``generate_batch`` (no JAX) when you only
+need bisection-reachable posteriors.
+"""
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from .handle import Argument
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _label(arg: "Argument", seed: int, used_optimize: bool) -> None:
+    """Attach ``_benchmark_meta`` to *arg* in-place."""
+    from ..metrics import (
+        circuit_rank,
+        d_separated_groups,
+        max_depth,
+        posterior_range,
+        upstream_size,
+    )
+
+    bn = arg.bn
+    p = arg.posterior(arg.target)
+
+    # posterior_range: exact=True on polytrees, outer bound on shared graphs.
+    # Report manipulability as the [min, max] achievable root posterior — far more
+    # informative than the width alone, since it shows *where* on [0,1] the judge can
+    # be pushed (e.g. a one-sided graph is pinned to one side of the 0.5 threshold).
+    pr = posterior_range(bn, arg.target, exact=True)
+
+    arg._benchmark_meta = {
+        "posterior": float(p),
+        "manipulability": [round(float(pr.lo), 4), round(float(pr.hi), 4)],
+        "manipulability_width": round(float(pr.hi - pr.lo), 4),
+        "d_sep_count": len(d_separated_groups(bn, arg.target)),
+        "max_depth_metric": max_depth(bn, arg.target),
+        "upstream_size": upstream_size(bn, arg.target),
+        "circuit_rank": circuit_rank(bn),
+        "seed": seed,
+        "is_exact_manipulability": bool(pr.is_exact),
+        "used_optimize_fallback": used_optimize,
+    }
+
+
+def _try_optimize(
+    arg: "Argument",
+    target: float,
+    posterior_tol: float,
+    param_limits: Optional[dict],
+) -> "Optional[Argument]":
+    """Try to hit *target* via ``optimize``; return the new arg or None on failure."""
+    try:
+        from .optimize import (
+            InfeasibleTargetError,
+            OptimizeError,
+            achievable_interval,
+            optimize,
+        )
+    except ImportError:
+        return None
+
+    try:
+        lo, hi = achievable_interval(arg, param_limits=param_limits)
+    except Exception:
+        return None
+
+    if not (lo - 1e-4 <= target <= hi + 1e-4):
+        return None  # infeasible for this structure
+
+    try:
+        return optimize(arg, target_posterior=target,
+                        posterior_tol=posterior_tol,
+                        param_limits=param_limits)
+    except (InfeasibleTargetError, OptimizeError):
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def generate_with_targets(
+    n: int,
+    structural,
+    posterior_targets: "list[float]",
+    posterior_tol: float = 0.02,
+    param_limits: "Optional[dict]" = None,
+    seeds: "Optional[list[int]]" = None,
+    max_attempts_per_graph: int = 200,
+    n_jobs: int = 1,
+    on_failure: str = "skip",
+    verbose: bool = False,
+) -> "list[Argument]":
+    """Generate *n* graphs, each tuned to a target root posterior.
+
+    Parameters
+    ----------
+    n : int
+        Number of graphs to produce.
+    structural : StructuralParams
+        Structural shape shared across all graphs.
+    posterior_targets : list[float]
+        Desired root posterior for each graph.  Must have length *n*.
+    posterior_tol : float
+        Accepted absolute deviation from each target posterior (default 0.02).
+    param_limits : dict or None
+        Passed to ``optimize`` on the fallback path.  None uses default boxes.
+    seeds : list[int] or None
+        Per-graph seeds.  Defaults to ``list(range(n))``.
+    max_attempts_per_graph : int
+        Rejection budget for the primary ``generate`` path.
+    n_jobs : int
+        Worker count.  Currently only 1 is fully supported for the optimize
+        fallback (JAX / fork safety, OQ3 in spec); n_jobs > 1 on the primary
+        path only is safe.
+    on_failure : str
+        ``"skip"`` (default) | ``"warn"`` | ``"raise"``.
+    verbose : bool
+        Print progress per graph.
+
+    Returns
+    -------
+    list[Argument]
+        Accepted and tuned arguments.  Each carries ``._benchmark_meta``.
+        Length <= *n* when ``on_failure != 'raise'``.
+
+    Notes
+    -----
+    The fallback threshold (when to call ``optimize`` instead of resampling) is:
+    whenever the bisection-only ``generate`` raises ``RuntimeError`` (budget
+    exhausted).  ``optimize`` is then tried on an *unconstrained* graph from the
+    same seed; if the target is achievable in the parameter box, it succeeds.
+    """
+    from .generate import DifficultyTargets, generate
+
+    if len(posterior_targets) != n:
+        raise ValueError(
+            f"posterior_targets has {len(posterior_targets)} entries but n={n}"
+        )
+    if seeds is None:
+        seeds = list(range(n))
+    if len(seeds) != n:
+        raise ValueError(f"seeds has {len(seeds)} entries but n={n}")
+
+    accepted: list[Argument] = []
+
+    for i, (seed, t) in enumerate(zip(seeds, posterior_targets)):
+        if verbose:
+            print(f"[{i + 1}/{n}] seed={seed} target={t:.3f}", end=" ... ")
+
+        # --- primary path: bisection inside generate() -----------------------
+        arg: Optional[Argument] = None
+        used_optimize = False
+        try:
+            arg = generate(
+                seed=seed,
+                structural=structural,
+                targets=DifficultyTargets(
+                    target_posterior=t,
+                    posterior_tol=posterior_tol,
+                ),
+                max_attempts=max_attempts_per_graph,
+                verbose=False,
+            )
+        except RuntimeError:
+            pass  # bisection budget exhausted — try the optimize fallback
+
+        # --- fallback: generate unconstrained then optimize ------------------
+        if arg is None:
+            try:
+                base_arg = generate(
+                    seed=seed,
+                    structural=structural,
+                    max_attempts=max_attempts_per_graph,
+                    verbose=False,
+                )
+                opt_arg = _try_optimize(base_arg, t, posterior_tol, param_limits)
+                if opt_arg is not None:
+                    arg = opt_arg
+                    used_optimize = True
+            except RuntimeError:
+                pass
+
+        # --- failure handling ------------------------------------------------
+        if arg is None:
+            msg = (
+                f"generate_with_targets: seed {seed} could not reach "
+                f"target {t:.4f} via bisection or optimize."
+            )
+            if on_failure == "raise":
+                raise RuntimeError(msg)
+            if on_failure == "warn":
+                warnings.warn(msg, stacklevel=2)
+            if verbose:
+                print("FAILED")
+            continue
+
+        # --- label and collect -----------------------------------------------
+        _label(arg, seed=seed, used_optimize=used_optimize)
+        accepted.append(arg)
+        if verbose:
+            m = arg._benchmark_meta
+            print(
+                f"ok  P={m['posterior']:.3f}  manip={m['manipulability']:.3f}  "
+                f"nodes={m['upstream_size'] + 1}  "
+                f"{'[opt]' if used_optimize else '[bisect]'}"
+            )
+
+    return accepted

{probability_flow-0.3.0 → probability_flow-0.4.1}/probability_flow/aspic/calibrate.py

@@ -67,6 +67,7 @@ def _forward(arg: "Argument"):
     current values (logit priors, log LRs), and `meta` labels each entry."""
     jax, jnp = _require_jax()
     from ..core import IndependentEvidenceCPD, NoisyOrCPD
+    from ..core.cpd.noisy_and import NoisyAndCPD
 
     bn = arg.bn
     nodes = list(bn.nodes)                         # topological (inputs before node)
@@ -105,6 +106,12 @@ def _forward(arg: "Argument"):
                 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 isinstance(cpd, NoisyAndCPD):
+                # NoisyAnd activations are fixed structural constants (all 1.0 from
+                # the ASPIC compiler); they are NOT free parameters in theta.
+                a = jnp.asarray(cpd.activations)
+                pin = jnp.stack([m[i] for i in ins])
+                m[n] = (1 - cpd.leak) * jnp.prod(a * pin)
             elif _is_and_not_splice(cpd):
                 e, u = ins
                 m[n] = m[e] * (1 - m[u])

{probability_flow-0.3.0 → probability_flow-0.4.1}/probability_flow/aspic/compile.py

@@ -40,6 +40,8 @@ def _upstream(c: "Node") -> 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", []))
+    for srcs, _lr, _leak, _rule in getattr(c, "_conjunctions", []):
+        out.extend(srcs)
     return out
 
 
@@ -86,7 +88,12 @@ def _validate(claims: list["Node"]) -> None:
         if not isinstance(c, _Claim):
             continue
 
-        if c.role == "conclusion" and not c._edges and not c._strict:
+        if (
+            c.role == "conclusion"
+            and not c._edges
+            and not c._strict
+            and not getattr(c, "_conjunctions", [])
+        ):
             warnings.warn(
                 f"{c.name!r} is a conclusion with no incoming argument; a leaf "
                 "should be a Premise.",
@@ -160,21 +167,32 @@ def _lower(c: "_Claim") -> None:
 
     defeasible = [(resolve(src), lr) for src, lr, _kind in c._edges]
     strict = [resolve(src) for src in c._strict]
+    conjunctions = getattr(c, "_conjunctions", [])
+
+    # `host` carries the defeasible inputs: the conclusion itself, or a hidden `D`
+    # when strict edges divorce it (parent-divorcing).
+    host = Node(f"{c.name}/defeasible", prior=c.prior) if strict else c
+
+    for src, lr in defeasible:
+        host.add_input(src, lr=lr)
+
+    # Conjunctive support: one hidden NoisyAnd gate per group, fed into `host` as a
+    # single lr support. The gate is the logical AND (each source at activation 1,
+    # firing iff all hold); `leak` is the rule's fallibility; `lr` is its strength.
+    for k, (srcs, lr, leak, _rule) in enumerate(conjunctions):
+        gate = Node(f"{c.name}/conj[{k}]")
+        gate.noisy_and(leak=leak)
+        for src in srcs:
+            gate.add_input(resolve(src), activation=1.0)
+        host.add_input(gate, lr=lr)
+
+    _apply_correlations(c, host, resolve)
 
     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)
-        _apply_correlations(c, d, resolve)
         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)
-        _apply_correlations(c, c, resolve)
+        c.add_input(host, activation=1.0)
 
 
 def _apply_correlations(c: "_Claim", host: "Node", resolve) -> None: