Plate graphs

plate_graph

Plate-graph extraction for QVR programs.

A plate diagram in the Bayesian-statistics tradition has three ingredients:

• Nodes: random variables (sample sites, observe sites) and deterministic intermediate values (let-bound terms). Observed nodes are conventionally shaded; latent nodes are unshaded; deterministic nodes are conventionally drawn as smaller dots.
• Edges: directed dependencies between nodes, computed from the family-argument list of each draw / observe step. z ~ Categorical(probs=theta) produces an edge from theta to z.
• Plates: rectangular regions enclosing groups of nodes that share a common indexing axis. A plate is induced by an index annotation on a draw (sample theta : Doc <- ... puts theta on the Doc plate), by iid_over on the same step, by the over axis on a marginalize block (which scopes every inner observation), or by the grouping fibration [via=...] on an observe step.

The extractor in this module walks a compiled :class:quivers.dsl.compiler.Compiler's program declarations and builds an immutable :class:PlateGraph describing exactly this structure. Renderers (in-TUI table view, Mermaid output, Daft / TikZ / DOT code generation) consume the graph; the graph itself is the single source of truth and never embeds layout decisions.

The walker handles arbitrarily nested marginalize blocks by treating each block as a new plate scope that lives inside the enclosing plates. [via=idx] on an observe inside a marginalize-over-G block puts the observe on G's plate in addition to its own index plate, matching the standard convention for grouped observations in pgm notation.

NodeKind module-attribute

NodeKind = Literal['latent', 'observed', 'marginalized', 'deterministic']

The kind tag of a PlateNode, which selects how a renderer draws it.

Edge didactic-model

Edge(*, src: str, dst: str)

Bases: Model

One directed dependency edge in the plate diagram.

Plate didactic-model

Plate(*, name: str, cardinality: int | None = None, parent: str | None = None)

Bases: Model

One plate (indexing axis) of the diagram.

ATTRIBUTE	DESCRIPTION
name	The plate object's name (Doc, Topic, Word). Multiple plates can share a name when the same axis is indexed at multiple levels in the program; the extractor de-duplicates them. TYPE: str
cardinality	The plate's size if statically inferable from the module's object NAME : FinSet N declarations; None otherwise. TYPE: int | None
parent	Name of an enclosing plate, or None when this plate is at the program's outermost level. Set when a marginalize block nests inside another marginalize. TYPE: str | None

PlateNode didactic-model

PlateNode(*, name: str, kind: NodeKind, family: str | None = None, plates: tuple[str, ...] = (), args: tuple[str, ...] = (), scope_path: str = '')

Bases: Model

One random variable / deterministic intermediate.

ATTRIBUTE	DESCRIPTION
name	The variable's bound name in the program (theta, z, w, mu). TYPE: str
kind	"latent" for unobserved samples, "observed" for observe sites (clamped at runtime), "marginalized" for sample sites that are integrated out by an enclosing marginalize block, "deterministic" for let bindings whose value is a closed-form function of upstream nodes (no stochastic draw). TYPE: NodeKind
family	The distribution family name ("Normal", "Dirichlet", "Categorical") for stochastic nodes. None for let nodes. TYPE: str | None
plates	The plates this node sits inside, ordered from outermost to innermost. Empty tuple for scalar nodes. TYPE: tuple[str, ...]
args	The variable names this node's family / let-expression references. The graph's edges are derived from these. TYPE: tuple[str, ...]
scope_path	The ::-separated scope path to this binding (e.g. lda::theta, lda::z::w). Used by renderers that link plate-diagram clicks back to :info lookups. TYPE: str

PlateGraph didactic-model

PlateGraph(*, program_name: str, domain: str = '', codomain: str = '', nodes: tuple[PlateNode, ...] = (), plates: tuple[Plate, ...] = (), edges: tuple[Edge, ...] = ())

Bases: Model

Structural plate diagram of a single QVR program.

build_plate_graph

build_plate_graph(compiler, program_name: str) -> PlateGraph | None

Build the :class:PlateGraph for program_name.

Returns None if program_name isn't a registered program on compiler. Walks the program's draws and any nested marginalize body recursively. Marginalized sample sites (the variable bound by the marginalize header) get kind="marginalized" so renderers can shade them differently from a plain "latent".

PARAMETER	DESCRIPTION
compiler	A quivers.dsl.Compiler that has compiled the module holding the program. TYPE: Compiler
program_name	The name of the program to diagram. TYPE: str

RETURNS	DESCRIPTION
PlateGraph or None	The program's plate graph, or None when no program of that name is registered.

Source code in src/quivers/analysis/plate_graph.py

def build_plate_graph(  # type: ignore[no-untyped-def]
    compiler,
    program_name: str,
) -> PlateGraph | None:
    """Build the :class:`PlateGraph` for ``program_name``.

    Returns ``None`` if ``program_name`` isn't a registered program
    on ``compiler``. Walks the program's ``draws`` and any nested
    ``marginalize`` body recursively. Marginalized sample sites
    (the variable bound by the marginalize header) get
    ``kind="marginalized"`` so renderers can shade them
    differently from a plain ``"latent"``.

    Parameters
    ----------
    compiler : Compiler
        A `quivers.dsl.Compiler` that has compiled the module holding
        the program.
    program_name : str
        The name of the program to diagram.

    Returns
    -------
    PlateGraph or None
        The program's plate graph, or ``None`` when no program of that
        name is registered.
    """
    # Locate the program declaration. Parametric programs live on
    # ``compiler.programs`` (the template registry); non-parametric
    # programs compile straight to a morphism and are accessible
    # only through the source module's statement list. Check both.
    programs = getattr(compiler, "programs", {}) or {}
    decl = programs.get(program_name)
    if not isinstance(decl, ProgramDecl):
        module = getattr(compiler, "_module", None)
        for stmt in getattr(module, "statements", ()) or ():
            if (
                isinstance(stmt, ProgramDecl)
                and getattr(stmt, "name", None) == program_name
            ):
                decl = stmt
                break
    if not isinstance(decl, ProgramDecl):
        return None

    nodes: list[PlateNode] = []
    plate_names: dict[str, str | None] = {}  # plate -> parent plate
    edges: list[Edge] = []

    def _walk(  # type: ignore[no-untyped-def]
        steps: tuple[object, ...],
        parent_plates: tuple[str, ...],
        path_prefix: str,
    ) -> None:
        for step in steps:
            if isinstance(step, SampleStep):
                name = step.vars[0] if step.vars else "?"
                plates = _plates_for_step(step, parent_plates)
                for p in plates:
                    if p not in plate_names:
                        plate_names[p] = parent_plates[-1] if parent_plates else None
                args = _step_args(step)
                nodes.append(
                    PlateNode(
                        name=name,
                        kind="latent",
                        family=step.morphism,
                        plates=plates,
                        args=args,
                        scope_path=f"{path_prefix}::{name}",
                    )
                )
                for a in args:
                    edges.append(Edge(src=a, dst=name))
            elif isinstance(step, ObserveStep):
                name = step.vars[0]
                plates = _plates_for_step(step, parent_plates)
                for p in plates:
                    if p not in plate_names:
                        plate_names[p] = parent_plates[-1] if parent_plates else None
                args = _step_args(step)
                nodes.append(
                    PlateNode(
                        name=name,
                        kind="observed",
                        family=step.morphism,
                        plates=plates,
                        args=args,
                        scope_path=f"{path_prefix}::{name}",
                    )
                )
                for a in args:
                    edges.append(Edge(src=a, dst=name))
            elif isinstance(step, MarginalizeStep):
                name = step.var
                plates = _plates_for_step(step, parent_plates)
                for p in plates:
                    if p not in plate_names:
                        plate_names[p] = parent_plates[-1] if parent_plates else None
                args = _step_args(step)
                nodes.append(
                    PlateNode(
                        name=name,
                        kind="marginalized",
                        family=step.morphism,
                        plates=plates,
                        args=args,
                        scope_path=f"{path_prefix}::{name}",
                    )
                )
                for a in args:
                    edges.append(Edge(src=a, dst=name))
                # Recurse into the marginalize body. Inner plates
                # inherit ``parent_plates + [over]`` rather than
                # the marginalize variable's own index plate,
                # because the marginalized variable is integrated
                # out and inner observations don't lie on its axis.
                inner_prefix = f"{path_prefix}::{name}"
                inner_plates = _child_plates_for_marginalize(step, parent_plates)
                _walk(step.scope, inner_plates, inner_prefix)
            elif isinstance(step, LetStep):
                name = step.name
                args = _let_value_args(step.value)
                nodes.append(
                    PlateNode(
                        name=name,
                        kind="deterministic",
                        family=None,
                        plates=parent_plates,
                        args=args,
                        scope_path=f"{path_prefix}::{name}",
                    )
                )
                for a in args:
                    edges.append(Edge(src=a, dst=name))
            elif isinstance(step, ReturnStep):
                # Return doesn't introduce a node; downstream
                # consumers see the return vars via the program
                # decl.
                continue

    _walk(decl.draws, (), program_name)

    plates_out: list[Plate] = []
    for p, parent in plate_names.items():
        plates_out.append(
            Plate(
                name=p,
                cardinality=_cardinality_lookup(compiler, p),
                parent=parent,
            )
        )

    domain = getattr(decl.domain, "name", "") if decl.domain is not None else ""
    codomain = getattr(decl.codomain, "name", "") if decl.codomain is not None else ""
    return PlateGraph(
        program_name=program_name,
        domain=domain,
        codomain=codomain,
        nodes=tuple(nodes),
        plates=tuple(plates_out),
        edges=tuple(edges),
    )