probability-flow 0.1.0__tar.gz → 0.2.0__tar.gz

{probability_flow-0.1.0 → probability_flow-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: probability-flow
-Version: 0.1.0
+Version: 0.2.0
 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
@@ -237,12 +237,15 @@ target posterior is hit by calibrating the root prior. See `docs/generation.md`.
 ## Visualization (optional, `[viz]`)
 
 With the `[viz]` extra installed, a compiled network and an argument both render to
-a matplotlib figure with an in-house layered layout (likelihood-ratio edges
-coloured red→blue by their LR):
+a matplotlib figure with an in-house layout — `radial` by default (the root at the
+centre, the d-separated evidence branches fanning out under a force model), or
+`layout="layered"` for the root-at-the-bottom tree. Likelihood-ratio edges are
+coloured red→blue by their LR, and a long node name shrinks its font to fit:
 
 ```python
-bn.render()                  # or render(bn) from probability_flow.visualization
-guilty.assemble().render()   # the argument view
+bn.render()                       # radial by default
+bn.render(layout="layered")       # root-at-bottom tree (orientation="horizontal" too)
+guilty.assemble().render()        # the argument view (same layouts)
 ```
 
 matplotlib is imported lazily, only when you draw, so importing `probability_flow`
@@ -286,10 +289,12 @@ Working today: the build/compile flow, all four distributions, both solvers with
 linear-time structured messages, evidence, the ASPIC argument-compilation layer,
 argument serialization, the metrics seam (d-separation grouping, depth/size,
 loopiness, difficulty, manipulability), a random argument generator with structural
-and difficulty controls, optional matplotlib renderers, and optional JAX-based
-posterior calibration and parameter sensitivities. Planned (see `docs/ROADMAP.md`):
-a core-network serializer, the loopy-BP "topology zoo" robustness harness, and the
-exact manipulability range.
+and difficulty controls, optional matplotlib renderers (radial and layered
+layouts), and optional JAX-based posterior calibration and parameter sensitivities.
+Loopy BP, the metrics seam, and the renderer are validated against the exact solver
+on non-tree (shared-parent) graphs by the "topology zoo" harness
+(`tools/topology_zoo.py`, `docs/topology_zoo.md`). Planned (see `docs/ROADMAP.md`):
+a core-network serializer and the exact manipulability range.
 
 ## Learning more
 

{probability_flow-0.1.0 → probability_flow-0.2.0}/README.md

@@ -203,12 +203,15 @@ target posterior is hit by calibrating the root prior. See `docs/generation.md`.
 ## Visualization (optional, `[viz]`)
 
 With the `[viz]` extra installed, a compiled network and an argument both render to
-a matplotlib figure with an in-house layered layout (likelihood-ratio edges
-coloured red→blue by their LR):
+a matplotlib figure with an in-house layout — `radial` by default (the root at the
+centre, the d-separated evidence branches fanning out under a force model), or
+`layout="layered"` for the root-at-the-bottom tree. Likelihood-ratio edges are
+coloured red→blue by their LR, and a long node name shrinks its font to fit:
 
 ```python
-bn.render()                  # or render(bn) from probability_flow.visualization
-guilty.assemble().render()   # the argument view
+bn.render()                       # radial by default
+bn.render(layout="layered")       # root-at-bottom tree (orientation="horizontal" too)
+guilty.assemble().render()        # the argument view (same layouts)
 ```
 
 matplotlib is imported lazily, only when you draw, so importing `probability_flow`
@@ -252,10 +255,12 @@ Working today: the build/compile flow, all four distributions, both solvers with
 linear-time structured messages, evidence, the ASPIC argument-compilation layer,
 argument serialization, the metrics seam (d-separation grouping, depth/size,
 loopiness, difficulty, manipulability), a random argument generator with structural
-and difficulty controls, optional matplotlib renderers, and optional JAX-based
-posterior calibration and parameter sensitivities. Planned (see `docs/ROADMAP.md`):
-a core-network serializer, the loopy-BP "topology zoo" robustness harness, and the
-exact manipulability range.
+and difficulty controls, optional matplotlib renderers (radial and layered
+layouts), and optional JAX-based posterior calibration and parameter sensitivities.
+Loopy BP, the metrics seam, and the renderer are validated against the exact solver
+on non-tree (shared-parent) graphs by the "topology zoo" harness
+(`tools/topology_zoo.py`, `docs/topology_zoo.md`). Planned (see `docs/ROADMAP.md`):
+a core-network serializer and the exact manipulability range.
 
 ## Learning more
 

{probability_flow-0.1.0 → probability_flow-0.2.0}/probability_flow/aspic/visualization.py

@@ -37,15 +37,21 @@ def render_argument(
     orientation: str = "vertical",
     title: Optional[str] = None,
     path: Optional[str] = None,
+    layout: str = "radial",
 ):
     """Draw the argument rooted at `target`. Support / rebut / undermine edges are
     coloured by their likelihood ratio on the same confirming/disconfirming scale
     as the Bayesian-network view (strict is bold black, undercut purple-dotted onto
     the edge it attacks). Each claim is annotated with its compiled belief
     `P(=1 | evidence)`; pass `beliefs=False` to skip. `colorbar` / `legend` add the
-    LR scale and the colour key (default off)."""
+    LR scale and the colour key (default off). `layout="radial"` puts the target at
+    the centre with the d-separated branches fanning outward (default "layered")."""
     _require_matplotlib()
     claims = [c for c in _reachable(target) if isinstance(c, _Claim)]
+    if layout == "radial":
+        from ..visualization.layout import radial_positions
+        positions = {**radial_positions(target.compile(), target),
+                     **(positions or {})}
     all_lrs = [lr for c in claims for _src, lr, _kind in c._edges]
     color_for_lr, _ = _lr_colormap(all_lrs)
 

{probability_flow-0.1.0 → probability_flow-0.2.0}/probability_flow/metrics/__init__.py

@@ -20,6 +20,7 @@ and docs/metrics.md.
 """
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING, Optional
 
 from ..core import ExactSolver
@@ -43,7 +44,20 @@ def posterior(bn_or_target, node: Optional["Node"] = None, evidence=None, *,
     target = node if node is not None else default
     if target is None:
         raise ValueError("posterior needs a target node")
-    return solver(bn).prob(target, 1, evidence=evidence)
+    s = solver(bn)
+    p = s.prob(target, 1, evidence=evidence)
+    # `ExactSolver` is always exact; an iterative solver (`LoopySolver`, the "for
+    # scale" path) is exact only on a polytree and merely a tight approximation on
+    # a loopy graph — and on a graph it cannot settle, a silently non-converged
+    # value is worse than that. Surface it; the flag was otherwise unread.
+    if getattr(s, "last_converged", True) is False:
+        warnings.warn(
+            f"the iterative solver did not converge in {getattr(s, 'last_iters', '?')} "
+            "iterations, so this posterior is an unconverged approximation; raise "
+            "max_iters or add damping, or use the exact solver on a smaller graph.",
+            stacklevel=2,
+        )
+    return p
 
 
 __all__ = [

{probability_flow-0.1.0 → probability_flow-0.2.0}/probability_flow/visualization/image.py

@@ -1,12 +1,18 @@
-"""matplotlib rendering: an in-house layered (DAG) layout and the network
-renderer. matplotlib is imported lazily inside the drawing functions, so importing
-this module (or `probability_flow`) never pulls it in.
-
-Layout: the root (sink) sits at the bottom and the graph grows upward to the
-leaves. A node's x is the mean of its inputs' x, so each node's inputs sit
-centred over it (a tidy tree, no edges reaching across the figure). Nodes are
-drawn as data-coordinate patches (circle, or box for a leaf premise), so their
-size relative to the spacing is fixed and arrowheads land exactly on the border.
+"""matplotlib rendering: the layered (DAG) layout and the network renderer.
+matplotlib is imported lazily inside the drawing functions, so importing this
+module (or `probability_flow`) never pulls it in.
+
+Two layouts: the default `radial` (root at the centre, evidence radiating outward,
+see `layout.py`) and the `layered` one here (`_positions`: root at the bottom, the
+graph growing upward; a node's x is the mean of its inputs' x, so its inputs sit
+centred over it). Nodes are drawn as data-coordinate patches (circle, or box for a
+leaf premise), so their size relative to the spacing is fixed and arrowheads land
+exactly on the border.
+
+A node label is fit to the node by shrinking its font down to `MIN_FONT` (see
+`_fit_label`); only a name too long to fit even then is truncated with an ellipsis.
+Names still read best kept short — the full claim text belongs in `desc`, not the
+name.
 """
 from __future__ import annotations
 
@@ -21,12 +27,12 @@ if TYPE_CHECKING:
     from ..core.node import Node
 
 # geometry, in data units (the axes use equal aspect)
-R_NODE = 0.72            # node radius / box half-extent
-X_SPACING = 2.1          # horizontal gap between adjacent leaves
-Y_SPACING = 2.5          # vertical gap between layers
+R_NODE = 0.80            # node radius / box half-extent
+X_SPACING = 2.2          # horizontal gap between adjacent leaves
+Y_SPACING = 2.6          # vertical gap between layers
 IN_PER_UNIT = 0.62       # figure inches per data unit
-FONT = 6.5               # default label point size (the BN renderer's)
-WRAP = 11                # characters before a label line wraps
+FONT = 6.5               # preferred label point size; shrinks (to MIN_FONT) to fit
+MIN_FONT = 3.6           # smallest label font before a name is truncated instead
 DPI = 200                # saved-image resolution
 _NEUTRAL_EDGE = "#9aa0a6"
 _OBSERVED_EDGE = "#c77d00"   # border for an observed (evidence) node
@@ -48,10 +54,18 @@ def _require_matplotlib():
 
 
 def marginals(bn, evidence=None) -> dict:
-    """`P(node = 1 | evidence)` for every node, via the exact solver (small graphs
-    only). With `evidence=None` these are the prior-propagated marginals."""
-    from ..core import ExactSolver
-
+    """`P(node = 1 | evidence)` for every node, for the node labels. Polytree-aware,
+    like the `Argument` handle: one exact `LoopySolver` propagation on a polytree
+    (linear, every node in a single pass), the brute-force `ExactSolver` only on a
+    loop (small graphs — it re-enumerates `2**n` per node, so it is what made
+    rendering a large graph hang). With `evidence=None` these are the
+    prior-propagated marginals."""
+    from ..core import ExactSolver, LoopySolver
+    from ..metrics import is_polytree
+
+    if is_polytree(bn):
+        m = LoopySolver(bn).marginals(evidence)
+        return {n: float(m[n][1]) for n in bn.nodes}
     solver = ExactSolver(bn)
     return {n: solver.prob(n, 1, evidence=evidence) for n in bn.nodes}
 
@@ -107,6 +121,28 @@ def _positions(nodes: Sequence, edges: Sequence[dict],
         if n not in x:
             x[n] = counter[0] * X_SPACING
             counter[0] += 1
+
+    # De-collide within each layer. The mean-of-inputs rule places a node centred
+    # over its parents, which is tidy for a tree but collapses distinct nodes onto
+    # the same coordinate once parents are *shared* (a non-polytree): two siblings
+    # that share an input get the same mean and draw on top of each other. Sweep
+    # each layer left to right enforcing a minimum gap, then recentre so the tidy
+    # shape is preserved. A graph with no overlaps (every tree) is unchanged.
+    by_layer: dict = defaultdict(list)
+    for n in nodes:
+        by_layer[layer[n]].append(n)
+    for ns in by_layer.values():
+        if len(ns) < 2:
+            continue
+        ns.sort(key=lambda n: (x[n], getattr(n, "id", id(n))))
+        before = sum(x[n] for n in ns) / len(ns)
+        for i in range(1, len(ns)):
+            if x[ns[i]] - x[ns[i - 1]] < X_SPACING:
+                x[ns[i]] = x[ns[i - 1]] + X_SPACING
+        shift = before - sum(x[n] for n in ns) / len(ns)
+        for n in ns:
+            x[n] += shift
+
     if orientation == "horizontal":       # root (layer 0) on the right, leaves left
         return {n: (-layer[n] * Y_SPACING, x[n]) for n in nodes}
     return {n: (x[n], layer[n] * Y_SPACING) for n in nodes}
@@ -121,24 +157,125 @@ def _boundary(center, toward, marker, r=R_NODE):
     return (center[0] + ux * t, center[1] + uy * t)
 
 
-def _wrap(text: str) -> str:
-    """Wrap a label to the node width. `->` and `/` (common in auto-named helper
-    nodes) become break opportunities, and words are never split mid-token."""
-    out = []
-    for line in text.split("\n"):
-        line = line.replace("->", " → ").replace("/", " / ")
-        out.append(textwrap.fill(line, WRAP, break_long_words=False))
-    return "\n".join(out)
+def _node_pt_radius(ax, fig, x, y) -> float:
+    """`R_NODE`'s on-screen size in points, so label fitting works at whatever scale
+    the figure ends up (the auto-sized canvas, or a big user-supplied one)."""
+    (x0, _), (x1, _) = ax.transData.transform([(x, y), (x + R_NODE, y)])
+    return abs(x1 - x0) * 72.0 / float(fig.dpi)
+
+
+def _split_label(text: str):
+    """Separate a node label into its name and a trailing `P=…` / `prior=…` belief
+    line (kept intact, drawn below the name)."""
+    lines = text.split("\n")
+    if len(lines) >= 2 and "=" in lines[-1]:
+        return "\n".join(lines[:-1]), lines[-1]
+    return text, ""
+
+
+def _wrap_to(name: str, cols: int) -> list[str]:
+    name = name.replace("->", " → ").replace("/", " / ")
+    return textwrap.fill(name, max(4, cols), break_long_words=False).split("\n")
+
+
+# a serif glyph is ~0.52 of the point size wide; lines sit ~1.2 point sizes apart.
+_CHAR_W, _LINE_H = 0.52, 1.2
+
+
+def _fit_label(ax, fig, x, y, text: str, is_box: bool):
+    """Fit a node label inside the node by shrinking the font (down to `MIN_FONT`)
+    so a long name stays readable rather than overflowing; only a name too long even
+    at `MIN_FONT` is truncated with an ellipsis. Returns `(text, fontsize)`."""
+    name, belief = _split_label(text)
+    half = _node_pt_radius(ax, fig, x, y) * (0.92 if is_box else 0.74)  # usable radius
+    extra = 1 if belief else 0
+
+    f = FONT
+    while f >= MIN_FONT - 1e-9:
+        cols = max(4, int(2 * half / (_CHAR_W * f)))
+        lines = _wrap_to(name, cols)
+        widest = max((len(s) for s in lines + ([belief] if belief else [])), default=0)
+        if ((len(lines) + extra) * _LINE_H * f <= 2 * half
+                and widest * _CHAR_W * f <= 2 * half):
+            break
+        f -= 0.5
+    else:                                   # doesn't fit even at MIN_FONT: truncate
+        f = MIN_FONT
+        cols = max(4, int(2 * half / (_CHAR_W * f)))
+        rows = max(1, int(2 * half / (_LINE_H * f)) - extra)
+        lines = _wrap_to(name, cols)
+        if len(lines) > rows:
+            lines = lines[:rows]
+            lines[-1] = lines[-1][:cols - 1].rstrip() + "…"
+
+    return "\n".join(lines + ([belief] if belief else [])), f
+
+
+_CLEAR_BUFFER = 1.45      # an edge should clear a node's centre by this * R_NODE
+_CURVE_TRIES = (0.12, 0.2, 0.3, 0.42, 0.58)   # arc3 rad magnitudes, smallest first
 
 
-def _arrow(ax, start, end, *, color, linestyle, lw):
-    """Draw an arrow. A dashed/dotted edge is split into a dashed shaft and a
-    separate SOLID arrowhead, so the head never looks dashed."""
+def _ctrl_for_rad(start, end, rad):
+    """The quadratic control point matplotlib's `arc3,rad` uses: the chord midpoint
+    displaced by `rad * (dy, -dx)` (its right normal). Computing it ourselves lets
+    the collision check below match the curve that actually gets drawn."""
+    mx, my = (start[0] + end[0]) / 2, (start[1] + end[1]) / 2
+    dx, dy = end[0] - start[0], end[1] - start[1]
+    return (mx + rad * dy, my - rad * dx)
+
+
+def _bezier_point(start, ctrl, end, t):
+    mt = 1.0 - t
+    return (mt * mt * start[0] + 2 * mt * t * ctrl[0] + t * t * end[0],
+            mt * mt * start[1] + 2 * mt * t * ctrl[1] + t * t * end[1])
+
+
+def _curve_clears(start, ctrl, end, obstacles, buffer):
+    pts = [_bezier_point(start, ctrl, end, i / 18.0) for i in range(19)]
+    for q, r in obstacles:
+        if min(math.hypot(q[0] - x, q[1] - y) for x, y in pts) < buffer * r:
+            return False
+    return True
+
+
+def _route_rad(start, end, obstacles, *, buffer=_CLEAR_BUFFER):
+    """Pick an `arc3` rad that bows the edge clear of any node its straight line
+    would pass through. `0.0` when the straight edge is already clear (the common
+    case, so trees and short edges are untouched). Bows away from the side the
+    blocking nodes sit on, taking the smallest curvature that clears them."""
+    dx, dy = end[0] - start[0], end[1] - start[1]
+    L2 = dx * dx + dy * dy or 1.0
+    blockers = []
+    for q, r in obstacles:
+        t = ((q[0] - start[0]) * dx + (q[1] - start[1]) * dy) / L2
+        if not (0.02 < t < 0.98):                 # only nodes alongside the shaft
+            continue
+        px, py = start[0] + t * dx, start[1] + t * dy
+        if math.hypot(q[0] - px, q[1] - py) < buffer * r:
+            blockers.append(dx * (q[1] - start[1]) - dy * (q[0] - start[0]))
+    if not blockers:
+        return 0.0
+    # `(dy, -dx)` is the chord's right normal, so a positive rad bows right; a node
+    # on the left (positive cross product) is escaped by bowing right.
+    sign = 1.0 if sum(blockers) >= 0 else -1.0
+    for mag in _CURVE_TRIES:
+        rad = mag * sign
+        if _curve_clears(start, _ctrl_for_rad(start, end, rad), end, obstacles,
+                         buffer * 0.92):
+            return rad
+    return _CURVE_TRIES[-1] * sign                 # best effort if nothing fully clears
+
+
+def _arrow(ax, start, end, *, color, linestyle, lw, rad=0.0):
+    """Draw an arrow, optionally bowed by `rad` (matplotlib `arc3` curvature) to
+    route around a node. A dashed/dotted edge is split into a dashed shaft and a
+    separate SOLID arrowhead, so the head never looks dashed (kept straight)."""
     from matplotlib.patches import FancyArrowPatch
 
     if linestyle in ("-", "solid", None):
         ax.add_patch(FancyArrowPatch(start, end, arrowstyle="-|>", mutation_scale=15,
-                                     shrinkA=0, shrinkB=0, lw=lw, color=color, zorder=1))
+                                     shrinkA=0, shrinkB=0, lw=lw, color=color,
+                                     connectionstyle=f"arc3,rad={rad}", zorder=1))
         return
     dx, dy = end[0] - start[0], end[1] - start[1]
     dist = math.hypot(dx, dy) or 1.0
@@ -211,9 +348,16 @@ def draw_layered(
             _, ax = plt.subplots(figsize=(w, h))
         fig = ax.figure
         ax.set_aspect("equal")
+        pad = R_NODE + 0.28
+        ax.set_xlim(min(xs) - pad, max(xs) + pad)
+        ax.set_ylim(min(ys) - pad, max(ys) + pad)
+        ax.axis("off")
+        fig.canvas.draw()                 # finalise the transform so label fitting
+        #                                   sees the real on-screen node size
 
         for e in edges:
             cu = pos[e["u"]]
+            rad = 0.0
             if "attacks" in e:                         # arrow lands on another edge
                 a, b = e["attacks"]
                 end = ((pos[a][0] + pos[b][0]) / 2, (pos[a][1] + pos[b][1]) / 2)
@@ -225,12 +369,15 @@ def draw_layered(
                 gdx, gdy = start[0] - end[0], start[1] - end[1]
                 gd = math.hypot(gdx, gdy) or 1.0
                 end = (end[0] + gdx / gd * 0.06, end[1] + gdy / gd * 0.06)  # small gap
+                # bow the (solid) edge around any node it would otherwise cross.
+                obstacles = [(pos[n], R_NODE) for n in nodes
+                             if n is not e["u"] and n is not e["v"]]
+                rad = _route_rad(start, end, obstacles)
             _arrow(ax, start, end, color=e.get("color", _NEUTRAL_EDGE),
-                   linestyle=e.get("style", "-"), lw=e.get("width", 1.7))
+                   linestyle=e.get("style", "-"), lw=e.get("width", 1.7), rad=rad)
             if e.get("label"):
                 f = e.get("label_pos", 0.5)             # fraction from start to end
-                mx = start[0] + f * (end[0] - start[0])
-                my = start[1] + f * (end[1] - start[1])
+                mx, my = _bezier_point(start, _ctrl_for_rad(start, end, rad), end, f)
                 ax.text(mx, my, e["label"], fontsize=font - 1, ha="center", va="center",
                         zorder=3, bbox=dict(boxstyle="round,pad=0.12", fc="white",
                                             ec="none", alpha=0.85))
@@ -247,13 +394,9 @@ def draw_layered(
                     ec=ec, lw=lw, zorder=2))
             else:
                 ax.add_patch(Circle((x, y), R_NODE, fc=color_of(n), ec=ec, lw=lw, zorder=2))
-            ax.text(x, y, _wrap(label_of(n)), fontsize=font, ha="center", va="center",
-                    zorder=4)
+            txt, fsz = _fit_label(ax, fig, x, y, label_of(n), marker_of(n) == "s")
+            ax.text(x, y, txt, fontsize=fsz, ha="center", va="center", zorder=4)
 
-        pad = R_NODE + 0.28
-        ax.set_xlim(min(xs) - pad, max(xs) + pad)
-        ax.set_ylim(min(ys) - pad, max(ys) + pad)
-        ax.axis("off")
         if title:
             ax.set_title(title, fontsize=font + 5, pad=3)
 
@@ -355,19 +498,26 @@ def _legend_entries(cpds) -> list:
 def render(bn, values: Optional[Mapping["Node", float]] = None, evidence=None,
            positions: Optional[Mapping] = None, colorbar: bool = False,
            legend: bool = False, ax=None, orientation: str = "vertical",
-           title: Optional[str] = None, path: Optional[str] = None):
-    """Draw a compiled `BayesianNetwork`. By default the root sits at the bottom;
-    pass `orientation="horizontal"` to lay it out left to right with the root on
-    the right.
+           title: Optional[str] = None, path: Optional[str] = None,
+           layout: str = "radial"):
+    """Draw a compiled `BayesianNetwork`. `layout="radial"` (default) puts the root
+    at the centre with the d-separated evidence branches fanning outward (see
+    `layout.py`); `layout="layered"` puts the root at the bottom and grows upward,
+    and `orientation="horizontal"` lays that form out left to right.
 
     A node with no parents shows its `prior`; a node with parents (or any node
     under evidence) shows its posterior `P`. Observed nodes get a bold amber
     border. Pass `values={}` to skip the numbers. Likelihood-ratio edges are
     always coloured by their LR (red disconfirming, blue confirming); set
-    `colorbar=True` for the LR scale and `legend=True` for the node-colour key."""
+    `colorbar=True` for the LR scale and `legend=True` for the node-colour key.
+    A long node name shrinks its font to fit the node (and truncates only if it
+    still will not fit); names read best kept short, with the full text in `desc`."""
     _require_matplotlib()
     if values is None:
         values = marginals(bn, evidence)
+    if layout == "radial":
+        from .layout import radial_positions
+        positions = {**radial_positions(bn), **(positions or {})}
     cpd_of = {n: bn.compiled_cpd(n).cpd for n in bn.nodes}
 
     all_lrs = [edge_lr(cpd_of[n], inp) for n in bn.nodes for inp in cpd_of[n].inputs]

probability_flow-0.2.0/probability_flow/visualization/layout.py

@@ -0,0 +1,166 @@
+"""Radial layout: the root at the centre, evidence radiating outward.
+
+An alternative to the layered (root-at-bottom) layout in `image._positions`. The
+root sits at the origin and radius grows with depth, so the root repels everything
+outward and depths separate into rings. The angle comes from a radial-tree
+assignment driven by **d-separation**: each independent evidence branch (a
+d-separated group) gets its own angular sector with a gap between sectors, leaves
+spread across the sector in DFS order so siblings stay adjacent, and every internal
+node sits at the circular mean of its cone's leaves — a node shared across branches
+naturally bridges the sectors it spans. A force relaxation then settles related
+nodes (argument edges and co-input siblings) at an equilibrium distance — pushed
+apart when too close, pulled together when too far — so each branch clusters
+instead of spreading maximally, while the radius stays pinned to the depth ring so
+the root remains centred.
+
+`radial_positions(bn, target)` returns `{node: (x, y)}` for feeding the renderer
+via `render(..., layout="radial")`.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .image import _layers
+
+
+def _sink(bn):
+    """The root: the one node that feeds nothing (no outgoing edge)."""
+    used = {i for n in bn.nodes for i in bn.compiled_cpd(n).inputs}
+    sinks = [n for n in bn.nodes if n not in used]
+    return sinks[0] if sinks else bn.nodes[-1]
+
+
+def _radial_angles(bn, target, inputs_of, gap_frac):
+    from ..metrics import d_separated_groups
+    from ..metrics._util import ancestors
+
+    groups = d_separated_groups(bn, target)
+
+    def dfs_leaves(parents):
+        seen, order = set(), []
+
+        def rec(n):
+            if n in seen:
+                return
+            seen.add(n)
+            if not inputs_of[n]:
+                order.append(n)
+            else:
+                for i in inputs_of[n]:
+                    rec(i)
+        for p in parents:
+            rec(p)
+        return order
+
+    group_leaves = [dfs_leaves(list(g.parents)) for g in groups]
+    total = sum(len(gl) for gl in group_leaves) or 1
+    ng = max(1, len(groups))
+    # gaps separate sectors only with more than one group; a single (fully shared)
+    # group gets the whole circle.
+    gap = 0.0 if ng <= 1 else gap_frac * 2 * np.pi / ng
+    usable = 2 * np.pi - gap * ng
+
+    leaf_angle: dict = {}
+    cur = 0.0
+    for gl in group_leaves:
+        span = usable * (len(gl) / total)
+        for k, leaf in enumerate(gl):
+            leaf_angle[leaf] = cur + span * ((k + 0.5) / max(1, len(gl)))
+        cur += span + gap
+
+    ang: dict = {}
+    for n in bn.nodes:
+        cone = ancestors(bn, n) | {n}
+        ls = [c for c in cone if not inputs_of[c] and c in leaf_angle]
+        if ls:
+            xs = float(np.mean([np.cos(leaf_angle[c]) for c in ls]))
+            ys = float(np.mean([np.sin(leaf_angle[c]) for c in ls]))
+            ang[n] = float(np.arctan2(ys, xs))
+        else:
+            ang[n] = 0.0
+    return ang
+
+
+def _relax(pos, R, root_i, ea, eb, sa, sb, mod, k_anchor, *, iters, L0,
+           k_rep=2.2, k_spring=0.7, k_sib=0.45):
+    """Equilibrium-distance relaxation on the angle. All pairs repel (short range,
+    harder across d-separated groups via `mod`); related pairs — argument edges
+    (`ea→eb`) and co-input siblings (`sa,sb`) — sit on springs with rest length
+    `L0`, so they are pushed apart when closer and pulled together when farther,
+    which clusters a branch instead of spreading it maximally. A weak anchor to the
+    radial-tree start (`pos0`) keeps a connected blob from collapsing off-centre,
+    and the radius is re-snapped to the depth ring each step so all of this acts
+    angularly and the root stays at the centre."""
+    keep = R > 0
+    pos0 = pos.copy()
+    for it in range(iters):
+        diff = pos[:, None, :] - pos[None, :, :]
+        d2 = (diff ** 2).sum(-1) + 1e-3
+        rep = k_rep * mod / d2
+        np.fill_diagonal(rep, 0.0)
+        disp = (rep[..., None] * diff / np.sqrt(d2)[..., None]).sum(1)
+        for ia, ib, k in ((ea, eb, k_spring), (sa, sb, k_sib)):
+            if len(ia):
+                dv = pos[ib] - pos[ia]
+                L = np.hypot(dv[:, 0], dv[:, 1]) + 1e-9
+                f = (k * (L - L0) / L)[:, None] * dv
+                np.add.at(disp, ia, f)
+                np.add.at(disp, ib, -f)
+        disp += k_anchor * (pos0 - pos)
+        dv = disp * (0.1 * (1.0 - 0.8 * it / iters))
+        dn = np.hypot(dv[:, 0], dv[:, 1]) + 1e-9
+        pos = pos + dv * np.minimum(1.0, 0.5 / dn)[:, None]
+        pos[root_i] = [0.0, 0.0]
+        r = np.hypot(pos[:, 0], pos[:, 1]) + 1e-9
+        pos[keep] = (pos[keep] / r[keep, None]) * R[keep, None]
+    return pos
+
+
+def radial_positions(bn, target=None, *, ring=2.9, relax=True, iters=500,
+                     gap_frac=0.22, k_anchor=0.18, dsep_rep=1.8):
+    """`{node: (x, y)}` for a root-centred radial layout. `target` defaults to the
+    network's sink (its root). With `relax`, related nodes settle at an equilibrium
+    distance (clustering each branch) while the root stays centred and depths stay
+    ringed. Deterministic."""
+    from ..metrics import d_separated_groups
+
+    target = target if target is not None else _sink(bn)
+    nodes = list(bn.nodes)
+    idx = {n: i for i, n in enumerate(nodes)}
+    inputs_of = {n: list(bn.compiled_cpd(n).inputs) for n in nodes}
+    edge_dicts = [{"u": inp, "v": n} for n in nodes for inp in inputs_of[n]]
+    depth = np.array([float(_layers(nodes, edge_dicts)[n]) for n in nodes])
+
+    ang = _radial_angles(bn, target, inputs_of, gap_frac)
+    R = depth * ring
+    pos = np.array([[R[idx[n]] * np.cos(ang[n]), R[idx[n]] * np.sin(ang[n])]
+                    for n in nodes])
+    root_i = idx[target]
+    pos[root_i] = [0.0, 0.0]
+
+    if relax:
+        groups = d_separated_groups(bn, target)
+        ng = max(1, len(groups))
+        # many independent branches balance themselves around the root, so a light
+        # anchor lets them cluster; few (a densely shared blob) have nothing to fan,
+        # so anchor hard to the radial-tree to keep the root central.
+        anchor = k_anchor * (1.0 + 8.0 / ng ** 2)
+        grp = np.full(len(nodes), -1)
+        for gi, g in enumerate(groups):
+            for n in g.nodes:
+                if grp[idx[n]] == -1:
+                    grp[idx[n]] = gi
+        same = grp[:, None] == grp[None, :]
+        known = (grp[:, None] >= 0) & (grp[None, :] >= 0)
+        mod = np.where(known & ~same, dsep_rep, 1.0)
+
+        ea = np.array([idx[i] for n in nodes for i in inputs_of[n]], dtype=int)
+        eb = np.array([idx[n] for n in nodes for _ in inputs_of[n]], dtype=int)
+        sibs = {(min(idx[i], idx[j]), max(idx[i], idx[j]))
+                for n in nodes for a, i in enumerate(inputs_of[n])
+                for j in inputs_of[n][a + 1:]}
+        sa = np.array([s[0] for s in sibs], dtype=int)
+        sb = np.array([s[1] for s in sibs], dtype=int)
+        pos = _relax(pos, R, root_i, ea, eb, sa, sb, mod, anchor,
+                     iters=iters, L0=ring)
+    return {n: (float(pos[idx[n], 0]), float(pos[idx[n], 1])) for n in nodes}

{probability_flow-0.1.0 → probability_flow-0.2.0}/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "probability-flow"
-version = "0.1.0"
+version = "0.2.0"
 description = "A from-scratch, modular discrete Bayesian-network library."
 readme = "README.md"
 requires-python = ">=3.11"