API reference

The shared objects live at the top level of ergodic, and the pillars under their own modules. Most people arrive here from search, so the table below maps a task to its entry point and the guide that teaches it.

To do this	Start with	Guide
Build a causal graph	dag, admg, cpdag, pag	Causal graphs
Turn a graph into an estimand	identify_effect	Identification
Estimate an effect from data	estimate_effect	Causal inference
Estimate per-unit effects (CATE)	cate	Causal inference
Stress-test an estimate	refute_effect	Checking an estimate
Learn a graph from data	discover, bootstrap	Causal discovery
Measure a rollout with no graph	did, synthetic_control, interrupted_time_series	Panel designs
Analyze a process from its log	process_map, case_table, discover_model	Process intelligence
Simulate or ask counterfactuals	SCM	Structural causal models

Causal graphs

The mixed-graph object and its builders, separation, do-surgery, and interop.

ergodic.graph

Causal graphs: one mixed-graph core with typed faces (DAG, ADMG, CPDAG, PAG).

Build graphs with the :func:dag, :func:admg, :func:cpdag, :func:pag, or :func:mixed helpers. Edges read as glyph strings, so dag(["X -> Y", "Y -> Z"]) and pag(["X o-> Y", "Z o-> Y"]) do what they look like. Separation (:func:m_separated), the orientation machinery (:func:meek_closure, :func:fci_closure), the projections between the faces (:func:cpdag_of_dag, :func:mag_of_dag, :func:pag_of_mag, :func:mag_of_pag), the class toolkit (:func:possible_descendants, :func:visible_edges, :func:dags_of_cpdag, :func:mags_of_pag), and the primitives the identification algorithms read (:func:districts, :func:topological_order, :func:latent_projection) live here too.

EdgeSpec = Edge | str | tuple[Any, ...] module-attribute

How a caller can name an edge: an Edge, a glyph string like "X o-> Y", a 2-tuple (u, v) (read as u -> v), or a 4-tuple (u, u_mark, v, v_mark).

Edge dataclass

One edge between two distinct nodes, with a mark at each end.

Stored canonically with u < v (string order) so an undirected pair maps to exactly one Edge. u_mark is the mark at u, v_mark at v.

Source code in src/ergodic/graph/_graph.py

@dataclass(frozen=True)
class Edge:
    """One edge between two distinct nodes, with a mark at each end.

    Stored canonically with ``u < v`` (string order) so an undirected pair maps
    to exactly one ``Edge``. ``u_mark`` is the mark at ``u``, ``v_mark`` at ``v``.
    """

    u: str
    v: str
    u_mark: Mark
    v_mark: Mark

    @staticmethod
    def make(a: str, a_mark: Mark, b: str, b_mark: Mark) -> Edge:
        """Build an edge between ``a`` and ``b``, normalizing endpoint order."""
        if a == b:
            raise ValueError(f"self-loops are not allowed: {a!r}")
        if a < b:
            return Edge(a, b, a_mark, b_mark)
        return Edge(b, a, b_mark, a_mark)

    def mark_at(self, node: str) -> Mark:
        """Return the mark this edge carries at ``node``."""
        if node == self.u:
            return self.u_mark
        if node == self.v:
            return self.v_mark
        raise KeyError(node)

    def other(self, node: str) -> str:
        """Return the endpoint that is not ``node``."""
        if node == self.u:
            return self.v
        if node == self.v:
            return self.u
        raise KeyError(node)

    @property
    def is_directed(self) -> bool:
        return {self.u_mark, self.v_mark} == {Mark.TAIL, Mark.ARROW}

    @property
    def is_bidirected(self) -> bool:
        return self.u_mark is Mark.ARROW and self.v_mark is Mark.ARROW

    @property
    def is_undirected(self) -> bool:
        return self.u_mark is Mark.TAIL and self.v_mark is Mark.TAIL

    @property
    def has_circle(self) -> bool:
        return Mark.CIRCLE in (self.u_mark, self.v_mark)

    def directed_pair(self) -> tuple[str, str] | None:
        """Return ``(parent, child)`` if this edge is directed, else ``None``."""
        if not self.is_directed:
            return None
        if self.u_mark is Mark.TAIL:
            return (self.u, self.v)
        return (self.v, self.u)

    def glyph(self) -> str:
        """Render the edge, e.g. ``"X o-> Y"``."""
        left = self.u_mark.char(left=True)
        right = self.v_mark.char(left=False)
        return f"{self.u} {left}-{right} {self.v}"

    def __repr__(self) -> str:
        return f"Edge({self.glyph()!r})"

directed_pair()

Return (parent, child) if this edge is directed, else None.

Source code in src/ergodic/graph/_graph.py

def directed_pair(self) -> tuple[str, str] | None:
    """Return ``(parent, child)`` if this edge is directed, else ``None``."""
    if not self.is_directed:
        return None
    if self.u_mark is Mark.TAIL:
        return (self.u, self.v)
    return (self.v, self.u)

glyph()

Render the edge, e.g. "X o-> Y".

Source code in src/ergodic/graph/_graph.py

def glyph(self) -> str:
    """Render the edge, e.g. ``"X o-> Y"``."""
    left = self.u_mark.char(left=True)
    right = self.v_mark.char(left=False)
    return f"{self.u} {left}-{right} {self.v}"

make(a, a_mark, b, b_mark) staticmethod

Build an edge between a and b, normalizing endpoint order.

Source code in src/ergodic/graph/_graph.py

@staticmethod
def make(a: str, a_mark: Mark, b: str, b_mark: Mark) -> Edge:
    """Build an edge between ``a`` and ``b``, normalizing endpoint order."""
    if a == b:
        raise ValueError(f"self-loops are not allowed: {a!r}")
    if a < b:
        return Edge(a, b, a_mark, b_mark)
    return Edge(b, a, b_mark, a_mark)

mark_at(node)

Return the mark this edge carries at node.

Source code in src/ergodic/graph/_graph.py

def mark_at(self, node: str) -> Mark:
    """Return the mark this edge carries at ``node``."""
    if node == self.u:
        return self.u_mark
    if node == self.v:
        return self.v_mark
    raise KeyError(node)

other(node)

Return the endpoint that is not node.

Source code in src/ergodic/graph/_graph.py

def other(self, node: str) -> str:
    """Return the endpoint that is not ``node``."""
    if node == self.u:
        return self.v
    if node == self.v:
        return self.u
    raise KeyError(node)

GraphKind

Bases: Enum

The kind of a mixed graph, which fixes the allowed marks and rules.

Source code in src/ergodic/graph/_graph.py

class GraphKind(Enum):
    """The kind of a mixed graph, which fixes the allowed marks and rules."""

    MIXED = "mixed"
    DAG = "dag"
    ADMG = "admg"
    MAG = "mag"
    CPDAG = "cpdag"
    PAG = "pag"

InvalidGraphError

Bases: ValueError

A graph violates the rules of its declared :class:GraphKind.

Source code in src/ergodic/graph/_graph.py

class InvalidGraphError(ValueError):
    """A graph violates the rules of its declared :class:`GraphKind`."""

Mark

Bases: Enum

A mark at one end of an edge.

A mixed graph edge has a mark at each end. TAIL (drawn -) and ARROW (drawn >) give directed (->), bidirected (<->), and undirected (--) edges. CIRCLE (drawn o) means the orientation is unknown, the marks an FCI search leaves on a PAG.

Source code in src/ergodic/graph/_marks.py

class Mark(Enum):
    """A mark at one end of an edge.

    A mixed graph edge has a mark at each end. ``TAIL`` (drawn ``-``) and
    ``ARROW`` (drawn ``>``) give directed (``->``), bidirected (``<->``), and
    undirected (``--``) edges. ``CIRCLE`` (drawn ``o``) means the orientation is
    unknown, the marks an FCI search leaves on a PAG.
    """

    TAIL = "tail"
    ARROW = "arrow"
    CIRCLE = "circle"

    def char(self, *, left: bool) -> str:
        """Return the glyph for this mark on the ``left`` or right side of an edge."""
        if self is Mark.TAIL:
            return "-"
        if self is Mark.CIRCLE:
            return "o"
        return "<" if left else ">"

char(*, left)

Return the glyph for this mark on the left or right side of an edge.

Source code in src/ergodic/graph/_marks.py

def char(self, *, left: bool) -> str:
    """Return the glyph for this mark on the ``left`` or right side of an edge."""
    if self is Mark.TAIL:
        return "-"
    if self is Mark.CIRCLE:
        return "o"
    return "<" if left else ">"

MixedGraph dataclass

An immutable mixed graph over named nodes.

Use the :func:dag, :func:admg, :func:cpdag, :func:pag, or :func:mixed builders rather than constructing this directly.

Source code in src/ergodic/graph/_graph.py

@dataclass(frozen=True)
class MixedGraph:
    """An immutable mixed graph over named nodes.

    Use the :func:`dag`, :func:`admg`, :func:`cpdag`, :func:`pag`, or
    :func:`mixed` builders rather than constructing this directly.
    """

    nodes: frozenset[str] = field(default_factory=frozenset)
    edges: frozenset[Edge] = field(default_factory=frozenset)
    latent: frozenset[str] = field(default_factory=frozenset)
    selection: frozenset[str] = field(default_factory=frozenset)
    kind: GraphKind = GraphKind.MIXED

    # -- lookups -------------------------------------------------------------

    def __contains__(self, node: object) -> bool:
        return node in self.nodes

    def __len__(self) -> int:
        return len(self.nodes)

    def node_list(self) -> list[str]:
        return sorted(self.nodes)

    def edge_list(self) -> list[Edge]:
        return sorted(self.edges, key=lambda e: (e.u, e.v))

    def has_edge(self, a: str, b: str) -> bool:
        return self.get_edge(a, b) is not None

    def get_edge(self, a: str, b: str) -> Edge | None:
        pair = frozenset((a, b))
        for edge in self.edges:
            if frozenset((edge.u, edge.v)) == pair:
                return edge
        return None

    def mark_at(self, a: str, b: str, node: str) -> Mark | None:
        edge = self.get_edge(a, b)
        return None if edge is None else edge.mark_at(node)

    # -- neighbourhoods ------------------------------------------------------

    def _incident(self, node: str) -> Iterator[Edge]:
        return (e for e in self.edges if node in (e.u, e.v))

    def neighbors(self, x: str) -> frozenset[str]:
        return frozenset(e.other(x) for e in self._incident(x))

    def parents(self, x: str) -> frozenset[str]:
        out: set[str] = set()
        for e in self._incident(x):
            pair = e.directed_pair()
            if pair is not None and pair[1] == x:
                out.add(pair[0])
        return frozenset(out)

    def children(self, x: str) -> frozenset[str]:
        out: set[str] = set()
        for e in self._incident(x):
            pair = e.directed_pair()
            if pair is not None and pair[0] == x:
                out.add(pair[1])
        return frozenset(out)

    def spouses(self, x: str) -> frozenset[str]:
        return frozenset(e.other(x) for e in self._incident(x) if e.is_bidirected)

    def ancestors(self, x: str) -> frozenset[str]:
        reverse: dict[str, set[str]] = {n: set() for n in self.nodes}
        for e in self.edges:
            pair = e.directed_pair()
            if pair is not None:
                reverse[pair[1]].add(pair[0])
        return self._reach(x, reverse)

    def descendants(self, x: str) -> frozenset[str]:
        forward: dict[str, set[str]] = {n: set() for n in self.nodes}
        for e in self.edges:
            pair = e.directed_pair()
            if pair is not None:
                forward[pair[0]].add(pair[1])
        return self._reach(x, forward)

    @staticmethod
    def _reach(start: str, adj: dict[str, set[str]]) -> frozenset[str]:
        seen: set[str] = set()
        stack = [start]
        while stack:
            node = stack.pop()
            for nxt in adj.get(node, ()):
                if nxt not in seen:
                    seen.add(nxt)
                    stack.append(nxt)
        seen.discard(start)
        return frozenset(seen)

    def markov_blanket(self, x: str) -> frozenset[str]:
        """Parents, children, and the other parents of children (DAG only)."""
        if not self.is_dag():
            raise NotImplementedError("markov_blanket is implemented for DAGs only")
        blanket: set[str] = set(self.parents(x)) | set(self.children(x))
        for child in self.children(x):
            blanket |= self.parents(child)
        blanket.discard(x)
        return frozenset(blanket)

    # -- immutable edits -----------------------------------------------------

    def add_nodes(self, *names: str) -> MixedGraph:
        return replace(self, nodes=self.nodes | frozenset(names))

    def add_edge(self, spec: EdgeSpec) -> MixedGraph:
        """Add an edge on a free node pair; adding the edge already there is a no-op.

        Raises ``ValueError`` when the pair already carries a different edge.
        Call :meth:`remove_edge` first, or :meth:`orient` to redirect.
        """
        edge = _parse_spec(spec)
        existing = self.get_edge(edge.u, edge.v)
        if existing == edge:
            return self
        if existing is not None:
            raise ValueError(
                f"the pair ({edge.u!r}, {edge.v!r}) already has edge {existing.glyph()!r}; "
                f"call remove_edge first to replace it with {edge.glyph()!r}"
            )
        return replace(
            self,
            nodes=self.nodes | {edge.u, edge.v},
            edges=self.edges | {edge},
            kind=GraphKind.MIXED,
        )

    def remove_edge(self, a: str, b: str) -> MixedGraph:
        edge = self.get_edge(a, b)
        if edge is None:
            return self
        return replace(self, edges=self.edges - {edge})

    def orient(self, parent: str, child: str) -> MixedGraph:
        """Return a graph with the edge between the two nodes set to ``parent -> child``."""
        return self.remove_edge(parent, child).add_edge((parent, Mark.TAIL, child, Mark.ARROW))

    def set_mark(self, a: str, b: str, *, at: str, mark: Mark) -> MixedGraph:
        """Return a graph with the ``{a, b}`` edge's mark at endpoint ``at`` set to ``mark``.

        The finer-grained sibling of :meth:`orient`: one endpoint changes, the
        other keeps its mark (how the FCI rules orient one mark at a time).
        Raises ``KeyError`` when the pair carries no edge or ``at`` is not one
        of its endpoints.
        """
        edge = self.get_edge(a, b)
        if edge is None:
            raise KeyError(f"no edge between {a!r} and {b!r}")
        if at not in (edge.u, edge.v):
            raise KeyError(f"{at!r} is not an endpoint of the edge between {a!r} and {b!r}")
        other = edge.other(at)
        new = Edge.make(at, mark, other, edge.mark_at(other))
        if new == edge:
            return self
        return replace(self, edges=(self.edges - {edge}) | {new})

    def as_kind(self, kind: GraphKind | str) -> MixedGraph:
        """Re-tag the graph as ``kind`` (an enum or its name) after checking it fits."""
        if isinstance(kind, str):
            try:
                kind = GraphKind(kind.lower())
            except ValueError:
                names = ", ".join(k.value for k in GraphKind)
                raise ValueError(f"unknown graph kind {kind!r}; expected one of: {names}") from None
        out = replace(self, kind=kind)
        out.validate()
        return out

    # -- kind checks ---------------------------------------------------------

    def _kahn(self) -> tuple[list[str], bool]:
        """Kahn's algorithm over the directed edges: a topological order and acyclicity.

        Iterative so deep graphs do not hit the recursion limit. Ties (several
        ready nodes) break by name, so the order is deterministic. The second
        element is ``True`` when every node was emitted, i.e. the directed part
        is acyclic; on a cycle the order is the prefix that could be emitted.
        """
        forward: dict[str, set[str]] = {n: set() for n in self.nodes}
        for e in self.edges:
            pair = e.directed_pair()
            if pair is not None:
                forward.setdefault(pair[0], set()).add(pair[1])
                forward.setdefault(pair[1], set())
        indegree: dict[str, int] = dict.fromkeys(forward, 0)
        for targets in forward.values():
            for nxt in targets:
                indegree[nxt] += 1
        ready = sorted(n for n, degree in indegree.items() if degree == 0)
        order: list[str] = []
        while ready:
            node = ready.pop(0)
            order.append(node)
            newly: list[str] = []
            for nxt in forward[node]:
                indegree[nxt] -= 1
                if indegree[nxt] == 0:
                    newly.append(nxt)
            if newly:
                ready = sorted(ready + newly)
        return order, len(order) == len(forward)

    def is_acyclic(self) -> bool:
        """Return ``True`` if the directed part of the graph has no cycle."""
        return self._kahn()[1]

    def is_dag(self) -> bool:
        """Every edge directed and no directed cycle. This check is complete."""
        return all(e.is_directed for e in self.edges) and self.is_acyclic()

    def is_admg(self) -> bool:
        """Directed or bidirected edges only, and no directed cycle. This check is complete."""
        marks_ok = all(e.is_directed or e.is_bidirected for e in self.edges)
        return marks_ok and self.is_acyclic()

    def is_cpdag(self) -> bool:
        """Return ``True`` when the graph really is a CPDAG. This check is complete.

        Three things must hold: every edge is directed or undirected, some DAG
        has exactly this skeleton and these v-structures (the Dor-Tarsi
        consistent extension exists), and the graph equals the CPDAG of that
        extension, so every compelled edge is directed and every reversible
        edge is left undirected.
        """
        if not all(e.is_directed or e.is_undirected for e in self.edges):
            return False
        if not self.is_acyclic():
            return False
        from ._meek import consistent_extension, cpdag_of_dag

        extension = consistent_extension(self)
        if extension is None:
            return False
        return cpdag_of_dag(extension).edges == self.edges

    def is_mag(self) -> bool:
        """The ancestral rules of Richardson and Spirtes (2002).

        Three edge types (directed, bidirected, undirected), no directed
        cycle, no almost-directed cycle (a bidirected edge between ancestors),
        and the endpoints of any undirected edge carry no arrowhead at all (no
        parents and no spouses), which is how selection bias shows up in a
        MAG. One gap: maximality (every non-adjacent pair separable) is not
        checked here; the :meth:`is_pag` round trip is what exercises it.
        """
        if not all(e.is_directed or e.is_bidirected or e.is_undirected for e in self.edges):
            return False
        if not self.is_acyclic():
            return False
        undirected = {n for e in self.edges if e.is_undirected for n in (e.u, e.v)}
        for e in self.edges:
            if e.is_bidirected and (e.u in self.ancestors(e.v) or e.v in self.ancestors(e.u)):
                return False
            if e.u_mark is Mark.ARROW and e.u in undirected:
                return False
            if e.v_mark is Mark.ARROW and e.v in undirected:
                return False
        return True

    def is_pag(self) -> bool:
        """Return ``True`` when the graph really is a complete PAG. This check is complete.

        The round trip decides it, exactly as :meth:`is_cpdag` does for
        CPDAGs: extract one member MAG (``mag_of_pag``), check it satisfies
        the MAG rules, and require the canonical PAG of that member
        (``pag_of_mag``) to be exactly this graph, so every invariant mark is
        oriented and every varying mark is a circle. Selection bias is in
        scope: undirected and circle-tail edges are read as Zhang's rules R5
        to R7 leave them. One known gap rides along from :meth:`is_mag`:
        maximality of the member is not checked.
        """
        if not self.is_acyclic():
            return False
        from ._pag import mag_of_pag, pag_of_mag

        try:
            member = mag_of_pag(self)
        except InvalidGraphError:
            return False
        return pag_of_mag(member).edges == self.edges

    def validate(self, kind: GraphKind | None = None) -> None:
        """Check the graph against ``kind`` (default: its own), raising on failure."""
        target = self.kind if kind is None else kind
        for e in self.edges:
            if e.u not in self.nodes or e.v not in self.nodes:
                raise InvalidGraphError(f"edge {e.glyph()!r} references an unknown node")
        checks = {
            GraphKind.MIXED: lambda: True,
            GraphKind.DAG: self.is_dag,
            GraphKind.ADMG: self.is_admg,
            GraphKind.MAG: self.is_mag,
            GraphKind.CPDAG: self.is_cpdag,
            GraphKind.PAG: self.is_pag,
        }
        if not checks[target]():
            raise InvalidGraphError(f"graph does not satisfy the rules of {target.value}")

    # -- algorithms (delegated, kept out of this class) ----------------------

    def m_separated(self, x: Any, y: Any, given: Any = frozenset()) -> bool:
        from ._separation import m_separated

        return m_separated(self, x, y, given)

    def d_separated(self, x: Any, y: Any, given: Any = frozenset()) -> bool:
        from ._separation import d_separated

        return d_separated(self, x, y, given)

    def do(self, targets: Any) -> MixedGraph:
        """Graph surgery for an intervention: cut every arrowhead into ``targets``."""
        from ._surgery import intervene

        return intervene(self, targets)

    def induced_subgraph(self, nodes: Any) -> MixedGraph:
        from ._surgery import induced_subgraph

        return induced_subgraph(self, nodes)

    def ancestral_subgraph(self, nodes: Any) -> MixedGraph:
        from ._surgery import ancestral_subgraph

        return ancestral_subgraph(self, nodes)

    # -- interop -------------------------------------------------------------

    def to_networkx(self) -> Any:
        from ._io import to_networkx

        return to_networkx(self)

    def to_adjacency(self, order: list[str] | None = None) -> tuple[np.ndarray, list[str]]:
        from ._io import to_adjacency

        return to_adjacency(self, order)

    def to_dict(self) -> dict[str, Any]:
        from ._io import to_dict

        return to_dict(self)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> MixedGraph:
        from ._io import from_dict

        return from_dict(data)

    @classmethod
    def from_networkx(cls, graph: Any) -> MixedGraph:
        from ._io import from_networkx

        return from_networkx(graph)

    def __repr__(self) -> str:
        glyphs = ", ".join(e.glyph() for e in self.edge_list())
        return f"MixedGraph(kind={self.kind.value}, nodes={len(self.nodes)}, edges=[{glyphs}])"

add_edge(spec)

Add an edge on a free node pair; adding the edge already there is a no-op.

Raises ValueError when the pair already carries a different edge. Call :meth:remove_edge first, or :meth:orient to redirect.

Source code in src/ergodic/graph/_graph.py

def add_edge(self, spec: EdgeSpec) -> MixedGraph:
    """Add an edge on a free node pair; adding the edge already there is a no-op.

    Raises ``ValueError`` when the pair already carries a different edge.
    Call :meth:`remove_edge` first, or :meth:`orient` to redirect.
    """
    edge = _parse_spec(spec)
    existing = self.get_edge(edge.u, edge.v)
    if existing == edge:
        return self
    if existing is not None:
        raise ValueError(
            f"the pair ({edge.u!r}, {edge.v!r}) already has edge {existing.glyph()!r}; "
            f"call remove_edge first to replace it with {edge.glyph()!r}"
        )
    return replace(
        self,
        nodes=self.nodes | {edge.u, edge.v},
        edges=self.edges | {edge},
        kind=GraphKind.MIXED,
    )

as_kind(kind)

Re-tag the graph as kind (an enum or its name) after checking it fits.

Source code in src/ergodic/graph/_graph.py

def as_kind(self, kind: GraphKind | str) -> MixedGraph:
    """Re-tag the graph as ``kind`` (an enum or its name) after checking it fits."""
    if isinstance(kind, str):
        try:
            kind = GraphKind(kind.lower())
        except ValueError:
            names = ", ".join(k.value for k in GraphKind)
            raise ValueError(f"unknown graph kind {kind!r}; expected one of: {names}") from None
    out = replace(self, kind=kind)
    out.validate()
    return out

do(targets)

Graph surgery for an intervention: cut every arrowhead into targets.

Source code in src/ergodic/graph/_graph.py

def do(self, targets: Any) -> MixedGraph:
    """Graph surgery for an intervention: cut every arrowhead into ``targets``."""
    from ._surgery import intervene

    return intervene(self, targets)

is_acyclic()

Return True if the directed part of the graph has no cycle.

Source code in src/ergodic/graph/_graph.py

def is_acyclic(self) -> bool:
    """Return ``True`` if the directed part of the graph has no cycle."""
    return self._kahn()[1]

is_admg()

Directed or bidirected edges only, and no directed cycle. This check is complete.

Source code in src/ergodic/graph/_graph.py

def is_admg(self) -> bool:
    """Directed or bidirected edges only, and no directed cycle. This check is complete."""
    marks_ok = all(e.is_directed or e.is_bidirected for e in self.edges)
    return marks_ok and self.is_acyclic()

is_cpdag()

Return True when the graph really is a CPDAG. This check is complete.

Three things must hold: every edge is directed or undirected, some DAG has exactly this skeleton and these v-structures (the Dor-Tarsi consistent extension exists), and the graph equals the CPDAG of that extension, so every compelled edge is directed and every reversible edge is left undirected.

Source code in src/ergodic/graph/_graph.py

def is_cpdag(self) -> bool:
    """Return ``True`` when the graph really is a CPDAG. This check is complete.

    Three things must hold: every edge is directed or undirected, some DAG
    has exactly this skeleton and these v-structures (the Dor-Tarsi
    consistent extension exists), and the graph equals the CPDAG of that
    extension, so every compelled edge is directed and every reversible
    edge is left undirected.
    """
    if not all(e.is_directed or e.is_undirected for e in self.edges):
        return False
    if not self.is_acyclic():
        return False
    from ._meek import consistent_extension, cpdag_of_dag

    extension = consistent_extension(self)
    if extension is None:
        return False
    return cpdag_of_dag(extension).edges == self.edges

is_dag()

Every edge directed and no directed cycle. This check is complete.

Source code in src/ergodic/graph/_graph.py

def is_dag(self) -> bool:
    """Every edge directed and no directed cycle. This check is complete."""
    return all(e.is_directed for e in self.edges) and self.is_acyclic()

is_mag()

The ancestral rules of Richardson and Spirtes (2002).

Three edge types (directed, bidirected, undirected), no directed cycle, no almost-directed cycle (a bidirected edge between ancestors), and the endpoints of any undirected edge carry no arrowhead at all (no parents and no spouses), which is how selection bias shows up in a MAG. One gap: maximality (every non-adjacent pair separable) is not checked here; the :meth:is_pag round trip is what exercises it.

Source code in src/ergodic/graph/_graph.py

def is_mag(self) -> bool:
    """The ancestral rules of Richardson and Spirtes (2002).

    Three edge types (directed, bidirected, undirected), no directed
    cycle, no almost-directed cycle (a bidirected edge between ancestors),
    and the endpoints of any undirected edge carry no arrowhead at all (no
    parents and no spouses), which is how selection bias shows up in a
    MAG. One gap: maximality (every non-adjacent pair separable) is not
    checked here; the :meth:`is_pag` round trip is what exercises it.
    """
    if not all(e.is_directed or e.is_bidirected or e.is_undirected for e in self.edges):
        return False
    if not self.is_acyclic():
        return False
    undirected = {n for e in self.edges if e.is_undirected for n in (e.u, e.v)}
    for e in self.edges:
        if e.is_bidirected and (e.u in self.ancestors(e.v) or e.v in self.ancestors(e.u)):
            return False
        if e.u_mark is Mark.ARROW and e.u in undirected:
            return False
        if e.v_mark is Mark.ARROW and e.v in undirected:
            return False
    return True

is_pag()

Return True when the graph really is a complete PAG. This check is complete.

The round trip decides it, exactly as :meth:is_cpdag does for CPDAGs: extract one member MAG (mag_of_pag), check it satisfies the MAG rules, and require the canonical PAG of that member (pag_of_mag) to be exactly this graph, so every invariant mark is oriented and every varying mark is a circle. Selection bias is in scope: undirected and circle-tail edges are read as Zhang's rules R5 to R7 leave them. One known gap rides along from :meth:is_mag: maximality of the member is not checked.

Source code in src/ergodic/graph/_graph.py

def is_pag(self) -> bool:
    """Return ``True`` when the graph really is a complete PAG. This check is complete.

    The round trip decides it, exactly as :meth:`is_cpdag` does for
    CPDAGs: extract one member MAG (``mag_of_pag``), check it satisfies
    the MAG rules, and require the canonical PAG of that member
    (``pag_of_mag``) to be exactly this graph, so every invariant mark is
    oriented and every varying mark is a circle. Selection bias is in
    scope: undirected and circle-tail edges are read as Zhang's rules R5
    to R7 leave them. One known gap rides along from :meth:`is_mag`:
    maximality of the member is not checked.
    """
    if not self.is_acyclic():
        return False
    from ._pag import mag_of_pag, pag_of_mag

    try:
        member = mag_of_pag(self)
    except InvalidGraphError:
        return False
    return pag_of_mag(member).edges == self.edges

markov_blanket(x)

Parents, children, and the other parents of children (DAG only).

Source code in src/ergodic/graph/_graph.py

def markov_blanket(self, x: str) -> frozenset[str]:
    """Parents, children, and the other parents of children (DAG only)."""
    if not self.is_dag():
        raise NotImplementedError("markov_blanket is implemented for DAGs only")
    blanket: set[str] = set(self.parents(x)) | set(self.children(x))
    for child in self.children(x):
        blanket |= self.parents(child)
    blanket.discard(x)
    return frozenset(blanket)

orient(parent, child)

Return a graph with the edge between the two nodes set to parent -> child.

Source code in src/ergodic/graph/_graph.py

def orient(self, parent: str, child: str) -> MixedGraph:
    """Return a graph with the edge between the two nodes set to ``parent -> child``."""
    return self.remove_edge(parent, child).add_edge((parent, Mark.TAIL, child, Mark.ARROW))

set_mark(a, b, *, at, mark)

Return a graph with the {a, b} edge's mark at endpoint at set to mark.

The finer-grained sibling of :meth:orient: one endpoint changes, the other keeps its mark (how the FCI rules orient one mark at a time). Raises KeyError when the pair carries no edge or at is not one of its endpoints.

Source code in src/ergodic/graph/_graph.py

def set_mark(self, a: str, b: str, *, at: str, mark: Mark) -> MixedGraph:
    """Return a graph with the ``{a, b}`` edge's mark at endpoint ``at`` set to ``mark``.

    The finer-grained sibling of :meth:`orient`: one endpoint changes, the
    other keeps its mark (how the FCI rules orient one mark at a time).
    Raises ``KeyError`` when the pair carries no edge or ``at`` is not one
    of its endpoints.
    """
    edge = self.get_edge(a, b)
    if edge is None:
        raise KeyError(f"no edge between {a!r} and {b!r}")
    if at not in (edge.u, edge.v):
        raise KeyError(f"{at!r} is not an endpoint of the edge between {a!r} and {b!r}")
    other = edge.other(at)
    new = Edge.make(at, mark, other, edge.mark_at(other))
    if new == edge:
        return self
    return replace(self, edges=(self.edges - {edge}) | {new})

validate(kind=None)

Check the graph against kind (default: its own), raising on failure.

Source code in src/ergodic/graph/_graph.py

def validate(self, kind: GraphKind | None = None) -> None:
    """Check the graph against ``kind`` (default: its own), raising on failure."""
    target = self.kind if kind is None else kind
    for e in self.edges:
        if e.u not in self.nodes or e.v not in self.nodes:
            raise InvalidGraphError(f"edge {e.glyph()!r} references an unknown node")
    checks = {
        GraphKind.MIXED: lambda: True,
        GraphKind.DAG: self.is_dag,
        GraphKind.ADMG: self.is_admg,
        GraphKind.MAG: self.is_mag,
        GraphKind.CPDAG: self.is_cpdag,
        GraphKind.PAG: self.is_pag,
    }
    if not checks[target]():
        raise InvalidGraphError(f"graph does not satisfy the rules of {target.value}")

admg(edges=(), **kwargs)

Build an acyclic directed mixed graph (directed and bidirected edges).

Source code in src/ergodic/graph/_graph.py

def admg(edges: Iterable[EdgeSpec] = (), **kwargs: Any) -> MixedGraph:
    """Build an acyclic directed mixed graph (directed and bidirected edges)."""
    return mixed(edges, kind=GraphKind.ADMG, **kwargs)

consistent_extension(pdag)

A DAG with exactly the PDAG's skeleton and v-structures, or None.

Dor and Tarsi (1992): repeatedly find a node x with no outgoing directed edges whose undirected neighbours are adjacent to every other neighbour of x, orient all undirected edges at x into x, and retire x. When every node retires the oriented graph is the extension; when no candidate exists none does.

Source code in src/ergodic/graph/_meek.py

def consistent_extension(pdag: MixedGraph) -> MixedGraph | None:
    """A DAG with exactly the PDAG's skeleton and v-structures, or ``None``.

    Dor and Tarsi (1992): repeatedly find a node ``x`` with no outgoing
    directed edges whose undirected neighbours are adjacent to every other
    neighbour of ``x``, orient all undirected edges at ``x`` into ``x``, and
    retire ``x``. When every node retires the oriented graph is the
    extension; when no candidate exists none does.
    """
    maps = _pdag_maps(pdag, "consistent_extension")
    parents, children, undirected = maps
    remaining = set(pdag.nodes)
    graph = pdag
    while remaining:
        target = None
        for x in sorted(remaining):
            if children[x] & remaining:
                continue
            local_und = undirected[x] & remaining
            local_all = (parents[x] | undirected[x]) & remaining
            if all(_adjacent(u, w, maps) for u in local_und for w in local_all if w != u):
                target = x
                break
        if target is None:
            return None
        for u in sorted(undirected[target] & remaining):
            graph = graph.orient(u, target)
        remaining.discard(target)
    return graph.as_kind(GraphKind.DAG)

cpdag(edges=(), **kwargs)

Build a CPDAG (directed and undirected edges).

Source code in src/ergodic/graph/_graph.py

def cpdag(edges: Iterable[EdgeSpec] = (), **kwargs: Any) -> MixedGraph:
    """Build a CPDAG (directed and undirected edges)."""
    return mixed(edges, kind=GraphKind.CPDAG, **kwargs)

cpdag_of_dag(dag)

The CPDAG of a DAG: undirected skeleton, restored v-structures, Meek closure.

Every DAG in the Markov equivalence class of dag maps to the same CPDAG, which keeps an edge directed exactly when every member orients it the same way.

Source code in src/ergodic/graph/_meek.py

def cpdag_of_dag(dag: MixedGraph) -> MixedGraph:
    """The CPDAG of a DAG: undirected skeleton, restored v-structures, Meek closure.

    Every DAG in the Markov equivalence class of ``dag`` maps to the same
    CPDAG, which keeps an edge directed exactly when every member orients it
    the same way.
    """
    if not dag.is_dag():
        raise ValueError("cpdag_of_dag needs a DAG: every edge directed and no cycle")
    skeleton = frozenset(Edge.make(e.u, Mark.TAIL, e.v, Mark.TAIL) for e in dag.edges)
    out = MixedGraph(nodes=dag.nodes, edges=skeleton, kind=GraphKind.MIXED)
    for z in sorted(dag.nodes):
        for x, y in combinations(sorted(dag.parents(z)), 2):
            if not dag.has_edge(x, y):
                out = out.orient(x, z).orient(y, z)
    return replace(meek_closure(out), kind=GraphKind.CPDAG)

d_separated(graph, x, y, given=frozenset())

d-separation, the DAG special case of :func:m_separated.

Source code in src/ergodic/graph/_separation.py

def d_separated(
    graph: MixedGraph,
    x: str | Iterable[str],
    y: str | Iterable[str],
    given: str | Iterable[str] = frozenset(),
) -> bool:
    """d-separation, the DAG special case of :func:`m_separated`."""
    if not graph.is_dag():
        raise ValueError("d_separated requires a DAG; use m_separated for mixed graphs")
    return m_separated(graph, x, y, given)

dag(edges=(), **kwargs)

Build a directed acyclic graph. Edges default to directed (u -> v).

Source code in src/ergodic/graph/_graph.py

def dag(edges: Iterable[EdgeSpec] = (), **kwargs: Any) -> MixedGraph:
    """Build a directed acyclic graph. Edges default to directed (``u -> v``)."""
    return mixed(edges, kind=GraphKind.DAG, **kwargs)

dags_of_cpdag(cpdag, *, max_members=None)

Every DAG in the CPDAG's Markov equivalence class, in a deterministic order.

Branch on the first undirected edge, orient it each way, close under Meek's rules, and recurse. Distinct branches fix that edge differently, so no member appears twice and none is missed (the rules are sound, so every member agrees with the closure of its own branch). Each leaf is verified against the class definition, cpdag_of_dag(leaf) equal to the input, so a branch that closes into a non-member dies quietly and the answer is exact. max_members raises once the count would pass it; there is no silent truncation.

Source code in src/ergodic/graph/_meek.py

def dags_of_cpdag(cpdag: MixedGraph, *, max_members: int | None = None) -> tuple[MixedGraph, ...]:
    """Every DAG in the CPDAG's Markov equivalence class, in a deterministic order.

    Branch on the first undirected edge, orient it each way, close under
    Meek's rules, and recurse. Distinct branches fix that edge differently,
    so no member appears twice and none is missed (the rules are sound, so
    every member agrees with the closure of its own branch). Each leaf is
    verified against the class definition, ``cpdag_of_dag(leaf)`` equal to
    the input, so a branch that closes into a non-member dies quietly and
    the answer is exact. ``max_members`` raises once the count would pass
    it; there is no silent truncation.
    """
    if not cpdag.is_cpdag():
        raise ValueError("dags_of_cpdag needs a CPDAG; build one with cpdag() or cpdag_of_dag()")
    found: list[MixedGraph] = []

    def expand(graph: MixedGraph) -> None:
        undirected = sorted((e.u, e.v) for e in graph.edges if e.is_undirected)
        if not undirected:
            candidate = replace(graph, kind=GraphKind.DAG)
            if not candidate.is_dag() or cpdag_of_dag(candidate).edges != cpdag.edges:
                return
            if max_members is not None and len(found) >= max_members:
                raise ValueError(
                    f"the class has more than {max_members} members; raise max_members "
                    "or work on the CPDAG itself"
                )
            found.append(candidate)
            return
        u, v = undirected[0]
        for a, b in ((u, v), (v, u)):
            expand(meek_closure(graph.orient(a, b)))

    expand(replace(cpdag, kind=GraphKind.MIXED))
    return tuple(found)

district_of(graph, node)

Return the district (c-component) containing node.

Source code in src/ergodic/graph/_primitives.py

def district_of(graph: MixedGraph, node: str) -> frozenset[str]:
    """Return the district (c-component) containing ``node``."""
    if node not in graph.nodes:
        raise ValueError(f"unknown node name: {node!r}")
    component: set[str] = set()
    stack = [node]
    while stack:
        cur = stack.pop()
        if cur in component:
            continue
        component.add(cur)
        for nbr in graph.spouses(cur):
            if nbr not in component:
                stack.append(nbr)
    return frozenset(component)

districts(graph)

The c-components: connected components of the bidirected (<->) subgraph.

Two observed nodes share a district when a <-> path joins them; a node with no bidirected edge is its own singleton. Every node lands in exactly one district and the districts partition graph.nodes. On a DAG (no bidirected edges) every district is a singleton. Defined over the nodes as they are: project latents first (:func:latent_projection) if the graph still carries latent= designations.

Source code in src/ergodic/graph/_primitives.py

def districts(graph: MixedGraph) -> frozenset[frozenset[str]]:
    """The c-components: connected components of the bidirected (``<->``) subgraph.

    Two observed nodes share a district when a ``<->`` path joins them; a node
    with no bidirected edge is its own singleton. Every node lands in exactly
    one district and the districts partition ``graph.nodes``. On a DAG (no
    bidirected edges) every district is a singleton. Defined over the nodes as
    they are: project latents first (:func:`latent_projection`) if the graph
    still carries ``latent=`` designations.
    """
    seen: set[str] = set()
    out: set[frozenset[str]] = set()
    for start in sorted(graph.nodes):
        if start in seen:
            continue
        component: set[str] = set()
        stack = [start]
        while stack:
            node = stack.pop()
            if node in component:
                continue
            component.add(node)
            for nbr in graph.spouses(node):
                if nbr not in component:
                    stack.append(nbr)
        seen |= component
        out.add(frozenset(component))
    return frozenset(out)

extendable(pdag)

Return True when the PDAG admits a consistent extension.

Source code in src/ergodic/graph/_meek.py

def extendable(pdag: MixedGraph) -> bool:
    """Return ``True`` when the PDAG admits a consistent extension."""
    return consistent_extension(pdag) is not None

fci_closure(graph, *, discriminating)

Close a partially oriented graph under Zhang's rules R1 to R10.

The input carries its v-structure arrowheads (and any knowledge-seeded marks); circles elsewhere. Rules only ever turn a circle into a tail or an arrowhead: a mark that is already definite never moves, so a seeded mark that disagrees with a rule simply stays and the conflict surfaces at the caller's validation. discriminating(d, a, b, c) answers R4's question, whether b is a collider on the discriminating path; None skips that path. Sweeps run in sorted order until a full pass changes nothing, so the closure is deterministic. R5 to R7 are the selection rules and never fire on a selection-free class; Zhang stages the rules (R1 to R4, then R5 once, then R6 to R10) but every rule is individually sound and marks only move circle to definite, so the joint fixpoint is the same closure, which the round-trip tests hold executably. Returns kind mixed.

Source code in src/ergodic/graph/_pag.py

def fci_closure(
    graph: MixedGraph,
    *,
    discriminating: Callable[[str, str, str, str], bool | None],
) -> MixedGraph:
    """Close a partially oriented graph under Zhang's rules R1 to R10.

    The input carries its v-structure arrowheads (and any knowledge-seeded
    marks); circles elsewhere. Rules only ever turn a circle into a tail or
    an arrowhead: a mark that is already definite never moves, so a seeded
    mark that disagrees with a rule simply stays and the conflict surfaces at
    the caller's validation. ``discriminating(d, a, b, c)`` answers R4's
    question, whether b is a collider on the discriminating path; ``None``
    skips that path. Sweeps run in sorted order until a full pass changes
    nothing, so the closure is deterministic. R5 to R7 are the selection
    rules and never fire on a selection-free class; Zhang stages the rules
    (R1 to R4, then R5 once, then R6 to R10) but every rule is individually
    sound and marks only move circle to definite, so the joint fixpoint is
    the same closure, which the round-trip tests hold executably. Returns
    kind ``mixed``.
    """
    marks = _mark_maps(graph)
    adj = {n: sorted(marks[n]) for n in marks}
    changed = True
    while changed:
        changed = False
        changed |= _r1(marks, adj)
        changed |= _r2(marks, adj)
        changed |= _r3(marks, adj)
        changed |= _r4(marks, adj, discriminating)
        changed |= _r5(marks, adj)
        changed |= _r6(marks, adj)
        changed |= _r7(marks, adj)
        changed |= _r8(marks, adj)
        changed |= _r9(marks, adj)
        changed |= _r10(marks, adj)
    return _rebuild(graph, marks)

latent_projection(graph, keep=None)

Project a DAG or ADMG with latents onto the observed margin, as an ADMG.

keep defaults to graph.nodes - graph.latent. The projected-out set is H = nodes - keep. Over keep the result is an ADMG (kind=ADMG, empty latent):

• a -> b iff a directed path a -> ... -> b has every intermediate node in H.
• a <-> b iff a path between a and b has every intermediate node in H and a non-collider, with arrowheads into both a and b.

Computed by iterative single-node elimination of each node in H, which is equivalent. The defining contract is m-separation preservation: for disjoint observed A, B, Z (none in H), the projection and graph agree on m_separated(A, B, Z). This is the ADMG projection the ID algorithm needs, distinct from the MAG ancestral projection :func:mag_of_dag.

The result can carry a bow as two parallel edges on one pair (a -> b and a <-> b), which the marks model otherwise forbids. This is correct for an ADMG, and the ID engine reads it through parents, spouses, districts, and m_separated, all of which see both edges. A projected graph is an internal ID input, not a serializable one: to_dict, to_adjacency, and to_networkx raise on a parallel-edge pair rather than drop or mangle it.

Source code in src/ergodic/graph/_primitives.py

def latent_projection(graph: MixedGraph, keep: Iterable[str] | None = None) -> MixedGraph:
    """Project a DAG or ADMG with latents onto the observed margin, as an ADMG.

    ``keep`` defaults to ``graph.nodes - graph.latent``. The projected-out set
    is ``H = nodes - keep``. Over ``keep`` the result is an ADMG (``kind=ADMG``,
    empty ``latent``):

    - ``a -> b`` iff a directed path ``a -> ... -> b`` has every intermediate
      node in ``H``.
    - ``a <-> b`` iff a path between ``a`` and ``b`` has every intermediate node
      in ``H`` and a non-collider, with arrowheads into both ``a`` and ``b``.

    Computed by iterative single-node elimination of each node in ``H``, which
    is equivalent. The defining contract is m-separation preservation: for
    disjoint observed ``A``, ``B``, ``Z`` (none in ``H``), the projection and
    ``graph`` agree on ``m_separated(A, B, Z)``. This is the ADMG projection
    the ID algorithm needs, distinct from the MAG ancestral projection
    :func:`mag_of_dag`.

    The result can carry a bow as two parallel edges on one pair (``a -> b`` and
    ``a <-> b``), which the marks model otherwise forbids. This is correct for an
    ADMG, and the ID engine reads it through ``parents``, ``spouses``,
    ``districts``, and ``m_separated``, all of which see both edges. A projected
    graph is an internal ID input, not a serializable one: ``to_dict``,
    ``to_adjacency``, and ``to_networkx`` raise on a parallel-edge pair rather
    than drop or mangle it.
    """
    if not graph.is_admg():
        raise ValueError(
            "latent_projection needs a DAG or ADMG (directed and bidirected "
            "edges, no directed cycle); project the graph you have onto one first"
        )
    keep_set = (graph.nodes - graph.latent) if keep is None else frozenset(keep)
    unknown = keep_set - graph.nodes
    if unknown:
        raise ValueError(f"unknown node names to keep: {', '.join(sorted(unknown))}")
    eliminate = graph.nodes - keep_set

    # Track directed and bidirected relations separately, since the projection
    # of a DAG/ADMG can produce a bow (both ``a -> b`` and ``a <-> b`` on one
    # pair) and m-separation needs both: the directed edge makes ``a`` a
    # non-collider for pass-through paths, the bidirected edge makes it a
    # collider that conditioning opens. They are kept as two edges; the
    # one-edge-per-pair invariant only constrains the glyph builders.
    directed: set[tuple[str, str]] = set()
    bidir: set[frozenset[str]] = set()
    for e in graph.edges:
        pair = e.directed_pair()
        if pair is not None:
            directed.add(pair)
        elif e.is_bidirected:
            bidir.add(frozenset((e.u, e.v)))

    def parents_of(node: str) -> set[str]:
        return {p for (p, c) in directed if c == node}

    def children_of(node: str) -> set[str]:
        return {c for (p, c) in directed if p == node}

    def spouses_of(node: str) -> set[str]:
        return {next(iter(pr - {node})) for pr in bidir if node in pr}

    for h in sorted(eliminate):
        pa = parents_of(h)
        ch = children_of(h)
        sp = spouses_of(h)
        # A path through h is a non-collider iff one side is a child of h
        # (the ``h -> c`` edge gives h a tail). Reconnect each child to every
        # other neighbour, arrowhead into the child, the neighbour keeping its
        # mark: a parent stays a tail (directed), a child or spouse an arrow
        # (bidirected).
        for c in ch:
            for p in pa:
                if p != c:
                    directed.add((p, c))
            for other in (ch | sp) - {c}:
                bidir.add(frozenset((c, other)))
        directed = {(p, c) for (p, c) in directed if p != h and c != h}
        bidir = {pr for pr in bidir if h not in pr}

    edges: list[Edge] = []
    for p, c in sorted(directed):
        if p in keep_set and c in keep_set:
            edges.append(Edge.make(p, Mark.TAIL, c, Mark.ARROW))
    for pr in sorted(sorted(pr) for pr in bidir):
        a, b = pr
        if a in keep_set and b in keep_set:
            edges.append(Edge.make(a, Mark.ARROW, b, Mark.ARROW))
    out = MixedGraph(nodes=keep_set, edges=frozenset(edges), kind=GraphKind.ADMG)
    out.validate()
    return out

m_connected(graph, x, y, given=frozenset())

Return True if some active path joins x and y given given.

Circle-free graphs (DAG, ADMG, MAG) use plain m-connection. A graph with circle marks (a PAG) uses definite m-connection, which on a valid complete PAG is the class-invariant answer. A CPDAG-kind graph gets the same class reading with its undirected edges recoded as circles; an undirected edge on any other kind is walked literally, as a tail at both ends.

Source code in src/ergodic/graph/_separation.py

def m_connected(
    graph: MixedGraph,
    x: str | Iterable[str],
    y: str | Iterable[str],
    given: str | Iterable[str] = frozenset(),
) -> bool:
    """Return ``True`` if some active path joins ``x`` and ``y`` given ``given``.

    Circle-free graphs (DAG, ADMG, MAG) use plain m-connection. A graph with
    circle marks (a PAG) uses definite m-connection, which on a valid complete
    PAG is the class-invariant answer. A CPDAG-kind graph gets the same class
    reading with its undirected edges recoded as circles; an undirected edge
    on any other kind is walked literally, as a tail at both ends.
    """
    sources = _as_set(x)
    targets = _as_set(y)
    conditioned = _as_set(given)

    unknown = sorted((sources | targets | conditioned) - graph.nodes)
    if unknown:
        raise ValueError(f"unknown node names (not in the graph): {unknown}")
    overlap = (sources & targets) | (sources & conditioned) | (targets & conditioned)
    if overlap:
        raise ValueError(f"x, y, and given must be disjoint; shared nodes: {sorted(overlap)}")

    from ._graph import Edge, GraphKind

    if graph.kind is GraphKind.CPDAG:
        recoded = replace(
            graph,
            edges=frozenset(
                Edge.make(e.u, Mark.CIRCLE, e.v, Mark.CIRCLE) if e.is_undirected else e
                for e in graph.edges
            ),
            kind=GraphKind.MIXED,
        )
        return _definitely_connected(recoded, sources, targets, conditioned)
    if any(e.has_circle for e in graph.edges):
        return _definitely_connected(graph, sources, targets, conditioned)
    return _plainly_connected(graph, sources, targets, conditioned)

m_separated(graph, x, y, given=frozenset())

Return True if x and y are m-separated given given.

Source code in src/ergodic/graph/_separation.py

def m_separated(
    graph: MixedGraph,
    x: str | Iterable[str],
    y: str | Iterable[str],
    given: str | Iterable[str] = frozenset(),
) -> bool:
    """Return ``True`` if ``x`` and ``y`` are m-separated given ``given``."""
    return not m_connected(graph, x, y, given)

mag_of_dag(graph)

Project a DAG or ADMG with latent and selection nodes onto its MAG.

Two observed nodes are adjacent exactly when no observed subset separates them given the selection nodes, decided by the one witness set that suffices (Richardson and Spirtes 2002, Theorem 4.2): the observed ancestors of either endpoint or of selection, plus the selection nodes themselves. Marks read ancestry in the full graph: a tail at X on the edge to Y means X is an ancestor of Y or of selection there, an arrowhead means it is neither; two tails make an undirected edge, which on a DAG means both endpoints are ancestors of selection. The result preserves every m-separation over the observed margin conditional on selection, which is the defining property of the projection.

Source code in src/ergodic/graph/_pag.py

def mag_of_dag(graph: MixedGraph) -> MixedGraph:
    """Project a DAG or ADMG with latent and selection nodes onto its MAG.

    Two observed nodes are adjacent exactly when no observed subset separates
    them given the selection nodes, decided by the one witness set that
    suffices (Richardson and Spirtes 2002, Theorem 4.2): the observed
    ancestors of either endpoint or of selection, plus the selection nodes
    themselves. Marks read ancestry in the full graph: a tail at ``X`` on the
    edge to ``Y`` means ``X`` is an ancestor of ``Y`` or of selection there,
    an arrowhead means it is neither; two tails make an undirected edge,
    which on a DAG means both endpoints are ancestors of selection. The
    result preserves every m-separation over the observed margin conditional
    on selection, which is the defining property of the projection.
    """
    if not graph.is_admg():
        raise ValueError(
            "mag_of_dag needs a DAG or ADMG (directed and bidirected edges, no "
            "directed cycle); project the graph you have onto one first"
        )
    stray = graph.selection - graph.nodes
    if stray:
        raise ValueError(f"selection names not in the graph: {', '.join(sorted(stray))}")
    overlap = graph.selection & graph.latent
    if overlap:
        raise ValueError(f"nodes cannot be both latent and selection: {', '.join(sorted(overlap))}")
    selection = set(graph.selection)
    observed = sorted(graph.nodes - graph.latent - graph.selection)
    ancestors = {v: graph.ancestors(v) for v in graph.nodes}
    sel_closure = selection | {a for s in selection for a in ancestors[s]}
    edges: list[Edge] = []
    for x, y in combinations(observed, 2):
        pool = (ancestors[x] | ancestors[y] | sel_closure) & (set(observed) | selection)
        if m_separated(graph, x, y, pool - {x, y}):
            continue
        x_mark = Mark.TAIL if (x in ancestors[y] or x in sel_closure) else Mark.ARROW
        y_mark = Mark.TAIL if (y in ancestors[x] or y in sel_closure) else Mark.ARROW
        edges.append(Edge.make(x, x_mark, y, y_mark))
    out = MixedGraph(nodes=frozenset(observed), edges=frozenset(edges), kind=GraphKind.MAG)
    out.validate()
    return out

mag_of_pag(p)

Extract one member MAG from a PAG (Zhang 2008, arrowhead augmentation).

Every o-> edge becomes directed into its arrowhead, every -o edge becomes directed out of its tail (the circle takes the arrowhead), undirected and bidirected and directed edges stay as they are, and the circle-circle component is oriented into a DAG that adds no unshielded collider (the Dor-Tarsi extension of its undirected skeleton; for a valid complete PAG that component is chordal, R5 guarantees it, so the extension exists). An inextricable circle component or a result that fails the MAG rules raises. Maximality of the member is not checked here because is_mag does not check it; the is_pag round trip is what guards it.

Source code in src/ergodic/graph/_pag.py

def mag_of_pag(p: MixedGraph) -> MixedGraph:
    """Extract one member MAG from a PAG (Zhang 2008, arrowhead augmentation).

    Every ``o->`` edge becomes directed into its arrowhead, every ``-o``
    edge becomes directed out of its tail (the circle takes the arrowhead),
    undirected and bidirected and directed edges stay as they are, and the
    circle-circle component is oriented into a DAG that adds no unshielded
    collider (the Dor-Tarsi extension of its undirected skeleton; for a
    valid complete PAG that component is chordal, R5 guarantees it, so the
    extension exists). An inextricable circle component or a result that
    fails the MAG rules raises. Maximality of the member is not checked here
    because ``is_mag`` does not check it; the ``is_pag`` round trip is what
    guards it.
    """
    fixed: list[Edge] = []
    circle_pairs: list[tuple[str, str]] = []
    for e in sorted(p.edges, key=lambda e: (e.u, e.v)):
        cu, cv = e.u_mark, e.v_mark
        if cu is Mark.CIRCLE and cv is Mark.CIRCLE:
            circle_pairs.append((e.u, e.v))
        elif cu is Mark.CIRCLE and cv is Mark.ARROW:
            fixed.append(Edge.make(e.u, Mark.TAIL, e.v, Mark.ARROW))
        elif cu is Mark.ARROW and cv is Mark.CIRCLE:
            fixed.append(Edge.make(e.u, Mark.ARROW, e.v, Mark.TAIL))
        elif cu is Mark.TAIL and cv is Mark.CIRCLE:
            fixed.append(Edge.make(e.u, Mark.TAIL, e.v, Mark.ARROW))
        elif cu is Mark.CIRCLE and cv is Mark.TAIL:
            fixed.append(Edge.make(e.u, Mark.ARROW, e.v, Mark.TAIL))
        else:
            fixed.append(e)
    if circle_pairs:
        component_nodes = frozenset(n for pair in circle_pairs for n in pair)
        undirected = MixedGraph(
            nodes=component_nodes,
            edges=frozenset(Edge.make(u, Mark.TAIL, v, Mark.TAIL) for u, v in circle_pairs),
            kind=GraphKind.MIXED,
        )
        extension = consistent_extension(undirected)
        if extension is None:
            raise InvalidGraphError(
                "the circle component of this graph is not extendable without new "
                "colliders, so it is not a valid PAG"
            )
        fixed.extend(extension.edges)
    mag = MixedGraph(nodes=p.nodes, edges=frozenset(fixed), kind=GraphKind.MAG)
    if not mag.is_mag():
        raise InvalidGraphError("the extracted member violates the MAG rules")
    return mag

mags_of_pag(p, *, max_members=None, selection=False)

Every MAG in the PAG's equivalence class, in a deterministic order.

Each circle mark takes a tail or an arrowhead; an assignment survives when it yields a valid MAG whose canonical PAG (:func:pag_of_mag) is exactly p. With selection=False (the default) only the members without undirected edges count: the reading every consumer that assumes no selection bias quantifies over, identification among them. Pass selection=True for the full Markov equivalence class, where a circle-circle edge may also resolve to an undirected selection edge. A PAG that itself carries selection marks (undirected or circle-tail edges) has no selection-free members and raises under the default. Exponential in the number of circle marks by design, which is fine at the scale the library targets; max_members raises once the count would pass it (no silent truncation). mag_of_pag(p) is always among the members.

Source code in src/ergodic/graph/_pag.py

def mags_of_pag(
    p: MixedGraph, *, max_members: int | None = None, selection: bool = False
) -> tuple[MixedGraph, ...]:
    """Every MAG in the PAG's equivalence class, in a deterministic order.

    Each circle mark takes a tail or an arrowhead; an assignment survives
    when it yields a valid MAG whose canonical PAG (:func:`pag_of_mag`) is
    exactly ``p``. With ``selection=False`` (the default) only the members
    without undirected edges count: the reading every consumer that assumes
    no selection bias quantifies over, identification among them. Pass
    ``selection=True`` for the full Markov equivalence class, where a
    circle-circle edge may also resolve to an undirected selection edge. A
    PAG that itself carries selection marks (undirected or circle-tail
    edges) has no selection-free members and raises under the default.
    Exponential in the number of circle marks by design, which is fine at
    the scale the library targets; ``max_members`` raises once the count
    would pass it (no silent truncation). ``mag_of_pag(p)`` is always among
    the members.
    """
    if not p.is_pag():
        raise ValueError("mags_of_pag needs a complete PAG; check is_pag or build with pag()")
    if not selection:
        for e in p.edges:
            if e.is_undirected or {e.u_mark, e.v_mark} == {Mark.TAIL, Mark.CIRCLE}:
                raise ValueError(
                    f"edge {e.glyph()!r} is a selection mark, so the class has no "
                    "selection-free members; pass selection=True to enumerate the "
                    "full Markov equivalence class"
                )
    slots: list[tuple[str, str]] = []
    for e in sorted(p.edges, key=lambda e: (e.u, e.v)):
        if e.u_mark is Mark.CIRCLE:
            slots.append((e.u, e.v))
        if e.v_mark is Mark.CIRCLE:
            slots.append((e.v, e.u))
    found: list[MixedGraph] = []
    for assignment in product((Mark.TAIL, Mark.ARROW), repeat=len(slots)):
        marks = _mark_maps(p)
        for (node, other), mark in zip(slots, assignment, strict=True):
            marks[node][other] = mark
        candidate = replace(
            p,
            edges=frozenset(Edge.make(e.u, marks[e.u][e.v], e.v, marks[e.v][e.u]) for e in p.edges),
            kind=GraphKind.MAG,
        )
        if not selection and any(e.is_undirected for e in candidate.edges):
            continue
        if not candidate.is_mag() or pag_of_mag(candidate).edges != p.edges:
            continue
        if max_members is not None and len(found) >= max_members:
            raise ValueError(
                f"the class has more than {max_members} members; raise max_members "
                "or work on the PAG itself"
            )
        found.append(candidate)
    return tuple(found)

meek_closure(pdag, *, forbidden=None, ambiguous=None)

Apply Meek's rules R1 to R4 to a PDAG until fixpoint.

Each rule turns an undirected edge into a directed one whenever the reverse orientation would create a cycle or a new v-structure in every consistent extension. forbidden(u, v) returning True blocks the orientation u -> v: the rule simply does not fire that way and the edge stays undirected unless another rule orients it.

ambiguous lists unshielded triples (x, z, y), z the middle node and the ends unordered, that the caller cannot certify as definite non-colliders (PC's conservative and majority collider rules leave such triples behind). A definite non-collider is exactly what R1 and R3 consume: R1 reads c -> a -- b with c and b nonadjacent as compelling a -> b only because the collider c -> a <- b was ruled out, and R3 rests on its two converging parents never colliding at the undirected partner. Both rules skip a gated triple; R2 (a cycle rule) and R4 run unchanged. Default None gates nothing.

The result is returned with kind mixed; re-face it with :meth:MixedGraph.as_kind when more is known.

Source code in src/ergodic/graph/_meek.py

def meek_closure(
    pdag: MixedGraph,
    *,
    forbidden: Callable[[str, str], bool] | None = None,
    ambiguous: Collection[tuple[str, str, str]] | None = None,
) -> MixedGraph:
    """Apply Meek's rules R1 to R4 to a PDAG until fixpoint.

    Each rule turns an undirected edge into a directed one whenever the
    reverse orientation would create a cycle or a new v-structure in every
    consistent extension. ``forbidden(u, v)`` returning ``True`` blocks the
    orientation ``u -> v``: the rule simply does not fire that way and the
    edge stays undirected unless another rule orients it.

    ``ambiguous`` lists unshielded triples ``(x, z, y)``, ``z`` the middle
    node and the ends unordered, that the caller cannot certify as definite
    non-colliders (PC's conservative and majority collider rules leave such
    triples behind). A definite non-collider is exactly what R1 and R3
    consume: R1 reads ``c -> a -- b`` with ``c`` and ``b`` nonadjacent as
    compelling ``a -> b`` only because the collider ``c -> a <- b`` was ruled
    out, and R3 rests on its two converging parents never colliding at the
    undirected partner. Both rules skip a gated triple; R2 (a cycle rule)
    and R4 run unchanged. Default ``None`` gates nothing.

    The result is returned with kind ``mixed``; re-face it with
    :meth:`MixedGraph.as_kind` when more is known.
    """
    maps = _pdag_maps(pdag, "meek_closure")
    parents, children, undirected = maps
    blocked: Callable[[str, str], bool] = forbidden if forbidden is not None else lambda u, v: False
    gated = frozenset(
        (min(x, y), z, max(x, y)) for x, z, y in (ambiguous if ambiguous is not None else ())
    )
    graph = pdag
    changed = True
    while changed:
        changed = False
        pairs = sorted({tuple(sorted((x, y))) for x in undirected for y in undirected[x]})
        for u, v in pairs:
            if v not in undirected[u]:
                continue  # oriented earlier in this pass
            for a, b in ((u, v), (v, u)):
                if blocked(a, b) or not _compelled(a, b, maps, gated):
                    continue
                graph = graph.orient(a, b)
                undirected[a].discard(b)
                undirected[b].discard(a)
                children[a].add(b)
                parents[b].add(a)
                changed = True
                break
    return replace(graph, kind=GraphKind.MIXED)

mixed(edges=(), *, nodes=(), latent=(), selection=(), kind=GraphKind.MIXED)

Build a mixed graph of the given kind and check it satisfies that kind.

Source code in src/ergodic/graph/_graph.py

def mixed(
    edges: Iterable[EdgeSpec] = (),
    *,
    nodes: Iterable[str] = (),
    latent: Iterable[str] = (),
    selection: Iterable[str] = (),
    kind: GraphKind = GraphKind.MIXED,
) -> MixedGraph:
    """Build a mixed graph of the given ``kind`` and check it satisfies that kind."""
    built = _build_edges(edges)
    node_set = frozenset(nodes) | frozenset(n for e in built for n in (e.u, e.v))
    graph = MixedGraph(
        nodes=node_set,
        edges=built,
        latent=frozenset(latent),
        selection=frozenset(selection),
        kind=kind,
    )
    graph.validate()
    return graph

pag(edges=(), **kwargs)

Build a partial ancestral graph, validated as a complete PAG.

The marks must be exactly what an FCI run would leave: every invariant mark oriented (closed under the orientation rules) and a circle wherever the equivalence class disagrees. Build with :func:mixed instead when you want circle marks without that guarantee.

Source code in src/ergodic/graph/_graph.py

def pag(edges: Iterable[EdgeSpec] = (), **kwargs: Any) -> MixedGraph:
    """Build a partial ancestral graph, validated as a complete PAG.

    The marks must be exactly what an FCI run would leave: every invariant
    mark oriented (closed under the orientation rules) and a circle wherever
    the equivalence class disagrees. Build with :func:`mixed` instead when
    you want circle marks without that guarantee.
    """
    return mixed(edges, kind=GraphKind.PAG, **kwargs)

pag_of_mag(m)

The canonical PAG of a MAG: the marks every Markov-equivalent MAG shares.

Same skeleton, every mark a circle; unshielded colliders copied from the MAG (they are class invariants); then the rule closure, with R4's discriminating-path question answered by the MAG's own marks. Zhang (2008) proved the ten rules complete with latent confounders and selection bias both, so the result carries exactly the invariant marks. The input should be maximal; is_mag cannot check that, and a non-maximal input gives a PAG for the wrong class.

Source code in src/ergodic/graph/_pag.py

def pag_of_mag(m: MixedGraph) -> MixedGraph:
    """The canonical PAG of a MAG: the marks every Markov-equivalent MAG shares.

    Same skeleton, every mark a circle; unshielded colliders copied from the
    MAG (they are class invariants); then the rule closure, with R4's
    discriminating-path question answered by the MAG's own marks. Zhang
    (2008) proved the ten rules complete with latent confounders and
    selection bias both, so the result carries exactly the invariant marks.
    The input should be maximal; ``is_mag`` cannot check that, and a
    non-maximal input gives a PAG for the wrong class.
    """
    if not m.is_mag():
        raise ValueError("pag_of_mag needs a MAG; project or validate the graph first")
    marks: _Marks = {n: {} for n in m.nodes}
    for e in m.edges:
        marks[e.u][e.v] = Mark.CIRCLE
        marks[e.v][e.u] = Mark.CIRCLE
    adjacency = {n: sorted(marks[n]) for n in marks}
    for z in sorted(m.nodes):
        for x, y in combinations(adjacency[z], 2):
            if y in marks[x]:
                continue  # shielded
            if m.mark_at(x, z, z) is Mark.ARROW and m.mark_at(y, z, z) is Mark.ARROW:
                marks[z][x] = Mark.ARROW
                marks[z][y] = Mark.ARROW

    def from_mag(d: str, a: str, b: str, c: str) -> bool | None:
        return m.mark_at(a, b, b) is Mark.ARROW and m.mark_at(b, c, b) is Mark.ARROW

    skeleton = MixedGraph(
        nodes=m.nodes,
        edges=frozenset(Edge.make(e.u, marks[e.u][e.v], e.v, marks[e.v][e.u]) for e in m.edges),
        kind=GraphKind.MIXED,
    )
    closed = fci_closure(skeleton, discriminating=from_mag)
    # replace, not as_kind: validating PAG kind would re-enter this function
    # through the is_pag round trip.
    return replace(closed, kind=GraphKind.PAG)

possible_ancestors(graph, nodes)

Nodes with a possibly directed path into nodes.

The mirror of :func:possible_descendants: w counts when some member of the class can orient a path from w all the way into a seed. Each seed is excluded from its own walk.

Source code in src/ergodic/graph/_pag.py

def possible_ancestors(graph: MixedGraph, nodes: str | Iterable[str]) -> frozenset[str]:
    """Nodes with a possibly directed path into ``nodes``.

    The mirror of :func:`possible_descendants`: ``w`` counts when some
    member of the class can orient a path from ``w`` all the way into a
    seed. Each seed is excluded from its own walk.
    """
    seeds = _seed_names(graph, nodes)
    marks = _mark_maps(graph)
    out: set[str] = set()
    for seed in sorted(seeds):
        seen = {seed}
        stack = [seed]
        while stack:
            cur = stack.pop()
            for nxt in marks[cur]:
                if nxt not in seen and marks[nxt][cur] is not Mark.ARROW:
                    seen.add(nxt)
                    stack.append(nxt)
        out |= seen - {seed}
    return frozenset(out)

possible_descendants(graph, nodes)

Nodes reachable from nodes along a possibly directed path.

A path is possibly directed when no edge on it carries an arrowhead at the endpoint nearer the start, so some member of the graph's class can point every edge away from the start. On a DAG or ADMG this is plain descendants. Each seed is excluded from its own walk, mirroring :meth:MixedGraph.descendants.

Source code in src/ergodic/graph/_pag.py

def possible_descendants(graph: MixedGraph, nodes: str | Iterable[str]) -> frozenset[str]:
    """Nodes reachable from ``nodes`` along a possibly directed path.

    A path is possibly directed when no edge on it carries an arrowhead at
    the endpoint nearer the start, so some member of the graph's class can
    point every edge away from the start. On a DAG or ADMG this is plain
    ``descendants``. Each seed is excluded from its own walk, mirroring
    :meth:`MixedGraph.descendants`.
    """
    seeds = _seed_names(graph, nodes)
    marks = _mark_maps(graph)
    out: set[str] = set()
    for seed in sorted(seeds):
        seen = {seed}
        stack = [seed]
        while stack:
            cur = stack.pop()
            for nxt, mark in marks[cur].items():
                if nxt not in seen and mark is not Mark.ARROW:
                    seen.add(nxt)
                    stack.append(nxt)
        out |= seen - {seed}
    return frozenset(out)

topological_order(graph)

A topological order over the directed edges, tie-broken by node name.

For every directed edge a -> b, a precedes b. Bidirected edges impose no order. The name tie-break makes the order stable across runs, so estimands that read it render reproducibly. Raises on a graph whose directed part has a cycle.

Source code in src/ergodic/graph/_primitives.py

def topological_order(graph: MixedGraph) -> tuple[str, ...]:
    """A topological order over the directed edges, tie-broken by node name.

    For every directed edge ``a -> b``, ``a`` precedes ``b``. Bidirected edges
    impose no order. The name tie-break makes the order stable across runs, so
    estimands that read it render reproducibly. Raises on a graph whose
    directed part has a cycle.
    """
    order, acyclic = graph._kahn()
    if not acyclic:
        raise ValueError("topological_order needs an acyclic directed part; the graph has a cycle")
    return tuple(order)

visible_edges(graph)

The visible directed edges, as (tail, head) pairs (Zhang 2008).

A -> B is visible when some C not adjacent to B has an edge into A, or reaches A through a collider path into A whose every interior node is a parent of B. Visibility certifies the edge carries no hidden confounding in any represented world, which is what adjustment amenability on MAGs and PAGs consumes. The marks are read as given (arrowheads must be definite); on a DAG or CPDAG every directed edge is visible by causal sufficiency, and callers there should not need this function.

Source code in src/ergodic/graph/_pag.py

def visible_edges(graph: MixedGraph) -> frozenset[tuple[str, str]]:
    """The visible directed edges, as ``(tail, head)`` pairs (Zhang 2008).

    ``A -> B`` is visible when some ``C`` not adjacent to ``B`` has an edge
    into ``A``, or reaches ``A`` through a collider path into ``A`` whose
    every interior node is a parent of ``B``. Visibility certifies the edge
    carries no hidden confounding in any represented world, which is what
    adjustment amenability on MAGs and PAGs consumes. The marks are read as
    given (arrowheads must be definite); on a DAG or CPDAG every directed
    edge is visible by causal sufficiency, and callers there should not need
    this function.
    """
    marks = _mark_maps(graph)
    adj = {n: sorted(marks[n]) for n in marks}
    out: set[tuple[str, str]] = set()
    for a in sorted(marks):
        for b in adj[a]:
            if _is_directed(marks, a, b) and _is_visible(marks, adj, a, b):
                out.add((a, b))
    return frozenset(out)

Domain knowledge

Hard prior constraints (required and forbidden edges, tiers, roots, sinks) that discovery and inference read.

ergodic.knowledge

Domain knowledge: prior constraints on causal structure.

DomainKnowledge records what a user knows before looking at the data, the edges that must or must not exist, the temporal order of variables, which variables are roots or sinks, and which pairs share a hidden confounder. A skeleton restriction (restrict_to) caps which pairs may be adjacent at all. It is an immutable value object built up with small methods, and two knowledge objects combine with & (more knowledge means fewer admissible graphs).

Hard constraints are the focus for now. The energy and log_prior hooks are where soft, probabilistic knowledge will attach once the Bayesian work starts; they raise for the moment so nothing depends on a half-built scoring rule.

DomainKnowledge dataclass

Prior constraints on a causal graph over named variables.

Source code in src/ergodic/knowledge.py

@dataclass(frozen=True)
class DomainKnowledge:
    """Prior constraints on a causal graph over named variables."""

    variables: frozenset[str] = field(default_factory=frozenset)
    required: frozenset[tuple[str, str]] = field(default_factory=frozenset)
    forbidden: frozenset[tuple[str, str]] = field(default_factory=frozenset)
    precedes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
    roots: frozenset[str] = field(default_factory=frozenset)
    sinks: frozenset[str] = field(default_factory=frozenset)
    confounded: frozenset[frozenset[str]] = field(default_factory=frozenset)
    allowed_adjacencies: frozenset[frozenset[str]] | None = None

    # -- builders ------------------------------------------------------------

    def with_variables(self, *names: str) -> DomainKnowledge:
        return replace(self, variables=self.variables | frozenset(names))

    def require(self, cause: str, effect: str) -> DomainKnowledge:
        """Require a directed edge ``cause -> effect``."""
        return replace(
            self,
            required=self.required | {(cause, effect)},
            variables=self.variables | {cause, effect},
        )

    def forbid(self, cause: str, effect: str) -> DomainKnowledge:
        """Forbid the directed edge ``cause -> effect`` (the reverse stays allowed)."""
        return replace(
            self,
            forbidden=self.forbidden | {(cause, effect)},
            variables=self.variables | {cause, effect},
        )

    def forbid_adjacency(self, a: str, b: str) -> DomainKnowledge:
        """Forbid any edge between two variables, in either direction."""
        return replace(
            self,
            forbidden=self.forbidden | {(a, b), (b, a)},
            variables=self.variables | {a, b},
        )

    def order(self, before: str, after: str) -> DomainKnowledge:
        """Record that ``before`` precedes ``after`` in time, so ``after`` cannot cause it."""
        return replace(
            self,
            precedes=self.precedes | {(before, after)},
            variables=self.variables | {before, after},
        )

    def with_tiers(self, tiers: Sequence[Iterable[str]]) -> DomainKnowledge:
        """Add a tiered order: no variable in a later tier causes one in an earlier tier."""
        groups = [_pairs(t) for t in tiers]
        precedes = set(self.precedes)
        variables = set(self.variables)
        for i, earlier in enumerate(groups):
            variables |= earlier
            for later in groups[i + 1 :]:
                for a in earlier:
                    for b in later:
                        precedes.add((a, b))
        return replace(self, precedes=frozenset(precedes), variables=frozenset(variables))

    def root(self, *names: str) -> DomainKnowledge:
        """Mark variables as roots: no incoming directed edges."""
        return replace(
            self, roots=self.roots | frozenset(names), variables=self.variables | frozenset(names)
        )

    def sink(self, *names: str) -> DomainKnowledge:
        """Mark variables as sinks: no outgoing directed edges."""
        return replace(
            self, sinks=self.sinks | frozenset(names), variables=self.variables | frozenset(names)
        )

    def confound(self, a: str, b: str) -> DomainKnowledge:
        """Record that ``a`` and ``b`` share a hidden common cause."""
        if a == b:
            raise ValueError(f"a confounded pair needs two distinct variables, got {a!r} twice")
        return replace(
            self,
            confounded=self.confounded | {frozenset({a, b})},
            variables=self.variables | {a, b},
        )

    def forbid_between(self, group_a: Iterable[str], group_b: Iterable[str]) -> DomainKnowledge:
        """Forbid every directed edge from ``group_a`` into ``group_b``."""
        a = list(group_a)
        b = list(group_b)
        bans = {(x, y) for x in a for y in b}
        return replace(
            self, forbidden=self.forbidden | bans, variables=self.variables | set(a) | set(b)
        )

    def restrict_to(self, skeleton: MixedGraph | Iterable[tuple[str, str]]) -> DomainKnowledge:
        """Allow adjacencies only on the pairs of ``skeleton``.

        ``skeleton`` is a :class:`~ergodic.graph.MixedGraph` (its edges name
        the pairs, whatever the marks) or an iterable of ``(u, v)`` pairs. A
        pair outside the allow-list is banned in every form: directed either
        way, bidirected, or undirected. Restricting an already restricted
        object keeps the intersection, so a second call only narrows the
        support. An empty iterable allows no adjacencies at all.
        """
        from .graph import MixedGraph

        if isinstance(skeleton, MixedGraph):
            raw: Iterable[tuple[str, str]] = ((e.u, e.v) for e in skeleton.edges)
        else:
            raw = skeleton
        pairs: set[frozenset[str]] = set()
        names: set[str] = set()
        for u, v in raw:
            if u == v:
                raise ValueError(
                    f"a skeleton adjacency needs two distinct variables, got {u!r} twice"
                )
            pairs.add(frozenset({u, v}))
            names.update((u, v))
        allowed = frozenset(pairs)
        if self.allowed_adjacencies is not None:
            allowed &= self.allowed_adjacencies
        return replace(
            self,
            allowed_adjacencies=allowed,
            variables=self.variables | names,
        )

    def __and__(self, other: DomainKnowledge) -> DomainKnowledge:
        """Combine two knowledge objects by taking the union of their constraints.

        Skeleton allow-lists intersect: an adjacency survives only when both
        sides allow it, and ``None`` (unrestricted) defers to the other side.
        """
        if self.allowed_adjacencies is None:
            allowed = other.allowed_adjacencies
        elif other.allowed_adjacencies is None:
            allowed = self.allowed_adjacencies
        else:
            allowed = self.allowed_adjacencies & other.allowed_adjacencies
        return DomainKnowledge(
            variables=self.variables | other.variables,
            required=self.required | other.required,
            forbidden=self.forbidden | other.forbidden,
            precedes=self.precedes | other.precedes,
            roots=self.roots | other.roots,
            sinks=self.sinks | other.sinks,
            confounded=self.confounded | other.confounded,
            allowed_adjacencies=allowed,
        )

    combine = __and__

    # -- validation ----------------------------------------------------------

    def conflicts(self) -> list[str]:
        """Return a list of internal contradictions, empty when consistent."""
        out: list[str] = []
        for u, v in sorted(self.required & self.forbidden):
            out.append(f"edge {u}->{v} is both required and forbidden")
        for u, v in sorted(self.required):
            if (v, u) in self.precedes:
                out.append(f"required edge {u}->{v} contradicts order {v} before {u}")
            if v in self.roots:
                out.append(f"required edge {u}->{v} contradicts root {v}")
            if u in self.sinks:
                out.append(f"required edge {u}->{v} contradicts sink {u}")
        if self.allowed_adjacencies is not None:
            for u, v in sorted(self.required):
                if frozenset({u, v}) not in self.allowed_adjacencies:
                    out.append(f"required edge {u}->{v} lies outside the skeleton restriction")
            for pair in sorted(self.confounded, key=lambda p: sorted(p)):
                if pair not in self.allowed_adjacencies:
                    a, b = sorted(pair)
                    out.append(f"confounded pair {a},{b} lies outside the skeleton restriction")
        required_cycles = self._cycle_conflicts(self.required, "required edges")
        precedence_cycles = self._cycle_conflicts(self.precedes, "precedence")
        out.extend(required_cycles)
        out.extend(precedence_cycles)
        if not required_cycles and not precedence_cycles:
            # A required edge can close a loop through the precedence order even
            # when each relation is acyclic on its own, so check their union too.
            out.extend(
                self._cycle_conflicts(
                    self.required | self.precedes, "required edges and precedence together"
                )
            )
        return out

    @staticmethod
    def _cycle_conflicts(pairs: frozenset[tuple[str, str]], label: str) -> list[str]:
        # Kahn's algorithm, iterative so long chains do not hit the recursion limit.
        forward: dict[str, set[str]] = {}
        indegree: dict[str, int] = {}
        for a, b in pairs:
            forward.setdefault(a, set()).add(b)
            indegree.setdefault(a, 0)
            indegree[b] = indegree.get(b, 0) + 1
        ready = [n for n, degree in indegree.items() if degree == 0]
        removed = 0
        while ready:
            node = ready.pop()
            removed += 1
            for nxt in forward.get(node, ()):
                indegree[nxt] -= 1
                if indegree[nxt] == 0:
                    ready.append(nxt)
        return [] if removed == len(indegree) else [f"{label} contain a cycle"]

    def validate(self) -> None:
        """Raise :class:`InconsistentKnowledgeError` if the knowledge contradicts itself."""
        found = self.conflicts()
        if found:
            raise InconsistentKnowledgeError("; ".join(found))

    # -- projections used by algorithms --------------------------------------

    def _precedence_closure(self) -> frozenset[tuple[str, str]]:
        """The transitive closure of ``precedes``, walked with an explicit stack."""
        forward: dict[str, set[str]] = {}
        for a, b in self.precedes:
            forward.setdefault(a, set()).add(b)
        closed: set[tuple[str, str]] = set()
        for start in forward:
            seen: set[str] = set()
            stack = list(forward[start])
            while stack:
                node = stack.pop()
                if node in seen:
                    continue
                seen.add(node)
                stack.extend(forward.get(node, ()))
            seen.discard(start)
            closed.update((start, later) for later in seen)
        return frozenset(closed)

    def _permits(self, cause: str, effect: str, precedes: frozenset[tuple[str, str]]) -> bool:
        if (
            self.allowed_adjacencies is not None
            and frozenset({cause, effect}) not in self.allowed_adjacencies
        ):
            return False
        if (cause, effect) in self.forbidden:
            return False
        if (effect, cause) in precedes:
            return False
        return not (effect in self.roots or cause in self.sinks)

    def permits_edge(self, cause: str, effect: str) -> bool:
        """Return ``True`` if a directed edge ``cause -> effect`` is allowed.

        Precedence is read transitively: with ``A`` before ``B`` and ``B``
        before ``C``, an edge ``C -> A`` is not permitted. Under a skeleton
        restriction, a pair outside the allow-list fails in both directions.
        """
        return self._permits(cause, effect, self._precedence_closure())

    def requires_edge(self, cause: str, effect: str) -> bool:
        return (cause, effect) in self.required

    def precedence(self) -> frozenset[tuple[str, str]]:
        """Every ``(before, after)`` pair, with the order read transitively.

        The projection discovery reads to orient: ``(a, b)`` here means ``b``
        cannot be an ancestor of ``a``, which in a latent-variable search
        (FCI) pins an arrowhead at ``b`` on any surviving adjacency.
        """
        return self._precedence_closure()

    def violations(self, graph: MixedGraph) -> list[str]:
        """Return the ways ``graph`` breaks this knowledge, empty when consistent."""
        out: list[str] = []
        for u, v in sorted(self.required):
            edge = graph.get_edge(u, v)
            if edge is None or edge.directed_pair() != (u, v):
                out.append(f"missing required edge {u}->{v}")
        # A pair forbidden in both directions is an adjacency ban, so any edge
        # between the two (bidirected and undirected included) violates it.
        reported: set[frozenset[str]] = set()
        for u, v in sorted(self.forbidden):
            edge = graph.get_edge(u, v)
            if edge is None:
                continue
            if (v, u) in self.forbidden:
                pair = frozenset((u, v))
                if pair not in reported:
                    reported.add(pair)
                    out.append(f"present edge between {u} and {v}, forbidden in both directions")
            elif edge.directed_pair() == (u, v):
                out.append(f"present forbidden edge {u}->{v}")
        if self.allowed_adjacencies is not None:
            for edge in graph.edge_list():
                if frozenset({edge.u, edge.v}) not in self.allowed_adjacencies:
                    a, b = sorted((edge.u, edge.v))
                    out.append(
                        f"present edge between {a} and {b}, outside the skeleton restriction"
                    )
        for before, after in sorted(self._precedence_closure()):
            if (
                before in graph.nodes
                and after in graph.nodes
                and before in graph.descendants(after)
            ):
                out.append(f"order violated: {after} is an ancestor of {before}")
        for r in sorted(self.roots):
            if r in graph.nodes and graph.parents(r):
                out.append(f"root {r} has parents")
        for s in sorted(self.sinks):
            if s in graph.nodes and graph.children(s):
                out.append(f"sink {s} has children")
        for pair in self.confounded:
            a, b = sorted(pair)
            edge = graph.get_edge(a, b)
            if edge is None or not edge.is_bidirected:
                out.append(f"confounded pair {a},{b} lacks a bidirected edge")
        return out

    def is_consistent(self, graph: MixedGraph) -> bool:
        """Return ``True`` if ``graph`` satisfies every constraint."""
        return not self.violations(graph)

    def apply_to(self, graph: MixedGraph) -> MixedGraph:
        """Enforce the hard constraints on ``graph``: drop forbidden, add required edges.

        A pair forbidden in both directions loses its edge whatever the marks; a
        single-direction ban only drops the matching directed edge. Under a
        skeleton restriction, every edge on a pair outside the allow-list is
        dropped too. Edges left undirected or with a circle are oriented when
        only one direction stays permitted. The result is returned as a new
        graph.
        """
        out = graph
        if self.allowed_adjacencies is not None:
            for present in list(out.edges):
                if frozenset({present.u, present.v}) not in self.allowed_adjacencies:
                    out = out.remove_edge(present.u, present.v)
        for u, v in self.forbidden:
            edge = out.get_edge(u, v)
            if edge is None:
                continue
            if (v, u) in self.forbidden or edge.directed_pair() == (u, v):
                out = out.remove_edge(u, v)
        for u, v in self.required:
            out = out.orient(u, v)
        precedes = self._precedence_closure()
        for edge in list(out.edges):
            if edge.is_directed:
                continue
            forward_ok = self._permits(edge.u, edge.v, precedes)
            backward_ok = self._permits(edge.v, edge.u, precedes)
            if forward_ok and not backward_ok:
                out = out.orient(edge.u, edge.v)
            elif backward_ok and not forward_ok:
                out = out.orient(edge.v, edge.u)
        return out

    def edge_mask(
        self, variables: Sequence[str] | None = None
    ) -> tuple[npt.NDArray[np.bool_], list[str]]:
        """Return a boolean matrix of permitted directed edges and its variable order.

        ``mask[i, j]`` is ``True`` when an edge from variable ``i`` to variable
        ``j`` is allowed. The diagonal is ``False``, and a skeleton restriction
        zeroes both directions of every pair outside its allow-list. Useful as
        a search mask for score-based and continuous-optimization discovery.
        """
        import numpy as np

        labels = list(variables) if variables is not None else sorted(self.variables)
        size = len(labels)
        precedes = self._precedence_closure()
        mask = np.zeros((size, size), dtype=bool)
        for i, cause in enumerate(labels):
            for j, effect in enumerate(labels):
                if i != j and self._permits(cause, effect, precedes):
                    mask[i, j] = True
        return mask, labels

    # -- soft knowledge, not built yet ---------------------------------------

    def energy(self, graph: MixedGraph) -> float:
        """Penalty for ``graph`` under soft knowledge. Planned for the Bayesian phase."""
        raise NotImplementedError(
            "soft/energy-based knowledge is not implemented yet; it is planned for "
            "the Bayesian discovery work. Use the hard-constraint methods for now."
        )

    def log_prior(self, graph: MixedGraph) -> float:
        """Log prior of ``graph`` under soft knowledge. Planned for the Bayesian phase."""
        raise NotImplementedError(
            "soft/probabilistic knowledge is not implemented yet; it is planned for "
            "the Bayesian discovery work. Use the hard-constraint methods for now."
        )

    def __repr__(self) -> str:
        restricted = (
            ""
            if self.allowed_adjacencies is None
            else f", allowed_adjacencies={len(self.allowed_adjacencies)}"
        )
        return (
            "DomainKnowledge("
            f"vars={len(self.variables)}, required={len(self.required)}, "
            f"forbidden={len(self.forbidden)}, precedes={len(self.precedes)}, "
            f"roots={len(self.roots)}, sinks={len(self.sinks)}, "
            f"confounded={len(self.confounded)}{restricted})"
        )

__and__(other)

Combine two knowledge objects by taking the union of their constraints.

Skeleton allow-lists intersect: an adjacency survives only when both sides allow it, and None (unrestricted) defers to the other side.

Source code in src/ergodic/knowledge.py

def __and__(self, other: DomainKnowledge) -> DomainKnowledge:
    """Combine two knowledge objects by taking the union of their constraints.

    Skeleton allow-lists intersect: an adjacency survives only when both
    sides allow it, and ``None`` (unrestricted) defers to the other side.
    """
    if self.allowed_adjacencies is None:
        allowed = other.allowed_adjacencies
    elif other.allowed_adjacencies is None:
        allowed = self.allowed_adjacencies
    else:
        allowed = self.allowed_adjacencies & other.allowed_adjacencies
    return DomainKnowledge(
        variables=self.variables | other.variables,
        required=self.required | other.required,
        forbidden=self.forbidden | other.forbidden,
        precedes=self.precedes | other.precedes,
        roots=self.roots | other.roots,
        sinks=self.sinks | other.sinks,
        confounded=self.confounded | other.confounded,
        allowed_adjacencies=allowed,
    )

apply_to(graph)

Enforce the hard constraints on graph: drop forbidden, add required edges.

A pair forbidden in both directions loses its edge whatever the marks; a single-direction ban only drops the matching directed edge. Under a skeleton restriction, every edge on a pair outside the allow-list is dropped too. Edges left undirected or with a circle are oriented when only one direction stays permitted. The result is returned as a new graph.

Source code in src/ergodic/knowledge.py

def apply_to(self, graph: MixedGraph) -> MixedGraph:
    """Enforce the hard constraints on ``graph``: drop forbidden, add required edges.

    A pair forbidden in both directions loses its edge whatever the marks; a
    single-direction ban only drops the matching directed edge. Under a
    skeleton restriction, every edge on a pair outside the allow-list is
    dropped too. Edges left undirected or with a circle are oriented when
    only one direction stays permitted. The result is returned as a new
    graph.
    """
    out = graph
    if self.allowed_adjacencies is not None:
        for present in list(out.edges):
            if frozenset({present.u, present.v}) not in self.allowed_adjacencies:
                out = out.remove_edge(present.u, present.v)
    for u, v in self.forbidden:
        edge = out.get_edge(u, v)
        if edge is None:
            continue
        if (v, u) in self.forbidden or edge.directed_pair() == (u, v):
            out = out.remove_edge(u, v)
    for u, v in self.required:
        out = out.orient(u, v)
    precedes = self._precedence_closure()
    for edge in list(out.edges):
        if edge.is_directed:
            continue
        forward_ok = self._permits(edge.u, edge.v, precedes)
        backward_ok = self._permits(edge.v, edge.u, precedes)
        if forward_ok and not backward_ok:
            out = out.orient(edge.u, edge.v)
        elif backward_ok and not forward_ok:
            out = out.orient(edge.v, edge.u)
    return out

conflicts()

Return a list of internal contradictions, empty when consistent.

Source code in src/ergodic/knowledge.py

def conflicts(self) -> list[str]:
    """Return a list of internal contradictions, empty when consistent."""
    out: list[str] = []
    for u, v in sorted(self.required & self.forbidden):
        out.append(f"edge {u}->{v} is both required and forbidden")
    for u, v in sorted(self.required):
        if (v, u) in self.precedes:
            out.append(f"required edge {u}->{v} contradicts order {v} before {u}")
        if v in self.roots:
            out.append(f"required edge {u}->{v} contradicts root {v}")
        if u in self.sinks:
            out.append(f"required edge {u}->{v} contradicts sink {u}")
    if self.allowed_adjacencies is not None:
        for u, v in sorted(self.required):
            if frozenset({u, v}) not in self.allowed_adjacencies:
                out.append(f"required edge {u}->{v} lies outside the skeleton restriction")
        for pair in sorted(self.confounded, key=lambda p: sorted(p)):
            if pair not in self.allowed_adjacencies:
                a, b = sorted(pair)
                out.append(f"confounded pair {a},{b} lies outside the skeleton restriction")
    required_cycles = self._cycle_conflicts(self.required, "required edges")
    precedence_cycles = self._cycle_conflicts(self.precedes, "precedence")
    out.extend(required_cycles)
    out.extend(precedence_cycles)
    if not required_cycles and not precedence_cycles:
        # A required edge can close a loop through the precedence order even
        # when each relation is acyclic on its own, so check their union too.
        out.extend(
            self._cycle_conflicts(
                self.required | self.precedes, "required edges and precedence together"
            )
        )
    return out

confound(a, b)

Record that a and b share a hidden common cause.

Source code in src/ergodic/knowledge.py

def confound(self, a: str, b: str) -> DomainKnowledge:
    """Record that ``a`` and ``b`` share a hidden common cause."""
    if a == b:
        raise ValueError(f"a confounded pair needs two distinct variables, got {a!r} twice")
    return replace(
        self,
        confounded=self.confounded | {frozenset({a, b})},
        variables=self.variables | {a, b},
    )

edge_mask(variables=None)

Return a boolean matrix of permitted directed edges and its variable order.

mask[i, j] is True when an edge from variable i to variable j is allowed. The diagonal is False, and a skeleton restriction zeroes both directions of every pair outside its allow-list. Useful as a search mask for score-based and continuous-optimization discovery.

Source code in src/ergodic/knowledge.py

def edge_mask(
    self, variables: Sequence[str] | None = None
) -> tuple[npt.NDArray[np.bool_], list[str]]:
    """Return a boolean matrix of permitted directed edges and its variable order.

    ``mask[i, j]`` is ``True`` when an edge from variable ``i`` to variable
    ``j`` is allowed. The diagonal is ``False``, and a skeleton restriction
    zeroes both directions of every pair outside its allow-list. Useful as
    a search mask for score-based and continuous-optimization discovery.
    """
    import numpy as np

    labels = list(variables) if variables is not None else sorted(self.variables)
    size = len(labels)
    precedes = self._precedence_closure()
    mask = np.zeros((size, size), dtype=bool)
    for i, cause in enumerate(labels):
        for j, effect in enumerate(labels):
            if i != j and self._permits(cause, effect, precedes):
                mask[i, j] = True
    return mask, labels

energy(graph)

Penalty for graph under soft knowledge. Planned for the Bayesian phase.

Source code in src/ergodic/knowledge.py

def energy(self, graph: MixedGraph) -> float:
    """Penalty for ``graph`` under soft knowledge. Planned for the Bayesian phase."""
    raise NotImplementedError(
        "soft/energy-based knowledge is not implemented yet; it is planned for "
        "the Bayesian discovery work. Use the hard-constraint methods for now."
    )

forbid(cause, effect)

Forbid the directed edge cause -> effect (the reverse stays allowed).

Source code in src/ergodic/knowledge.py

def forbid(self, cause: str, effect: str) -> DomainKnowledge:
    """Forbid the directed edge ``cause -> effect`` (the reverse stays allowed)."""
    return replace(
        self,
        forbidden=self.forbidden | {(cause, effect)},
        variables=self.variables | {cause, effect},
    )

forbid_adjacency(a, b)

Forbid any edge between two variables, in either direction.

Source code in src/ergodic/knowledge.py

def forbid_adjacency(self, a: str, b: str) -> DomainKnowledge:
    """Forbid any edge between two variables, in either direction."""
    return replace(
        self,
        forbidden=self.forbidden | {(a, b), (b, a)},
        variables=self.variables | {a, b},
    )

forbid_between(group_a, group_b)

Forbid every directed edge from group_a into group_b.

Source code in src/ergodic/knowledge.py

def forbid_between(self, group_a: Iterable[str], group_b: Iterable[str]) -> DomainKnowledge:
    """Forbid every directed edge from ``group_a`` into ``group_b``."""
    a = list(group_a)
    b = list(group_b)
    bans = {(x, y) for x in a for y in b}
    return replace(
        self, forbidden=self.forbidden | bans, variables=self.variables | set(a) | set(b)
    )

is_consistent(graph)

Return True if graph satisfies every constraint.

Source code in src/ergodic/knowledge.py

def is_consistent(self, graph: MixedGraph) -> bool:
    """Return ``True`` if ``graph`` satisfies every constraint."""
    return not self.violations(graph)

log_prior(graph)

Log prior of graph under soft knowledge. Planned for the Bayesian phase.

Source code in src/ergodic/knowledge.py

def log_prior(self, graph: MixedGraph) -> float:
    """Log prior of ``graph`` under soft knowledge. Planned for the Bayesian phase."""
    raise NotImplementedError(
        "soft/probabilistic knowledge is not implemented yet; it is planned for "
        "the Bayesian discovery work. Use the hard-constraint methods for now."
    )

order(before, after)

Record that before precedes after in time, so after cannot cause it.

Source code in src/ergodic/knowledge.py

def order(self, before: str, after: str) -> DomainKnowledge:
    """Record that ``before`` precedes ``after`` in time, so ``after`` cannot cause it."""
    return replace(
        self,
        precedes=self.precedes | {(before, after)},
        variables=self.variables | {before, after},
    )

permits_edge(cause, effect)

Return True if a directed edge cause -> effect is allowed.

Precedence is read transitively: with A before B and B before C, an edge C -> A is not permitted. Under a skeleton restriction, a pair outside the allow-list fails in both directions.

Source code in src/ergodic/knowledge.py

def permits_edge(self, cause: str, effect: str) -> bool:
    """Return ``True`` if a directed edge ``cause -> effect`` is allowed.

    Precedence is read transitively: with ``A`` before ``B`` and ``B``
    before ``C``, an edge ``C -> A`` is not permitted. Under a skeleton
    restriction, a pair outside the allow-list fails in both directions.
    """
    return self._permits(cause, effect, self._precedence_closure())

precedence()

Every (before, after) pair, with the order read transitively.

The projection discovery reads to orient: (a, b) here means b cannot be an ancestor of a, which in a latent-variable search (FCI) pins an arrowhead at b on any surviving adjacency.

Source code in src/ergodic/knowledge.py

def precedence(self) -> frozenset[tuple[str, str]]:
    """Every ``(before, after)`` pair, with the order read transitively.

    The projection discovery reads to orient: ``(a, b)`` here means ``b``
    cannot be an ancestor of ``a``, which in a latent-variable search
    (FCI) pins an arrowhead at ``b`` on any surviving adjacency.
    """
    return self._precedence_closure()

require(cause, effect)

Require a directed edge cause -> effect.

Source code in src/ergodic/knowledge.py

def require(self, cause: str, effect: str) -> DomainKnowledge:
    """Require a directed edge ``cause -> effect``."""
    return replace(
        self,
        required=self.required | {(cause, effect)},
        variables=self.variables | {cause, effect},
    )

restrict_to(skeleton)

Allow adjacencies only on the pairs of skeleton.

skeleton is a :class:~ergodic.graph.MixedGraph (its edges name the pairs, whatever the marks) or an iterable of (u, v) pairs. A pair outside the allow-list is banned in every form: directed either way, bidirected, or undirected. Restricting an already restricted object keeps the intersection, so a second call only narrows the support. An empty iterable allows no adjacencies at all.

Source code in src/ergodic/knowledge.py

def restrict_to(self, skeleton: MixedGraph | Iterable[tuple[str, str]]) -> DomainKnowledge:
    """Allow adjacencies only on the pairs of ``skeleton``.

    ``skeleton`` is a :class:`~ergodic.graph.MixedGraph` (its edges name
    the pairs, whatever the marks) or an iterable of ``(u, v)`` pairs. A
    pair outside the allow-list is banned in every form: directed either
    way, bidirected, or undirected. Restricting an already restricted
    object keeps the intersection, so a second call only narrows the
    support. An empty iterable allows no adjacencies at all.
    """
    from .graph import MixedGraph

    if isinstance(skeleton, MixedGraph):
        raw: Iterable[tuple[str, str]] = ((e.u, e.v) for e in skeleton.edges)
    else:
        raw = skeleton
    pairs: set[frozenset[str]] = set()
    names: set[str] = set()
    for u, v in raw:
        if u == v:
            raise ValueError(
                f"a skeleton adjacency needs two distinct variables, got {u!r} twice"
            )
        pairs.add(frozenset({u, v}))
        names.update((u, v))
    allowed = frozenset(pairs)
    if self.allowed_adjacencies is not None:
        allowed &= self.allowed_adjacencies
    return replace(
        self,
        allowed_adjacencies=allowed,
        variables=self.variables | names,
    )

root(*names)

Mark variables as roots: no incoming directed edges.

Source code in src/ergodic/knowledge.py

def root(self, *names: str) -> DomainKnowledge:
    """Mark variables as roots: no incoming directed edges."""
    return replace(
        self, roots=self.roots | frozenset(names), variables=self.variables | frozenset(names)
    )

sink(*names)

Mark variables as sinks: no outgoing directed edges.

Source code in src/ergodic/knowledge.py

def sink(self, *names: str) -> DomainKnowledge:
    """Mark variables as sinks: no outgoing directed edges."""
    return replace(
        self, sinks=self.sinks | frozenset(names), variables=self.variables | frozenset(names)
    )

validate()

Raise :class:InconsistentKnowledgeError if the knowledge contradicts itself.

Source code in src/ergodic/knowledge.py

def validate(self) -> None:
    """Raise :class:`InconsistentKnowledgeError` if the knowledge contradicts itself."""
    found = self.conflicts()
    if found:
        raise InconsistentKnowledgeError("; ".join(found))

violations(graph)

Return the ways graph breaks this knowledge, empty when consistent.

Source code in src/ergodic/knowledge.py

def violations(self, graph: MixedGraph) -> list[str]:
    """Return the ways ``graph`` breaks this knowledge, empty when consistent."""
    out: list[str] = []
    for u, v in sorted(self.required):
        edge = graph.get_edge(u, v)
        if edge is None or edge.directed_pair() != (u, v):
            out.append(f"missing required edge {u}->{v}")
    # A pair forbidden in both directions is an adjacency ban, so any edge
    # between the two (bidirected and undirected included) violates it.
    reported: set[frozenset[str]] = set()
    for u, v in sorted(self.forbidden):
        edge = graph.get_edge(u, v)
        if edge is None:
            continue
        if (v, u) in self.forbidden:
            pair = frozenset((u, v))
            if pair not in reported:
                reported.add(pair)
                out.append(f"present edge between {u} and {v}, forbidden in both directions")
        elif edge.directed_pair() == (u, v):
            out.append(f"present forbidden edge {u}->{v}")
    if self.allowed_adjacencies is not None:
        for edge in graph.edge_list():
            if frozenset({edge.u, edge.v}) not in self.allowed_adjacencies:
                a, b = sorted((edge.u, edge.v))
                out.append(
                    f"present edge between {a} and {b}, outside the skeleton restriction"
                )
    for before, after in sorted(self._precedence_closure()):
        if (
            before in graph.nodes
            and after in graph.nodes
            and before in graph.descendants(after)
        ):
            out.append(f"order violated: {after} is an ancestor of {before}")
    for r in sorted(self.roots):
        if r in graph.nodes and graph.parents(r):
            out.append(f"root {r} has parents")
    for s in sorted(self.sinks):
        if s in graph.nodes and graph.children(s):
            out.append(f"sink {s} has children")
    for pair in self.confounded:
        a, b = sorted(pair)
        edge = graph.get_edge(a, b)
        if edge is None or not edge.is_bidirected:
            out.append(f"confounded pair {a},{b} lacks a bidirected edge")
    return out

with_tiers(tiers)

Add a tiered order: no variable in a later tier causes one in an earlier tier.

Source code in src/ergodic/knowledge.py

def with_tiers(self, tiers: Sequence[Iterable[str]]) -> DomainKnowledge:
    """Add a tiered order: no variable in a later tier causes one in an earlier tier."""
    groups = [_pairs(t) for t in tiers]
    precedes = set(self.precedes)
    variables = set(self.variables)
    for i, earlier in enumerate(groups):
        variables |= earlier
        for later in groups[i + 1 :]:
            for a in earlier:
                for b in later:
                    precedes.add((a, b))
    return replace(self, precedes=frozenset(precedes), variables=frozenset(variables))

InconsistentKnowledgeError

Bases: ValueError

Domain knowledge contradicts itself (for example a required edge cycle).

Source code in src/ergodic/knowledge.py

class InconsistentKnowledgeError(ValueError):
    """Domain knowledge contradicts itself (for example a required edge cycle)."""

Data

The typed dataset family across five shapes, and its accessors.

ergodic.data

Typed datasets: the data objects discovery and inference read.

Each shape embeds a column :class:Schema (continuous, categorical, ordinal, datetime, or excluded) and shares the accessors on :class:Dataset. The shapes are tabular, time-series, panel, and hierarchical time-series; the event log follows.

Build a dataset with :func:tabular, :func:timeseries, :func:panel, or :func:hts, each of which takes a pandas or polars dataframe, a dict of columns, or a numpy array.

ColumnType dataclass

The type of one column: its kind, and the categories when it is discrete.

Source code in src/ergodic/data/_schema.py

@dataclass(frozen=True)
class ColumnType:
    """The type of one column: its kind, and the categories when it is discrete."""

    name: str
    kind: VariableType
    categories: tuple[Any, ...] | None = None
    inferred: bool = True

    @property
    def ordered(self) -> bool:
        return self.kind is VariableType.ORDINAL

    @property
    def is_continuous(self) -> bool:
        return self.kind is VariableType.CONTINUOUS

    @property
    def is_categorical(self) -> bool:
        return self.kind is VariableType.CATEGORICAL

    @property
    def is_ordinal(self) -> bool:
        return self.kind is VariableType.ORDINAL

    @property
    def is_datetime(self) -> bool:
        return self.kind is VariableType.DATETIME

    @property
    def is_excluded(self) -> bool:
        return self.kind is VariableType.EXCLUDED

    @property
    def is_discrete(self) -> bool:
        """Categorical or ordinal: backed by a finite set of categories."""
        return self.kind in (VariableType.CATEGORICAL, VariableType.ORDINAL)

    @property
    def cardinality(self) -> int | None:
        return None if self.categories is None else len(self.categories)

    @property
    def is_binary(self) -> bool:
        return self.is_discrete and self.cardinality == 2

is_discrete property

Categorical or ordinal: backed by a finite set of categories.

Dataset dataclass

A typed column store. Built through the shape factories, not directly.

Source code in src/ergodic/data/_dataset.py

@dataclass(frozen=True, eq=False)
class Dataset:
    """A typed column store. Built through the shape factories, not directly."""

    schema: Schema
    n_rows: int
    _store: Mapping[str, npt.NDArray[Any]] = field(repr=False)

    # -- shape ---------------------------------------------------------------

    @property
    def names(self) -> tuple[str, ...]:
        return self.schema.names

    @property
    def n_columns(self) -> int:
        return len(self.schema)

    @property
    def shape(self) -> tuple[int, int]:
        return (self.n_rows, self.n_columns)

    def __len__(self) -> int:
        return self.n_rows

    @property
    def index_columns(self) -> tuple[str, ...]:
        """Columns that index the rows (time, entity, case), not modeled variables."""
        return ()

    @property
    def variables(self) -> tuple[str, ...]:
        """Column names the algorithms see: not the excluded ones, not the index."""
        index = set(self.index_columns)
        return tuple(name for name in self.schema.variables if name not in index)

    # -- column access -------------------------------------------------------

    def column(self, name: str) -> npt.NDArray[Any]:
        """A read-only copy of one column's raw values."""
        if name not in self.schema:
            raise KeyError(name)
        arr = np.array(self._store[name], copy=True)
        arr.setflags(write=False)
        return arr

    def codes(self, name: str) -> npt.NDArray[np.int64]:
        """Integer category codes for a categorical or ordinal column, -1 for missing.

        A non-missing value outside the column's categories raises a ``ValueError``.
        """
        col = self.schema[name]
        if not col.is_discrete:
            raise ValueError(f"column {name!r} is {col.kind.value}, not categorical or ordinal")
        return _encode(self._store[name], col.categories, col.ordered, name)

    def categories(self, name: str) -> tuple[Any, ...] | None:
        return self.schema[name].categories

    def matrix(self, names: Sequence[str] | None = None) -> npt.NDArray[np.float64]:
        """A float matrix of the named columns. Discrete columns become their codes,
        and a missing value (an NA category or NaT) becomes NaN."""
        chosen = list(names) if names is not None else list(self.variables)
        if not chosen:
            return np.empty((self.n_rows, 0), dtype=np.float64)
        cols: list[npt.NDArray[Any]] = []
        for name in chosen:
            col = self.schema[name]
            if col.is_continuous:
                cols.append(np.asarray(self._store[name], dtype=np.float64))
            elif col.is_discrete:
                encoded = self.codes(name).astype(np.float64)
                encoded[encoded == -1.0] = np.nan
                cols.append(encoded)
            elif col.is_datetime:
                stamps = np.asarray(self._store[name], dtype="datetime64[ns]")
                floats = stamps.view(np.int64).astype(np.float64)
                floats[np.isnat(stamps)] = np.nan
                cols.append(floats)
            else:
                raise ValueError(f"column {name!r} is excluded; it has no numeric form")
        return np.asarray(np.column_stack(cols), dtype=np.float64)

    def to_numpy(self) -> npt.NDArray[np.float64]:
        return self.matrix()

    def to_pandas(self) -> pd.DataFrame:
        import pandas as pd

        return pd.DataFrame({name: self._store[name] for name in self.schema.names})

    # -- missing data --------------------------------------------------------

    def missing(self) -> dict[str, int]:
        """Missing-value count per column."""
        return {name: int(_missing_mask(self._store[name]).sum()) for name in self.schema.names}

    @property
    def has_missing(self) -> bool:
        return any(count > 0 for count in self.missing().values())

    def dropna(self) -> Self:
        """Drop every row with a missing value in any column."""
        keep = np.ones(self.n_rows, dtype=bool)
        for arr in self._store.values():
            keep &= ~_missing_mask(arr)
        store = {name: arr[keep] for name, arr in self._store.items()}
        return replace(self, _store=store, n_rows=int(keep.sum()))

    # -- immutable edits -----------------------------------------------------

    def take(self, rows: Sequence[int] | npt.NDArray[np.int64]) -> Self:
        """A new dataset with the rows at the given positions, in the given order.

        Positions may repeat (a bootstrap resample) and may be negative
        (Python indexing). Row identity is all this knows about; shapes with
        row semantics (time series, panel) keep their index columns but it is
        the caller's job to resample in a way their structure survives.
        """
        idx = np.asarray(rows, dtype=np.int64)
        if idx.ndim != 1:
            raise ValueError(f"rows must be one-dimensional, got shape {idx.shape}")
        if idx.size and (int(idx.min()) < -self.n_rows or int(idx.max()) >= self.n_rows):
            raise IndexError(f"row positions out of range for {self.n_rows} rows")
        store = {name: arr[idx] for name, arr in self._store.items()}
        return replace(self, _store=store, n_rows=int(idx.size))

    def select(self, names: Iterable[str]) -> Self:
        """A new dataset with the named columns, plus any index columns, in order."""
        chosen = [str(n) for n in names]
        absent = [n for n in chosen if n not in self.schema]
        if absent:
            raise KeyError(f"no such columns: {absent}")
        for index_name in self.index_columns:
            if index_name not in chosen:
                chosen.append(index_name)
        schema = Schema(tuple(self.schema[n] for n in chosen))
        store = {n: self._store[n] for n in chosen}
        return replace(self, schema=schema, _store=store)

    def with_type(
        self,
        name: str,
        kind: VariableType | str | Ordinal,
        categories: Sequence[Any] | None = None,
    ) -> Self:
        """A new dataset with ``name`` retyped, categories read from data when needed."""
        if name not in self.schema:
            raise KeyError(name)
        new = _resolve_type(name, _unique_nonnull(self._store[name]), kind, categories)
        if new.is_discrete and (categories is not None or isinstance(kind, Ordinal)):
            # encode once now, so declared categories that miss a value fail here
            _encode(self._store[name], new.categories, new.ordered, name)
        cols = tuple(new if c.name == name else c for c in self.schema.columns)
        return replace(self, schema=Schema(cols))

    def exclude(self, *names: str) -> Self:
        """A new dataset with the named columns marked excluded."""
        out = self
        for name in names:
            out = out.with_type(name, VariableType.EXCLUDED)
        return out

    # -- graph alignment -----------------------------------------------------

    def alignment(self, graph: MixedGraph) -> list[str]:
        """Where the variables and the graph's nodes disagree, empty when they match."""
        nodes = set(graph.nodes)
        cols = set(self.variables)
        out = [f"graph node {n!r} has no column" for n in sorted(nodes - cols)]
        out += [f"column {c!r} is not a graph node" for c in sorted(cols - nodes)]
        return out

    def aligns_with(self, graph: MixedGraph) -> bool:
        """True when the variables are exactly the graph's nodes."""
        return not self.alignment(graph)

    def __repr__(self) -> str:
        parts = [
            f"{len(names)} {kind.value}"
            for kind in VariableType
            if (names := self.schema.of_kind(kind))
        ]
        body = ", ".join(parts)
        return f"{type(self).__name__}(rows={self.n_rows}, cols={self.n_columns}: {body})"

index_columns property

Columns that index the rows (time, entity, case), not modeled variables.

variables property

Column names the algorithms see: not the excluded ones, not the index.

alignment(graph)

Where the variables and the graph's nodes disagree, empty when they match.

Source code in src/ergodic/data/_dataset.py

def alignment(self, graph: MixedGraph) -> list[str]:
    """Where the variables and the graph's nodes disagree, empty when they match."""
    nodes = set(graph.nodes)
    cols = set(self.variables)
    out = [f"graph node {n!r} has no column" for n in sorted(nodes - cols)]
    out += [f"column {c!r} is not a graph node" for c in sorted(cols - nodes)]
    return out

aligns_with(graph)

True when the variables are exactly the graph's nodes.

Source code in src/ergodic/data/_dataset.py

def aligns_with(self, graph: MixedGraph) -> bool:
    """True when the variables are exactly the graph's nodes."""
    return not self.alignment(graph)

codes(name)

Integer category codes for a categorical or ordinal column, -1 for missing.

A non-missing value outside the column's categories raises a ValueError.

Source code in src/ergodic/data/_dataset.py

def codes(self, name: str) -> npt.NDArray[np.int64]:
    """Integer category codes for a categorical or ordinal column, -1 for missing.

    A non-missing value outside the column's categories raises a ``ValueError``.
    """
    col = self.schema[name]
    if not col.is_discrete:
        raise ValueError(f"column {name!r} is {col.kind.value}, not categorical or ordinal")
    return _encode(self._store[name], col.categories, col.ordered, name)

column(name)

A read-only copy of one column's raw values.

Source code in src/ergodic/data/_dataset.py

def column(self, name: str) -> npt.NDArray[Any]:
    """A read-only copy of one column's raw values."""
    if name not in self.schema:
        raise KeyError(name)
    arr = np.array(self._store[name], copy=True)
    arr.setflags(write=False)
    return arr

dropna()

Drop every row with a missing value in any column.

Source code in src/ergodic/data/_dataset.py

def dropna(self) -> Self:
    """Drop every row with a missing value in any column."""
    keep = np.ones(self.n_rows, dtype=bool)
    for arr in self._store.values():
        keep &= ~_missing_mask(arr)
    store = {name: arr[keep] for name, arr in self._store.items()}
    return replace(self, _store=store, n_rows=int(keep.sum()))

exclude(*names)

A new dataset with the named columns marked excluded.

Source code in src/ergodic/data/_dataset.py

def exclude(self, *names: str) -> Self:
    """A new dataset with the named columns marked excluded."""
    out = self
    for name in names:
        out = out.with_type(name, VariableType.EXCLUDED)
    return out

matrix(names=None)

A float matrix of the named columns. Discrete columns become their codes, and a missing value (an NA category or NaT) becomes NaN.

Source code in src/ergodic/data/_dataset.py

def matrix(self, names: Sequence[str] | None = None) -> npt.NDArray[np.float64]:
    """A float matrix of the named columns. Discrete columns become their codes,
    and a missing value (an NA category or NaT) becomes NaN."""
    chosen = list(names) if names is not None else list(self.variables)
    if not chosen:
        return np.empty((self.n_rows, 0), dtype=np.float64)
    cols: list[npt.NDArray[Any]] = []
    for name in chosen:
        col = self.schema[name]
        if col.is_continuous:
            cols.append(np.asarray(self._store[name], dtype=np.float64))
        elif col.is_discrete:
            encoded = self.codes(name).astype(np.float64)
            encoded[encoded == -1.0] = np.nan
            cols.append(encoded)
        elif col.is_datetime:
            stamps = np.asarray(self._store[name], dtype="datetime64[ns]")
            floats = stamps.view(np.int64).astype(np.float64)
            floats[np.isnat(stamps)] = np.nan
            cols.append(floats)
        else:
            raise ValueError(f"column {name!r} is excluded; it has no numeric form")
    return np.asarray(np.column_stack(cols), dtype=np.float64)

missing()

Missing-value count per column.

Source code in src/ergodic/data/_dataset.py

def missing(self) -> dict[str, int]:
    """Missing-value count per column."""
    return {name: int(_missing_mask(self._store[name]).sum()) for name in self.schema.names}

select(names)

A new dataset with the named columns, plus any index columns, in order.

Source code in src/ergodic/data/_dataset.py

def select(self, names: Iterable[str]) -> Self:
    """A new dataset with the named columns, plus any index columns, in order."""
    chosen = [str(n) for n in names]
    absent = [n for n in chosen if n not in self.schema]
    if absent:
        raise KeyError(f"no such columns: {absent}")
    for index_name in self.index_columns:
        if index_name not in chosen:
            chosen.append(index_name)
    schema = Schema(tuple(self.schema[n] for n in chosen))
    store = {n: self._store[n] for n in chosen}
    return replace(self, schema=schema, _store=store)

take(rows)

A new dataset with the rows at the given positions, in the given order.

Positions may repeat (a bootstrap resample) and may be negative (Python indexing). Row identity is all this knows about; shapes with row semantics (time series, panel) keep their index columns but it is the caller's job to resample in a way their structure survives.

Source code in src/ergodic/data/_dataset.py

def take(self, rows: Sequence[int] | npt.NDArray[np.int64]) -> Self:
    """A new dataset with the rows at the given positions, in the given order.

    Positions may repeat (a bootstrap resample) and may be negative
    (Python indexing). Row identity is all this knows about; shapes with
    row semantics (time series, panel) keep their index columns but it is
    the caller's job to resample in a way their structure survives.
    """
    idx = np.asarray(rows, dtype=np.int64)
    if idx.ndim != 1:
        raise ValueError(f"rows must be one-dimensional, got shape {idx.shape}")
    if idx.size and (int(idx.min()) < -self.n_rows or int(idx.max()) >= self.n_rows):
        raise IndexError(f"row positions out of range for {self.n_rows} rows")
    store = {name: arr[idx] for name, arr in self._store.items()}
    return replace(self, _store=store, n_rows=int(idx.size))

with_type(name, kind, categories=None)

A new dataset with name retyped, categories read from data when needed.

Source code in src/ergodic/data/_dataset.py

def with_type(
    self,
    name: str,
    kind: VariableType | str | Ordinal,
    categories: Sequence[Any] | None = None,
) -> Self:
    """A new dataset with ``name`` retyped, categories read from data when needed."""
    if name not in self.schema:
        raise KeyError(name)
    new = _resolve_type(name, _unique_nonnull(self._store[name]), kind, categories)
    if new.is_discrete and (categories is not None or isinstance(kind, Ordinal)):
        # encode once now, so declared categories that miss a value fail here
        _encode(self._store[name], new.categories, new.ordered, name)
    cols = tuple(new if c.name == name else c for c in self.schema.columns)
    return replace(self, schema=Schema(cols))

EventLog dataclass

Bases: Dataset

Events keyed by case, activity, and timestamp; one row per event.

Source code in src/ergodic/data/_eventlog.py

@dataclass(frozen=True, eq=False)
class EventLog(Dataset):
    """Events keyed by case, activity, and timestamp; one row per event."""

    case: str
    activity: str
    timestamp: str

    @property
    def index_columns(self) -> tuple[str, ...]:
        return (self.case, self.activity, self.timestamp)

    @property
    def cases(self) -> tuple[Any, ...]:
        return tuple(sorted(set(self._store[self.case].tolist())))

    @property
    def n_cases(self) -> int:
        return len(set(self._store[self.case].tolist()))

    @property
    def activities(self) -> tuple[str, ...]:
        return tuple(sorted(set(self._store[self.activity].tolist())))

    @property
    def n_activities(self) -> int:
        return len(set(self._store[self.activity].tolist()))

    def traces(self) -> dict[Any, tuple[str, ...]]:
        """The ordered activity sequence per case."""
        cases = self._store[self.case].tolist()
        acts = self._store[self.activity].tolist()
        out: dict[Any, list[str]] = {}
        for case_id, act in zip(cases, acts, strict=True):
            out.setdefault(case_id, []).append(act)
        return {case_id: tuple(seq) for case_id, seq in out.items()}

    def variants(self) -> Counter[tuple[str, ...]]:
        """How many cases follow each distinct activity sequence."""
        return Counter(self.traces().values())

    def directly_follows(self) -> Counter[tuple[str, str]]:
        """Counts of each ``(a, b)`` where ``b`` immediately follows ``a`` in a case."""
        out: Counter[tuple[str, str]] = Counter()
        for trace in self.traces().values():
            for first, second in pairwise(trace):
                out[(first, second)] += 1
        return out

    def __repr__(self) -> str:
        return (
            f"EventLog(cases={self.n_cases}, events={self.n_rows}, activities={self.n_activities})"
        )

directly_follows()

Counts of each (a, b) where b immediately follows a in a case.

Source code in src/ergodic/data/_eventlog.py

def directly_follows(self) -> Counter[tuple[str, str]]:
    """Counts of each ``(a, b)`` where ``b`` immediately follows ``a`` in a case."""
    out: Counter[tuple[str, str]] = Counter()
    for trace in self.traces().values():
        for first, second in pairwise(trace):
            out[(first, second)] += 1
    return out

traces()

The ordered activity sequence per case.

Source code in src/ergodic/data/_eventlog.py

def traces(self) -> dict[Any, tuple[str, ...]]:
    """The ordered activity sequence per case."""
    cases = self._store[self.case].tolist()
    acts = self._store[self.activity].tolist()
    out: dict[Any, list[str]] = {}
    for case_id, act in zip(cases, acts, strict=True):
        out.setdefault(case_id, []).append(act)
    return {case_id: tuple(seq) for case_id, seq in out.items()}

variants()

How many cases follow each distinct activity sequence.

Source code in src/ergodic/data/_eventlog.py

def variants(self) -> Counter[tuple[str, ...]]:
    """How many cases follow each distinct activity sequence."""
    return Counter(self.traces().values())

HierarchicalTSData dataclass

Bases: TimeSeriesData

A time-series whose columns form an aggregation hierarchy.

Source code in src/ergodic/data/_hts.py

@dataclass(frozen=True, eq=False)
class HierarchicalTSData(TimeSeriesData):
    """A time-series whose columns form an aggregation hierarchy."""

    hierarchy: Hierarchy

    def aggregate(self, node: str) -> npt.NDArray[np.float64]:
        """Build a node's series from its direct children, recursively.

        A child contributes its own column when it has one, else its aggregate,
        and the parent is the sum or mean of those child series. This is the same
        definition :meth:`coherence` checks, so the two cannot disagree.
        """
        if self.hierarchy.is_leaf(node):
            if node not in self._store:
                raise KeyError(f"no column for leaf {node!r}")
            return np.asarray(self._store[node], dtype=np.float64)
        parts = [
            np.asarray(self._store[kid], dtype=np.float64)
            if kid in self._store
            else self.aggregate(kid)
            for kid in self.hierarchy.children_of(node)
        ]
        stacked = np.vstack(parts)
        if self.hierarchy.aggregation == "mean":
            reduced = stacked.mean(axis=0)
        else:
            reduced = stacked.sum(axis=0)
        return np.asarray(reduced, dtype=np.float64)

    def coherence(self, rtol: float = 1e-05, atol: float = 1e-08) -> list[str]:
        """Where a present parent column does not match :meth:`aggregate` of its children."""
        out: list[str] = []
        for parent in self.hierarchy.internal_nodes:
            if parent not in self._store:
                continue
            try:
                expected = self.aggregate(parent)
            except KeyError:
                continue
            present = np.asarray(self._store[parent], dtype=np.float64)
            if not np.allclose(present, expected, rtol=rtol, atol=atol):
                kids = ", ".join(self.hierarchy.children_of(parent))
                out.append(f"{parent} != {self.hierarchy.aggregation}({kids})")
        return out

    def is_coherent(self, rtol: float = 1e-05, atol: float = 1e-08) -> bool:
        """True when every present parent equals the aggregate of its children."""
        return not self.coherence(rtol=rtol, atol=atol)

aggregate(node)

Build a node's series from its direct children, recursively.

A child contributes its own column when it has one, else its aggregate, and the parent is the sum or mean of those child series. This is the same definition :meth:coherence checks, so the two cannot disagree.

Source code in src/ergodic/data/_hts.py

def aggregate(self, node: str) -> npt.NDArray[np.float64]:
    """Build a node's series from its direct children, recursively.

    A child contributes its own column when it has one, else its aggregate,
    and the parent is the sum or mean of those child series. This is the same
    definition :meth:`coherence` checks, so the two cannot disagree.
    """
    if self.hierarchy.is_leaf(node):
        if node not in self._store:
            raise KeyError(f"no column for leaf {node!r}")
        return np.asarray(self._store[node], dtype=np.float64)
    parts = [
        np.asarray(self._store[kid], dtype=np.float64)
        if kid in self._store
        else self.aggregate(kid)
        for kid in self.hierarchy.children_of(node)
    ]
    stacked = np.vstack(parts)
    if self.hierarchy.aggregation == "mean":
        reduced = stacked.mean(axis=0)
    else:
        reduced = stacked.sum(axis=0)
    return np.asarray(reduced, dtype=np.float64)

coherence(rtol=1e-05, atol=1e-08)

Where a present parent column does not match :meth:aggregate of its children.

Source code in src/ergodic/data/_hts.py

def coherence(self, rtol: float = 1e-05, atol: float = 1e-08) -> list[str]:
    """Where a present parent column does not match :meth:`aggregate` of its children."""
    out: list[str] = []
    for parent in self.hierarchy.internal_nodes:
        if parent not in self._store:
            continue
        try:
            expected = self.aggregate(parent)
        except KeyError:
            continue
        present = np.asarray(self._store[parent], dtype=np.float64)
        if not np.allclose(present, expected, rtol=rtol, atol=atol):
            kids = ", ".join(self.hierarchy.children_of(parent))
            out.append(f"{parent} != {self.hierarchy.aggregation}({kids})")
    return out

is_coherent(rtol=1e-05, atol=1e-08)

True when every present parent equals the aggregate of its children.

Source code in src/ergodic/data/_hts.py

def is_coherent(self, rtol: float = 1e-05, atol: float = 1e-08) -> bool:
    """True when every present parent equals the aggregate of its children."""
    return not self.coherence(rtol=rtol, atol=atol)

Hierarchy dataclass

A tree of named series, with edges running parent to child, and an aggregation rule.

Source code in src/ergodic/data/_hierarchy.py

@dataclass(frozen=True)
class Hierarchy:
    """A tree of named series, with edges running parent to child, and an aggregation rule."""

    edges: frozenset[tuple[str, str]]
    aggregation: str = "sum"

    def __post_init__(self) -> None:
        if self.aggregation not in _AGGREGATIONS:
            raise ValueError(
                f"aggregation must be one of {_AGGREGATIONS}, got {self.aggregation!r}"
            )
        dups = sorted(child for child, n in Counter(c for _, c in self.edges).items() if n > 1)
        if dups:
            raise ValueError(f"nodes with more than one parent: {dups}")
        self._check_acyclic()

    def _check_acyclic(self) -> None:
        parent = {c: p for p, c in self.edges}
        for start in parent:
            node = start
            steps = 0
            while node in parent:
                node = parent[node]
                steps += 1
                if node == start or steps > len(parent):
                    raise ValueError("hierarchy has a cycle")

    @classmethod
    def tree(cls, adjacency: Mapping[str, Iterable[str]], aggregation: str = "sum") -> Hierarchy:
        """Build from a parent-to-children mapping, for example ``{"total": ["a", "b"]}``."""
        edges = frozenset((parent, child) for parent, kids in adjacency.items() for child in kids)
        return cls(edges=edges, aggregation=aggregation)

    @cached_property
    def _children(self) -> dict[str, tuple[str, ...]]:
        out: dict[str, list[str]] = {}
        for parent, child in self.edges:
            out.setdefault(parent, []).append(child)
        return {parent: tuple(sorted(kids)) for parent, kids in out.items()}

    @cached_property
    def _parent(self) -> dict[str, str]:
        return {child: parent for parent, child in self.edges}

    @cached_property
    def nodes(self) -> frozenset[str]:
        out: set[str] = set()
        for parent, child in self.edges:
            out.add(parent)
            out.add(child)
        return frozenset(out)

    def children_of(self, node: str) -> tuple[str, ...]:
        return self._children.get(node, ())

    def parent_of(self, node: str) -> str | None:
        return self._parent.get(node)

    def is_leaf(self, node: str) -> bool:
        return node not in self._children

    @property
    def roots(self) -> tuple[str, ...]:
        """Nodes with no parent."""
        return tuple(sorted(n for n in self.nodes if n not in self._parent))

    @property
    def leaves(self) -> tuple[str, ...]:
        """Nodes with no children: the base series."""
        return tuple(sorted(n for n in self.nodes if n not in self._children))

    @property
    def internal_nodes(self) -> tuple[str, ...]:
        """Nodes with children: the aggregates."""
        return tuple(sorted(self._children))

    def leaf_descendants(self, node: str) -> tuple[str, ...]:
        """The base series that aggregate up to ``node`` (``node`` itself when it is a leaf)."""
        if self.is_leaf(node):
            return (node,)
        out: list[str] = []
        for child in self.children_of(node):
            out.extend(self.leaf_descendants(child))
        return tuple(out)

    def levels(self) -> dict[str, int]:
        """Depth of each node, with roots at 0."""
        depth: dict[str, int] = {}
        frontier = [(root, 0) for root in self.roots]
        while frontier:
            node, current = frontier.pop()
            depth[node] = current
            for child in self.children_of(node):
                frontier.append((child, current + 1))
        return depth

    def __repr__(self) -> str:
        return (
            f"Hierarchy(nodes={len(self.nodes)}, leaves={len(self.leaves)}, "
            f"agg={self.aggregation!r})"
        )

internal_nodes property

Nodes with children: the aggregates.

leaves property

Nodes with no children: the base series.

roots property

Nodes with no parent.

leaf_descendants(node)

The base series that aggregate up to node (node itself when it is a leaf).

Source code in src/ergodic/data/_hierarchy.py

def leaf_descendants(self, node: str) -> tuple[str, ...]:
    """The base series that aggregate up to ``node`` (``node`` itself when it is a leaf)."""
    if self.is_leaf(node):
        return (node,)
    out: list[str] = []
    for child in self.children_of(node):
        out.extend(self.leaf_descendants(child))
    return tuple(out)

levels()

Depth of each node, with roots at 0.

Source code in src/ergodic/data/_hierarchy.py

def levels(self) -> dict[str, int]:
    """Depth of each node, with roots at 0."""
    depth: dict[str, int] = {}
    frontier = [(root, 0) for root in self.roots]
    while frontier:
        node, current = frontier.pop()
        depth[node] = current
        for child in self.children_of(node):
            frontier.append((child, current + 1))
    return depth

tree(adjacency, aggregation='sum') classmethod

Build from a parent-to-children mapping, for example {"total": ["a", "b"]}.

Source code in src/ergodic/data/_hierarchy.py

@classmethod
def tree(cls, adjacency: Mapping[str, Iterable[str]], aggregation: str = "sum") -> Hierarchy:
    """Build from a parent-to-children mapping, for example ``{"total": ["a", "b"]}``."""
    edges = frozenset((parent, child) for parent, kids in adjacency.items() for child in kids)
    return cls(edges=edges, aggregation=aggregation)

Ordinal dataclass

Mark a column ordinal, listing its categories in increasing order.

Use it in a type override: types={"stage": Ordinal(["I", "II", "III"])}.

Source code in src/ergodic/data/_schema.py

@dataclass(frozen=True)
class Ordinal:
    """Mark a column ordinal, listing its categories in increasing order.

    Use it in a type override: ``types={"stage": Ordinal(["I", "II", "III"])}``.
    """

    levels: Sequence[Any]

    def __post_init__(self) -> None:
        object.__setattr__(self, "levels", tuple(self.levels))

PanelData dataclass

Bases: Dataset

Entities observed over time, keyed by a unique (entity, time) pair.

Source code in src/ergodic/data/_panel.py

@dataclass(frozen=True, eq=False)
class PanelData(Dataset):
    """Entities observed over time, keyed by a unique (entity, time) pair."""

    entity: str
    time: str
    freq: str | None = None

    @property
    def index_columns(self) -> tuple[str, ...]:
        return (self.entity, self.time)

    @property
    def entities(self) -> tuple[Any, ...]:
        return tuple(sorted(set(self._store[self.entity].tolist())))

    @property
    def periods(self) -> tuple[Any, ...]:
        return tuple(sorted(set(self._store[self.time].tolist())))

    @property
    def n_entities(self) -> int:
        return len(set(self._store[self.entity].tolist()))

    @property
    def n_periods(self) -> int:
        return len(set(self._store[self.time].tolist()))

    @property
    def is_balanced(self) -> bool:
        """True when every entity is seen in every period (a full grid)."""
        return self.n_rows == self.n_entities * self.n_periods

    def time_invariant(self) -> tuple[str, ...]:
        """Variables whose value never changes within an entity."""
        frame = self.to_pandas()
        out: list[str] = []
        for name in self.variables:
            counts = frame.groupby(self.entity)[name].nunique(dropna=False)
            if bool((counts <= 1).all()):
                out.append(name)
        return tuple(out)

is_balanced property

True when every entity is seen in every period (a full grid).

time_invariant()

Variables whose value never changes within an entity.

Source code in src/ergodic/data/_panel.py

def time_invariant(self) -> tuple[str, ...]:
    """Variables whose value never changes within an entity."""
    frame = self.to_pandas()
    out: list[str] = []
    for name in self.variables:
        counts = frame.groupby(self.entity)[name].nunique(dropna=False)
        if bool((counts <= 1).all()):
            out.append(name)
    return tuple(out)

Schema dataclass

An ordered set of column types, queryable by name and by kind.

Source code in src/ergodic/data/_schema.py

@dataclass(frozen=True)
class Schema:
    """An ordered set of column types, queryable by name and by kind."""

    columns: tuple[ColumnType, ...]

    @cached_property
    def _by_name(self) -> dict[str, ColumnType]:
        return {c.name: c for c in self.columns}

    def __getitem__(self, name: str) -> ColumnType:
        return self._by_name[name]

    def __contains__(self, name: object) -> bool:
        return name in self._by_name

    def __iter__(self) -> Iterator[ColumnType]:
        return iter(self.columns)

    def __len__(self) -> int:
        return len(self.columns)

    def get(self, name: str) -> ColumnType | None:
        return self._by_name.get(name)

    @property
    def names(self) -> tuple[str, ...]:
        return tuple(c.name for c in self.columns)

    def of_kind(self, kind: VariableType) -> tuple[str, ...]:
        return tuple(c.name for c in self.columns if c.kind is kind)

    @property
    def continuous(self) -> tuple[str, ...]:
        return self.of_kind(VariableType.CONTINUOUS)

    @property
    def categorical(self) -> tuple[str, ...]:
        return self.of_kind(VariableType.CATEGORICAL)

    @property
    def ordinal(self) -> tuple[str, ...]:
        return self.of_kind(VariableType.ORDINAL)

    @property
    def datetime(self) -> tuple[str, ...]:
        return self.of_kind(VariableType.DATETIME)

    @property
    def excluded(self) -> tuple[str, ...]:
        return self.of_kind(VariableType.EXCLUDED)

    @property
    def discrete(self) -> tuple[str, ...]:
        """Categorical and ordinal columns together."""
        return tuple(c.name for c in self.columns if c.is_discrete)

    @property
    def variables(self) -> tuple[str, ...]:
        """Every column the algorithms see: all but the excluded ones."""
        return tuple(c.name for c in self.columns if not c.is_excluded)

    def with_type(
        self, name: str, kind: VariableType | str | Ordinal, categories: Sequence[Any] | None = None
    ) -> Schema:
        """Return a new schema with ``name`` retyped. Marks the column user-set."""
        if name not in self._by_name:
            raise KeyError(name)
        new = _resolve_type(name, self[name].categories, kind, categories)
        cols = tuple(new if c.name == name else c for c in self.columns)
        return Schema(cols)

    def __repr__(self) -> str:
        parts = [
            f"{len(names)} {kind.value}" for kind in VariableType if (names := self.of_kind(kind))
        ]
        return f"Schema({len(self)} columns: {', '.join(parts)})"

discrete property

Categorical and ordinal columns together.

variables property

Every column the algorithms see: all but the excluded ones.

with_type(name, kind, categories=None)

Return a new schema with name retyped. Marks the column user-set.

Source code in src/ergodic/data/_schema.py

def with_type(
    self, name: str, kind: VariableType | str | Ordinal, categories: Sequence[Any] | None = None
) -> Schema:
    """Return a new schema with ``name`` retyped. Marks the column user-set."""
    if name not in self._by_name:
        raise KeyError(name)
    new = _resolve_type(name, self[name].categories, kind, categories)
    cols = tuple(new if c.name == name else c for c in self.columns)
    return Schema(cols)

SchemaInferenceWarning

Bases: UserWarning

A column type was guessed from ambiguous data (for example integer levels).

Source code in src/ergodic/data/_schema.py

class SchemaInferenceWarning(UserWarning):
    """A column type was guessed from ambiguous data (for example integer levels)."""

TabularData dataclass

Bases: Dataset

Cross-sectional data: rows are independent units, with no index column.

Source code in src/ergodic/data/_tabular.py

@dataclass(frozen=True, eq=False)
class TabularData(Dataset):
    """Cross-sectional data: rows are independent units, with no index column."""

TimeSeriesData dataclass

Bases: Dataset

Series ordered by a time column, with an optional sampling frequency.

Source code in src/ergodic/data/_timeseries.py

@dataclass(frozen=True, eq=False)
class TimeSeriesData(Dataset):
    """Series ordered by a time column, with an optional sampling frequency."""

    time: str
    freq: str | None

    @property
    def index_columns(self) -> tuple[str, ...]:
        return (self.time,)

    @property
    def times(self) -> npt.NDArray[Any]:
        """The time index, in ascending order."""
        return self.column(self.time)

times property

The time index, in ascending order.

VariableType

Bases: Enum

How a column is measured, which decides the methods that apply to it.

Source code in src/ergodic/data/_schema.py

class VariableType(Enum):
    """How a column is measured, which decides the methods that apply to it."""

    CONTINUOUS = "continuous"
    CATEGORICAL = "categorical"
    ORDINAL = "ordinal"
    DATETIME = "datetime"
    EXCLUDED = "excluded"

    def __str__(self) -> str:
        return self.value

event_log(data, *, case, activity, timestamp, columns=None, types=None, exclude=None, categorical_threshold=10)

Build an :class:EventLog, grouped by case and ordered by timestamp.

Source code in src/ergodic/data/_eventlog.py

def event_log(
    data: Any,
    *,
    case: str,
    activity: str,
    timestamp: str,
    columns: Sequence[str] | None = None,
    types: Mapping[str, VariableType | str | Ordinal] | None = None,
    exclude: Iterable[str] | None = None,
    categorical_threshold: int = 10,
) -> EventLog:
    """Build an :class:`EventLog`, grouped by case and ordered by timestamp."""
    store, n_rows, declared = to_columns(data, columns)
    for role, name in (("case", case), ("activity", activity), ("timestamp", timestamp)):
        if name not in store:
            raise KeyError(f"{role} column {name!r} not found in {sorted(store)}")
    overrides: dict[str, VariableType | str | Ordinal] = dict(types or {})
    overrides.setdefault(case, "categorical")
    overrides.setdefault(activity, "categorical")
    if timestamp not in overrides:
        overrides[timestamp] = "datetime" if store[timestamp].dtype.kind == "M" else "continuous"
    schema = infer_schema(
        store,
        n_rows=n_rows,
        declared=declared,
        overrides=overrides,
        exclude=exclude,
        categorical_threshold=categorical_threshold,
    )
    order = np.lexsort((store[timestamp], store[case]))
    ordered = {name: arr[order] for name, arr in store.items()}
    return EventLog(
        schema=schema,
        n_rows=n_rows,
        _store=ordered,
        case=case,
        activity=activity,
        timestamp=timestamp,
    )

hts(data, *, time, hierarchy, freq=None, columns=None, types=None, exclude=None, categorical_threshold=10)

Build a :class:HierarchicalTSData from series columns and a hierarchy.

Source code in src/ergodic/data/_hts.py

def hts(
    data: Any,
    *,
    time: str,
    hierarchy: Hierarchy | Mapping[str, Iterable[str]],
    freq: str | None = None,
    columns: Sequence[str] | None = None,
    types: Mapping[str, VariableType | str | Ordinal] | None = None,
    exclude: Iterable[str] | None = None,
    categorical_threshold: int = 10,
) -> HierarchicalTSData:
    """Build a :class:`HierarchicalTSData` from series columns and a hierarchy."""
    tree = hierarchy if isinstance(hierarchy, Hierarchy) else Hierarchy.tree(hierarchy)
    store, n_rows, declared = to_columns(data, columns)
    if time not in store:
        raise KeyError(f"time column {time!r} not found in {sorted(store)}")
    missing_leaves = [leaf for leaf in tree.leaves if leaf not in store]
    if missing_leaves:
        raise ValueError(f"hierarchy leaves have no column: {sorted(missing_leaves)}")
    overrides: dict[str, VariableType | str | Ordinal] = dict(types or {})
    if time not in overrides:
        overrides[time] = "datetime" if store[time].dtype.kind == "M" else "continuous"
    for node in tree.nodes:
        if node in store and node not in overrides:
            overrides[node] = "continuous"
    schema = infer_schema(
        store,
        n_rows=n_rows,
        declared=declared,
        overrides=overrides,
        exclude=exclude,
        categorical_threshold=categorical_threshold,
    )
    order = np.argsort(store[time], kind="stable")
    ordered = {name: arr[order] for name, arr in store.items()}
    return HierarchicalTSData(
        schema=schema, n_rows=n_rows, _store=ordered, time=time, freq=freq, hierarchy=tree
    )

infer_schema(store, *, n_rows, declared=None, overrides=None, exclude=None, categorical_threshold=10)

Build a schema from columns: user overrides beat source-declared beat inferred.

Source code in src/ergodic/data/_schema.py

def infer_schema(
    store: Mapping[str, npt.NDArray[Any]],
    *,
    n_rows: int,
    declared: Mapping[str, ColumnType] | None = None,
    overrides: Mapping[str, VariableType | str | Ordinal] | None = None,
    exclude: Iterable[str] | None = None,
    categorical_threshold: int = 10,
) -> Schema:
    """Build a schema from columns: user overrides beat source-declared beat inferred."""
    declared = declared or {}
    overrides = overrides or {}
    excluded = set(exclude or ())
    out: list[ColumnType] = []
    for name, values in store.items():
        if name in excluded:
            out.append(ColumnType(name, VariableType.EXCLUDED, inferred=False))
        elif name in overrides:
            spec = overrides[name]
            resolved = _resolve_type(name, _unique_nonnull(values), spec, None)
            if isinstance(spec, Ordinal):
                # encode once now, so a level list that misses a value fails at build
                _encode(values, resolved.categories, resolved.ordered, name)
            out.append(resolved)
        elif name in declared:
            out.append(declared[name])
        else:
            out.append(_infer_one(name, values, n_rows, categorical_threshold))
    return Schema(tuple(out))

panel(data, *, entity, time, freq=None, columns=None, types=None, exclude=None, categorical_threshold=10)

Build a :class:PanelData, keyed and sorted by (entity, time).

Source code in src/ergodic/data/_panel.py

def panel(
    data: Any,
    *,
    entity: str,
    time: str,
    freq: str | None = None,
    columns: Sequence[str] | None = None,
    types: Mapping[str, VariableType | str | Ordinal] | None = None,
    exclude: Iterable[str] | None = None,
    categorical_threshold: int = 10,
) -> PanelData:
    """Build a :class:`PanelData`, keyed and sorted by ``(entity, time)``."""
    store, n_rows, declared = to_columns(data, columns)
    for role, name in (("entity", entity), ("time", time)):
        if name not in store:
            raise KeyError(f"{role} column {name!r} not found in {sorted(store)}")
    overrides: dict[str, VariableType | str | Ordinal] = dict(types or {})
    if time not in overrides:
        overrides[time] = "datetime" if store[time].dtype.kind == "M" else "continuous"
    if entity not in overrides:
        overrides[entity] = "categorical"
    schema = infer_schema(
        store,
        n_rows=n_rows,
        declared=declared,
        overrides=overrides,
        exclude=exclude,
        categorical_threshold=categorical_threshold,
    )
    pairs = set(zip(store[entity].tolist(), store[time].tolist(), strict=True))
    if len(pairs) != n_rows:
        raise ValueError(
            "(entity, time) pairs are not unique; a panel needs one row per entity-period"
        )
    order = np.lexsort((store[time], store[entity]))
    ordered = {name: arr[order] for name, arr in store.items()}
    return PanelData(
        schema=schema, n_rows=n_rows, _store=ordered, entity=entity, time=time, freq=freq
    )

tabular(data, *, columns=None, types=None, exclude=None, categorical_threshold=10)

Build a :class:TabularData from a dataframe, a dict of columns, or an array.

Parameters:

Name	Type	Description	Default
data	Any	A pandas or polars dataframe, a mapping of name to column, or a 2-D numpy array (name the columns with columns).	required
columns	Sequence[str] | None	Column names, used when data is a bare numpy array.	None
types	Mapping[str, VariableType | str | Ordinal] | None	Per-column type overrides, for example {"age": "continuous", "stage": Ordinal(["I", "II", "III"])}.	None
exclude	Iterable[str] | None	Columns to mark excluded, so they are kept but not treated as variables (identifiers, free text).	None
categorical_threshold	int	Integer columns with at most this many distinct values are inferred categorical, with a warning.	10

Source code in src/ergodic/data/_tabular.py

def tabular(
    data: Any,
    *,
    columns: Sequence[str] | None = None,
    types: Mapping[str, VariableType | str | Ordinal] | None = None,
    exclude: Iterable[str] | None = None,
    categorical_threshold: int = 10,
) -> TabularData:
    """Build a :class:`TabularData` from a dataframe, a dict of columns, or an array.

    Args:
        data: A pandas or polars dataframe, a mapping of name to column, or a
            2-D numpy array (name the columns with ``columns``).
        columns: Column names, used when ``data`` is a bare numpy array.
        types: Per-column type overrides, for example
            ``{"age": "continuous", "stage": Ordinal(["I", "II", "III"])}``.
        exclude: Columns to mark excluded, so they are kept but not treated as
            variables (identifiers, free text).
        categorical_threshold: Integer columns with at most this many distinct
            values are inferred categorical, with a warning.
    """
    store, n_rows, declared = to_columns(data, columns)
    schema = infer_schema(
        store,
        n_rows=n_rows,
        declared=declared,
        overrides=types,
        exclude=exclude,
        categorical_threshold=categorical_threshold,
    )
    return TabularData(schema=schema, n_rows=n_rows, _store=store)

timeseries(data, *, time, freq=None, columns=None, types=None, exclude=None, categorical_threshold=10)

Build a :class:TimeSeriesData, sorting rows by the time column.

The time column is typed from its data (datetime, else continuous) unless you override it, so an integer step index is not mistaken for a category.

Source code in src/ergodic/data/_timeseries.py

def timeseries(
    data: Any,
    *,
    time: str,
    freq: str | None = None,
    columns: Sequence[str] | None = None,
    types: Mapping[str, VariableType | str | Ordinal] | None = None,
    exclude: Iterable[str] | None = None,
    categorical_threshold: int = 10,
) -> TimeSeriesData:
    """Build a :class:`TimeSeriesData`, sorting rows by the ``time`` column.

    The time column is typed from its data (datetime, else continuous) unless
    you override it, so an integer step index is not mistaken for a category.
    """
    store, n_rows, declared = to_columns(data, columns)
    if time not in store:
        raise KeyError(f"time column {time!r} not found in {sorted(store)}")
    overrides: dict[str, VariableType | str | Ordinal] = dict(types or {})
    if time not in overrides:
        overrides[time] = "datetime" if store[time].dtype.kind == "M" else "continuous"
    schema = infer_schema(
        store,
        n_rows=n_rows,
        declared=declared,
        overrides=overrides,
        exclude=exclude,
        categorical_threshold=categorical_threshold,
    )
    order = np.argsort(store[time], kind="stable")
    ordered = {name: arr[order] for name, arr in store.items()}
    return TimeSeriesData(schema=schema, n_rows=n_rows, _store=ordered, time=time, freq=freq)

Identification

Graph to estimand: adjustment, front-door, instruments, the general ID algorithm, and mediation.

ergodic.identification

Identification: turn a causal graph into an estimand.

Given a graph, a treatment, and an outcome, the question is whether the interventional distribution P(Y | do(X)) can be written in terms of the observational distribution, and if so, how. The work is graph-only and deterministic: no data, no statistics, no estimator choices yet.

The named criteria (adjustment, front-door, instruments) live in :mod:._criteria; the :class:Identification record in :mod:._record. :func:identify_effect tries them in order, falling through to the general do-calculus ID algorithm on a single graph.

The general engine is exposed too: :func:identify_distribution runs ID (or IDC) for any query and returns an :class:IDResult carrying the typed estimand :class:Functional (or the :class:Hedge witness when not identifiable). :func:ida covers an equivalence class by enumerating members.

Functional

The sealed base of the estimand algebra. Use the frozen subclasses below.

Three operations are common to every node: :meth:render (the readable estimand string), :meth:free_vars (the variable scope, used by recognition and the plug-in evaluator), and :meth:simplify (conservative rewrites only). Subclasses are frozen value objects, so two structurally equal functionals are equal and hash alike.

Source code in src/ergodic/identification/_functional.py

class Functional:
    """The sealed base of the estimand algebra. Use the frozen subclasses below.

    Three operations are common to every node: :meth:`render` (the readable
    estimand string), :meth:`free_vars` (the variable scope, used by recognition
    and the plug-in evaluator), and :meth:`simplify` (conservative rewrites
    only). Subclasses are frozen value objects, so two structurally equal
    functionals are equal and hash alike.
    """

    def render(self) -> str:
        """The readable estimand string, e.g. ``sum_z P(Y | X, z) P(z)``."""
        raise NotImplementedError

    def free_vars(self) -> frozenset[str]:
        """The variable scope: heads and givens, minus bound marginal indices."""
        raise NotImplementedError

    def simplify(self) -> Functional:
        """A conservatively simplified functional (safe rewrites only).

        Flatten nested products, drop a marginal over an empty index, drop a
        factor that is identically one. No cancellations, no clever algebra: two
        runs render the same estimand.
        """
        raise NotImplementedError

free_vars()

The variable scope: heads and givens, minus bound marginal indices.

Source code in src/ergodic/identification/_functional.py

def free_vars(self) -> frozenset[str]:
    """The variable scope: heads and givens, minus bound marginal indices."""
    raise NotImplementedError

render()

The readable estimand string, e.g. sum_z P(Y | X, z) P(z).

Source code in src/ergodic/identification/_functional.py

def render(self) -> str:
    """The readable estimand string, e.g. ``sum_z P(Y | X, z) P(z)``."""
    raise NotImplementedError

simplify()

A conservatively simplified functional (safe rewrites only).

Flatten nested products, drop a marginal over an empty index, drop a factor that is identically one. No cancellations, no clever algebra: two runs render the same estimand.

Source code in src/ergodic/identification/_functional.py

def simplify(self) -> Functional:
    """A conservatively simplified functional (safe rewrites only).

    Flatten nested products, drop a marginal over an empty index, drop a
    factor that is identically one. No cancellations, no clever algebra: two
    runs render the same estimand.
    """
    raise NotImplementedError

Hedge dataclass

Bases: Exception

The structural witness that an effect is not identifiable (a C-forest pair).

ID fails at a subproblem whose whole graph is a single c-component over nodes and whose treatment-free part is the district district. That pair is a hedge in the sense of Shpitser-Pearl: it is the obstruction that proves P(y | do(x)) is not a function of the observational distribution. The SCM bounds layer reads this object to decide where to return bounds instead of a point.

The witness is the (nodes, district) summary of the failing subproblem, enough to mark non-identifiability and to seed the bounds. The explicit minimal C-forest pair (the rooted forests of Shpitser-Pearl) is not extracted yet; that is a refinement for later.

Source code in src/ergodic/identification/_id.py

@dataclass(frozen=True)
class Hedge(Exception):
    """The structural witness that an effect is not identifiable (a C-forest pair).

    ID fails at a subproblem whose whole graph is a single c-component over
    ``nodes`` and whose treatment-free part is the district ``district``. That
    pair is a hedge in the sense of Shpitser-Pearl: it is the obstruction that
    proves ``P(y | do(x))`` is not a function of the observational distribution.
    The SCM bounds layer reads this object to decide where to return bounds
    instead of a point.

    The witness is the ``(nodes, district)`` summary of the failing subproblem,
    enough to mark non-identifiability and to seed the bounds. The explicit
    minimal C-forest pair (the rooted forests of Shpitser-Pearl) is not extracted
    yet; that is a refinement for later.
    """

    nodes: frozenset[str]
    district: frozenset[str]

    def render(self) -> str:
        """A human-readable reason."""
        return (
            f"not identifiable: a hedge formed by {sorted(self.nodes)} "
            f"and the c-component {sorted(self.district)}"
        )

    def __str__(self) -> str:
        return self.render()

render()

A human-readable reason.

Source code in src/ergodic/identification/_id.py

def render(self) -> str:
    """A human-readable reason."""
    return (
        f"not identifiable: a hedge formed by {sorted(self.nodes)} "
        f"and the c-component {sorted(self.district)}"
    )

IDA dataclass

Per-member identification across an equivalence class.

members is the tuple of per-member :class:Identification results, in the enumerator's deterministic order. For a CPDAG that is one entry per DAG member; for a PAG or MAG handled through the class-level adjustment criterion it is a single entry standing for the whole class (the criterion's answer holds in every member). class_kind records the face that was enumerated.

The summary properties read the multiset: how many members identify, the distinct strategies and estimands seen, and whether the class agrees.

Source code in src/ergodic/identification/_ida.py

@dataclass(frozen=True)
class IDA:
    """Per-member identification across an equivalence class.

    ``members`` is the tuple of per-member :class:`Identification` results, in
    the enumerator's deterministic order. For a CPDAG that is one entry per DAG
    member; for a PAG or MAG handled through the class-level adjustment criterion
    it is a single entry standing for the whole class (the criterion's answer
    holds in every member). ``class_kind`` records the face that was enumerated.

    The summary properties read the multiset: how many members identify, the
    distinct strategies and estimands seen, and whether the class agrees.
    """

    treatment: frozenset[str]
    outcome: frozenset[str]
    class_kind: GraphKind
    members: tuple[Identification, ...]

    @property
    def n_members(self) -> int:
        """The number of members examined."""
        return len(self.members)

    @property
    def n_identifiable(self) -> int:
        """How many members have an identifiable effect."""
        return sum(1 for m in self.members if m.identifiable)

    @property
    def all_identifiable(self) -> bool:
        """``True`` when every member identifies."""
        return all(m.identifiable for m in self.members)

    @property
    def strategies(self) -> frozenset[str]:
        """The distinct strategies seen across members (e.g. ``{"adjustment", "id"}``)."""
        return frozenset(m.strategy for m in self.members)

    @property
    def estimands(self) -> frozenset[str]:
        """The distinct estimand strings across the identifiable members."""
        return frozenset(m.estimand for m in self.members if m.identifiable)

    @property
    def agree(self) -> bool:
        """``True`` when every member identifies and shares one estimand.

        The classic IDA reading: a single, member-invariant answer. ``False``
        when some member fails or the members disagree about the estimand.
        """
        return self.all_identifiable and len(self.estimands) == 1

    def __bool__(self) -> bool:
        return self.all_identifiable

    def __repr__(self) -> str:
        arrow = f"{sorted(self.treatment)} -> {sorted(self.outcome)}"
        return (
            f"IDA({arrow}, {self.class_kind.value}: "
            f"{self.n_identifiable}/{self.n_members} identify, "
            f"strategies={sorted(self.strategies)}, agree={self.agree})"
        )

agree property

True when every member identifies and shares one estimand.

The classic IDA reading: a single, member-invariant answer. False when some member fails or the members disagree about the estimand.

all_identifiable property

True when every member identifies.

estimands property

The distinct estimand strings across the identifiable members.

n_identifiable property

How many members have an identifiable effect.

n_members property

The number of members examined.

strategies property

The distinct strategies seen across members (e.g. {"adjustment", "id"}).

IDResult dataclass

The outcome of running ID or IDC on a single graph.

identifiable is the verdict; functional is the estimand when it holds (None otherwise); hedge is the witness when it fails (None otherwise). query records the (outcome, do, given) that was asked.

Source code in src/ergodic/identification/_id.py

@dataclass(frozen=True)
class IDResult:
    """The outcome of running ID or IDC on a single graph.

    ``identifiable`` is the verdict; ``functional`` is the estimand when it
    holds (``None`` otherwise); ``hedge`` is the witness when it fails (``None``
    otherwise). ``query`` records the ``(outcome, do, given)`` that was asked.
    """

    identifiable: bool
    functional: Functional | None
    hedge: Hedge | None
    query: tuple[frozenset[str], frozenset[str], frozenset[str]]

    def __bool__(self) -> bool:
        return self.identifiable

    def render(self) -> str:
        if self.identifiable and self.functional is not None:
            return self.functional.render()
        if self.hedge is not None:
            return self.hedge.render()
        return "not identifiable"

    def __repr__(self) -> str:
        y, x, z = self.query
        cond_part = f", given={sorted(z)}" if z else ""
        head = f"P({sorted(y)} | do({sorted(x)}){cond_part})"
        if not self.identifiable:
            return f"IDResult(not identifiable, {head})"
        return f"IDResult({head} = {self.render()})"

Identification dataclass

The result of trying to identify an effect from a graph.

The named strategies ("null", "adjustment", "frontdoor", "iv") carry the set or instruments their estimator reads and a string estimand. When only the general do-calculus ID algorithm identifies an effect (outside the three named criteria, the napkin and Verma graphs and their kin), strategy is "id" and functional carries the typed estimand, with estimand set to functional.render(). A not-identifiable result built by the ID fallthrough carries the hedge witness. Existing strategies set neither, and estimand stays a string, so the field is a pure extension.

Source code in src/ergodic/identification/_record.py

@dataclass(frozen=True)
class Identification:
    """The result of trying to identify an effect from a graph.

    The named strategies (``"null"``, ``"adjustment"``, ``"frontdoor"``,
    ``"iv"``) carry the set or instruments their estimator reads and a string
    ``estimand``. When only the general do-calculus ID algorithm identifies an
    effect (outside the three named criteria, the napkin and Verma graphs and
    their kin), ``strategy`` is ``"id"`` and ``functional`` carries the typed
    estimand, with ``estimand`` set to ``functional.render()``. A
    not-identifiable result built by the ID fallthrough carries the ``hedge``
    witness. Existing strategies set neither, and ``estimand`` stays a string,
    so the field is a pure extension.
    """

    treatment: frozenset[str]
    outcome: frozenset[str]
    identifiable: bool
    strategy: str  # "null", "adjustment", "frontdoor", "iv", "id", or "none"
    adjustment_set: frozenset[str] | None = None
    mediators: frozenset[str] | None = None
    instruments: tuple[tuple[str, frozenset[str]], ...] = ()
    estimand: str = ""
    functional: Functional | None = None
    hedge: Hedge | None = None

    def __bool__(self) -> bool:
        return self.identifiable

    def __repr__(self) -> str:
        arrow = f"{sorted(self.treatment)} -> {sorted(self.outcome)}"
        if not self.identifiable:
            return f"Identification(not identifiable, {arrow})"
        return f"Identification({self.strategy}: {self.estimand})"

Marginal dataclass

Bases: Functional

Sum index out of body (a marginalization).

Source code in src/ergodic/identification/_functional.py

@dataclass(frozen=True)
class Marginal(Functional):
    """Sum ``index`` out of ``body`` (a marginalization)."""

    index: frozenset[str]
    body: Functional

    def render(self) -> str:
        if not self.index:
            return self.body.render()
        sums = " ".join(f"sum_{v.lower()}" for v in sorted(self.index))
        return f"{sums} {self.body.render()}"

    def free_vars(self) -> frozenset[str]:
        return self.body.free_vars() - self.index

    def simplify(self) -> Functional:
        body = self.body.simplify()
        # Only sum over indices the body actually mentions; if nothing remains,
        # the marginal is the identity.
        index = self.index & body.free_vars()
        if not index:
            return body
        if isinstance(body, Marginal):
            return Marginal(index | body.index, body.body)  # merge nested sums
        return Marginal(index, body)

MediationID dataclass

The result of trying to identify a mediation decomposition from a graph.

identifiable is the verdict for the natural direct and indirect effects: True when a covariate set covariates makes the mediation formula valid, False when none exists (the cross-world / post-treatment-confounder failure). nde_estimand and nie_estimand render the mediation formula with that set; they are empty strings when not identifiable.

The controlled direct effect is tracked separately, since it is a plain interventional contrast (no cross-world assumption): cde_identifiable is its verdict and cde_estimand its formula, decided through the general ID algorithm on P(Y | do(X), do(M)). CDE can be identified when NDE/NIE are not, and vice versa.

assumptions states, in words, the (untestable) claims the decomposition rests on, the same honesty the :class:~ergodic.inference.Design records carry for the quasi-experimental designs.

Source code in src/ergodic/identification/_mediation.py

@dataclass(frozen=True)
class MediationID:
    """The result of trying to identify a mediation decomposition from a graph.

    ``identifiable`` is the verdict for the natural direct and indirect effects:
    ``True`` when a covariate set ``covariates`` makes the mediation formula
    valid, ``False`` when none exists (the cross-world / post-treatment-confounder
    failure). ``nde_estimand`` and ``nie_estimand`` render the mediation formula
    with that set; they are empty strings when not identifiable.

    The controlled direct effect is tracked separately, since it is a plain
    interventional contrast (no cross-world assumption): ``cde_identifiable`` is
    its verdict and ``cde_estimand`` its formula, decided through the general ID
    algorithm on ``P(Y | do(X), do(M))``. CDE can be identified when NDE/NIE are
    not, and vice versa.

    ``assumptions`` states, in words, the (untestable) claims the decomposition
    rests on, the same honesty the :class:`~ergodic.inference.Design` records carry
    for the quasi-experimental designs.
    """

    treatment: frozenset[str]
    outcome: frozenset[str]
    mediators: frozenset[str]
    identifiable: bool
    nde_estimand: str
    nie_estimand: str
    cde_identifiable: bool
    cde_estimand: str
    covariates: frozenset[str] | None
    assumptions: tuple[str, ...]

    def __bool__(self) -> bool:
        return self.identifiable

    def __repr__(self) -> str:
        arrow = (
            f"{sorted(self.treatment)} -> [{', '.join(sorted(self.mediators))}] "
            f"-> {sorted(self.outcome)}"
        )
        if not self.identifiable:
            return f"MediationID(NDE/NIE not identifiable, {arrow}, cde={self.cde_identifiable})"
        return (
            f"MediationID(identifiable, {arrow}, "
            f"W={{{_join(self.covariates or frozenset())}}}, cde={self.cde_identifiable})"
        )

Prob dataclass

Bases: Functional

An atomic term P(head | given) (a marginal when given is empty).

Source code in src/ergodic/identification/_functional.py

@dataclass(frozen=True)
class Prob(Functional):
    """An atomic term ``P(head | given)`` (a marginal when ``given`` is empty)."""

    head: frozenset[str]
    given: frozenset[str] = frozenset()

    def render(self) -> str:
        if self.given:
            return f"P({_join(self.head)} | {_join(self.given)})"
        return f"P({_join(self.head)})"

    def free_vars(self) -> frozenset[str]:
        return self.head | self.given

    def simplify(self) -> Functional:
        return self

    def _is_one(self) -> bool:
        """``P(s | s)`` and ``P()`` are identically one and drop from a product."""
        return not self.head or self.head <= self.given

Product dataclass

Bases: Functional

A product of factors.

Source code in src/ergodic/identification/_functional.py

@dataclass(frozen=True)
class Product(Functional):
    """A product of factors."""

    factors: tuple[Functional, ...]

    def render(self) -> str:
        if not self.factors:
            return "1"
        return " ".join(f.render() for f in self.factors)

    def free_vars(self) -> frozenset[str]:
        out: frozenset[str] = frozenset()
        for f in self.factors:
            out |= f.free_vars()
        return out

    def simplify(self) -> Functional:
        flat: list[Functional] = []
        for f in self.factors:
            s = f.simplify()
            if isinstance(s, Product):
                flat.extend(s.factors)  # flatten nested products
            elif isinstance(s, Prob) and s._is_one():
                continue  # drop an identically-one factor
            else:
                flat.append(s)
        if len(flat) == 1:
            return flat[0]
        return Product(tuple(flat))

Ratio dataclass

Bases: Functional

A ratio num / den (the IDC normalization P_x(y, z) / P_x(z)).

Source code in src/ergodic/identification/_functional.py

@dataclass(frozen=True)
class Ratio(Functional):
    """A ratio ``num / den`` (the IDC normalization ``P_x(y, z) / P_x(z)``)."""

    num: Functional
    den: Functional

    def render(self) -> str:
        return f"[{self.num.render()}] / [{self.den.render()}]"

    def free_vars(self) -> frozenset[str]:
        return self.num.free_vars() | self.den.free_vars()

    def simplify(self) -> Functional:
        return Ratio(self.num.simplify(), self.den.simplify())

adjustment_set(graph, treatment, outcome)

Return the canonical valid adjustment set, or None if none exists.

The canonical set is the ancestors of treatment and outcome (possible ancestors, on a class face), minus the forbidden nodes and the latent nodes. If any valid adjustment set over the observed nodes exists, this one is valid; on a class, validity means valid in every member.

Source code in src/ergodic/identification/_criteria.py

def adjustment_set(
    graph: MixedGraph, treatment: str | Iterable[str], outcome: str | Iterable[str]
) -> frozenset[str] | None:
    """Return the canonical valid adjustment set, or ``None`` if none exists.

    The canonical set is the ancestors of treatment and outcome (possible
    ancestors, on a class face), minus the forbidden nodes and the latent
    nodes. If any valid adjustment set over the observed nodes exists, this
    one is valid; on a class, validity means valid in every member.
    """
    x = _as_set(treatment)
    y = _as_set(outcome)
    _validate(graph, x, y)
    if _is_class_kind(graph):
        candidate = possible_ancestors(graph, x | y) - _possible_forbidden(graph, x, y) - x - y
    else:
        candidate = (
            (_ancestors(graph, x | y) | x | y) - _forbidden(graph, x, y) - x - y - graph.latent
        )
    if is_valid_adjustment_set(graph, x, y, candidate):
        return frozenset(candidate)
    return None

frontdoor_set(graph, treatment, outcome)

Return a front-door mediator set, or None if none is found.

Tries the full set of observed nodes between treatment and outcome first, then single mediators. This is a search over likely candidates, not every subset.

Source code in src/ergodic/identification/_criteria.py

def frontdoor_set(
    graph: MixedGraph, treatment: str | Iterable[str], outcome: str | Iterable[str]
) -> frozenset[str] | None:
    """Return a front-door mediator set, or ``None`` if none is found.

    Tries the full set of observed nodes between treatment and outcome first,
    then single mediators. This is a search over likely candidates, not every
    subset.
    """
    x = _as_set(treatment)
    y = _as_set(outcome)
    _validate(graph, x, y)
    _check_single_graph(graph, "the front-door criterion")
    mediators = _causal_nodes(graph, x, y) - y - x - graph.latent
    if mediators and is_frontdoor_set(graph, x, y, mediators):
        return frozenset(mediators)
    for single in sorted(mediators):
        if is_frontdoor_set(graph, x, y, {single}):
            return frozenset({single})
    return None

ida(graph, treatment, outcome, *, max_members=_DEFAULT_MAX_MEMBERS)

Identify treatment -> outcome across every member of an equivalence class.

On a CPDAG, enumerate the DAG members (:func:dags_of_cpdag) and run :func:identify_effect per member, returning the multiset as an :class:IDA. On a PAG or MAG, return the class-level generalized-adjustment answer where it applies (a single member-invariant result), and raise when the case needs general per-member ID (deferred: it needs a DAG realization of each MAG member). A single graph (DAG or ADMG) raises with a pointer to :func:identify_effect.

Source code in src/ergodic/identification/_ida.py

def ida(
    graph: MixedGraph,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    max_members: int = _DEFAULT_MAX_MEMBERS,
) -> IDA:
    """Identify ``treatment -> outcome`` across every member of an equivalence class.

    On a CPDAG, enumerate the DAG members (:func:`dags_of_cpdag`) and run
    :func:`identify_effect` per member, returning the multiset as an :class:`IDA`.
    On a PAG or MAG, return the class-level generalized-adjustment answer where it
    applies (a single member-invariant result), and raise when the case needs
    general per-member ID (deferred: it needs a DAG realization of each MAG
    member). A single graph (DAG or ADMG) raises with a pointer to
    :func:`identify_effect`.
    """
    x = _as_set(treatment)
    y = _as_set(outcome)
    _validate(graph, x, y)
    kind = graph.kind

    if kind is GraphKind.MIXED:
        raise ValueError(
            "ida needs a verified class face (cpdag, pag, or mag); re-face the "
            "graph with as_kind once it satisfies a face's rules"
        )
    if kind in (GraphKind.DAG, GraphKind.ADMG):
        raise ValueError(
            "a single graph has one identify_effect; ida is for an equivalence "
            "class (cpdag, pag, or mag). Call identify_effect(graph, treatment, "
            "outcome) instead"
        )
    if kind is GraphKind.CPDAG:
        members = tuple(
            identify_effect(member, x, y)
            for member in dags_of_cpdag(graph, max_members=max_members)
        )
        return IDA(x, y, kind, members)
    # PAG or MAG: the class-level adjustment criterion gives a member-invariant
    # answer where it applies; otherwise general per-member ID is needed, which
    # needs a DAG realization of each MAG member (not built).
    class_level = identify_effect(graph, x, y)
    if class_level.identifiable:
        return IDA(x, y, kind, (class_level,))
    raise NotImplementedError(
        f"general ID over {kind.value} members needs a DAG realization of each "
        "MAG member (not built); the class-level adjustment criterion did not "
        f"identify {sorted(x)} -> {sorted(y)} here. Use the per-member adjustment "
        "criterion (identify_effect on the class) or enumerate the class "
        "(mags_of_pag) and reason per member"
    )

identify_distribution(graph, outcome, do, given=frozenset())

Run the general do-calculus ID (or IDC) algorithm on a single graph.

Decides whether P(outcome | do(do), given) is identifiable from the observational distribution and, when it is, returns the estimand as a typed functional inside an :class:~ergodic.identification._id.IDResult. With given empty this is ID; with given nonempty it is the conditional algorithm IDC. The graph may be a DAG or an ADMG, and may carry latent= designations or bows (it is latent-projected to an ADMG first).

A CPDAG, MAG, or PAG raises: ID is defined on a single graph, and the class is covered by :func:~ergodic.identification.ida, which enumerates members. A mixed-face graph raises, like the rest of identification.

Source code in src/ergodic/identification/_distribution.py

def identify_distribution(
    graph: MixedGraph,
    outcome: str | Iterable[str],
    do: str | Iterable[str],
    given: str | Iterable[str] = frozenset(),
) -> IDResult:
    """Run the general do-calculus ID (or IDC) algorithm on a single graph.

    Decides whether ``P(outcome | do(do), given)`` is identifiable from the
    observational distribution and, when it is, returns the estimand as a typed
    functional inside an :class:`~ergodic.identification._id.IDResult`. With
    ``given`` empty this is ID; with ``given`` nonempty it is the conditional
    algorithm IDC. The graph may be a DAG or an ADMG, and may carry ``latent=``
    designations or bows (it is latent-projected to an ADMG first).

    A CPDAG, MAG, or PAG raises: ID is defined on a single graph, and the class
    is covered by :func:`~ergodic.identification.ida`, which enumerates members.
    A ``mixed``-face graph raises, like the rest of identification.
    """
    y = _as_set(outcome)
    x = _as_set(do)
    z = _as_set(given)
    # An empty do-set is allowed here: P(y | do(), z) is the observational query
    # P(y | z), which the ID engine handles (step 1 returns the marginal).
    _validate(graph, x, y, z, allow_empty_treatment=True)
    if x & z:
        raise ValueError("the do-set and the conditioning set must be disjoint")
    if y & z:
        raise ValueError("the outcome and the conditioning set must be disjoint")
    if _is_class_kind(graph):
        raise NotImplementedError(
            f"identify_distribution runs the ID algorithm on a single graph, not on "
            f"an equivalence class ({graph.kind.value}); use ida() to enumerate the "
            "class and identify per member, or extract a member first"
        )
    return _identify_admg(graph, y, x, z)

identify_effect(graph, treatment, outcome)

Try adjustment, then the front-door, then instruments, and report the result.

Before any strategy runs, the trivial case is settled: when no directed path runs from the treatment to the outcome, the intervention cannot move the outcome, so P(y | do(x)) = P(y) and the strategy is "null".

On a class face (CPDAG, MAG, PAG) only the null and adjustment strategies apply, and a success holds in every graph the class represents: the null when no member has a causal path, adjustment through the generalized criterion. The front-door and instruments need a single graph, so a class that adjustment cannot identify reports not-identifiable and says whether amenability failed or no uniform set exists.

Source code in src/ergodic/identification/_criteria.py

def identify_effect(
    graph: MixedGraph, treatment: str | Iterable[str], outcome: str | Iterable[str]
) -> Identification:
    """Try adjustment, then the front-door, then instruments, and report the result.

    Before any strategy runs, the trivial case is settled: when no directed
    path runs from the treatment to the outcome, the intervention cannot move
    the outcome, so P(y | do(x)) = P(y) and the strategy is ``"null"``.

    On a class face (CPDAG, MAG, PAG) only the null and adjustment strategies
    apply, and a success holds in every graph the class represents: the null
    when no member has a causal path, adjustment through the generalized
    criterion. The front-door and instruments need a single graph, so a class
    that adjustment cannot identify reports not-identifiable and says whether
    amenability failed or no uniform set exists.
    """
    x = _as_set(treatment)
    y = _as_set(outcome)
    _validate(graph, x, y)
    is_class = _is_class_kind(graph)

    if is_class:
        no_path = not (possible_descendants(graph, x) & y)
        reason = "no possibly causal path from treatment to outcome in any member of the class"
    else:
        no_path = not (_descendants(graph, x) & y)
        reason = "no directed path from treatment to outcome"
    if no_path:
        return Identification(
            x,
            y,
            True,
            "null",
            estimand=f"P({_join(y)} | do({_join(x)})) = P({_join(y)}); {reason}",
        )

    via_adjustment = adjustment_set(graph, x, y)
    if via_adjustment is not None:
        return Identification(
            x,
            y,
            True,
            "adjustment",
            adjustment_set=via_adjustment,
            estimand=_adjustment_formula(x, y, via_adjustment),
        )

    if is_class:
        violations = _amenability_violations(graph, x, _possible_causal_nodes(graph, x, y))
        if violations:
            detail = (
                f"not amenable: the possibly causal path leaving the treatment "
                f"through {violations[0]!r} does not start with a visible directed "
                "edge, so members of the class disagree about it"
            )
        else:
            detail = "no adjustment set is valid in every member of the class"
        return Identification(
            x,
            y,
            False,
            "none",
            estimand=(
                f"{detail}; the front-door and instruments need a single graph "
                "(enumerate members with dags_of_cpdag or mags_of_pag)"
            ),
        )

    via_frontdoor = frontdoor_set(graph, x, y)
    if via_frontdoor is not None:
        return Identification(
            x,
            y,
            True,
            "frontdoor",
            mediators=via_frontdoor,
            estimand=_frontdoor_formula(x, y, via_frontdoor),
        )

    if len(x) == 1 and len(y) == 1:
        found = instruments(graph, next(iter(x)), next(iter(y)))
        if found:
            names = ", ".join(i for i, _ in found)
            return Identification(
                x,
                y,
                True,
                "iv",
                instruments=tuple(found),
                estimand=(
                    f"identified through instrument(s) {names} under a linear "
                    "(parametric) model; needs an IV estimator"
                ),
            )

    # The named ladder missed. On a single graph the general do-calculus ID
    # algorithm decides every remaining case: it identifies (returning the
    # functional, or a recognized named reduction) or exhibits the hedge that
    # proves non-identifiability. The import is function-local to avoid a load
    # cycle (the ID engine and recognition import this module's criteria).
    from ._distribution import id_fallthrough

    return id_fallthrough(graph, x, y)

identify_mediation(graph, treatment, outcome, mediators)

Identify the mediation decomposition of treatment -> outcome through mediators.

Searches for a covariate set W of non-descendants of the treatment under which the natural direct and indirect effects are identified by the mediation formula (Pearl, 2001): W is a valid adjustment set for X -> Y and for X -> M, and W u {X} is a valid adjustment set for M -> Y. When such a W exists, NDE and NIE are identifiable and their estimands render in the mediation-formula style; when none exists, they are not identifiable, the canonical case being a mediator-outcome confounder that is itself affected by the treatment (a cross-world obstruction).

The controlled direct effect is identified separately, through the general ID algorithm on P(Y | do(X), do(M)): a plain interventional contrast that needs no cross-world assumption, and may identify when NDE/NIE do not.

Single graphs only (a DAG or ADMG). A CPDAG, MAG, or PAG raises (extract a member or enumerate the class).

Source code in src/ergodic/identification/_mediation.py

def identify_mediation(
    graph: MixedGraph,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    mediators: str | Iterable[str],
) -> MediationID:
    """Identify the mediation decomposition of ``treatment -> outcome`` through ``mediators``.

    Searches for a covariate set ``W`` of non-descendants of the treatment under
    which the natural direct and indirect effects are identified by the mediation
    formula (Pearl, 2001): ``W`` is a valid adjustment set for ``X -> Y`` and for
    ``X -> M``, and ``W u {X}`` is a valid adjustment set for ``M -> Y``. When such
    a ``W`` exists, NDE and NIE are identifiable and their estimands render in the
    mediation-formula style; when none exists, they are not identifiable, the
    canonical case being a mediator-outcome confounder that is itself affected by
    the treatment (a cross-world obstruction).

    The controlled direct effect is identified separately, through the general ID
    algorithm on ``P(Y | do(X), do(M))``: a plain interventional contrast that
    needs no cross-world assumption, and may identify when NDE/NIE do not.

    Single graphs only (a DAG or ADMG). A CPDAG, MAG, or PAG raises (extract a
    member or enumerate the class).
    """
    x = _as_set(treatment)
    y = _as_set(outcome)
    m = _as_set(mediators)
    _validate(graph, x, y, m)
    _check_single_graph(graph, "mediation")
    _check_mediators(graph, x, y, m)

    cde_identifiable, cde_estimand = _cde_through_id(graph, x, y, m)

    w = _find_w(graph, x, y, m)
    if w is None:
        assumptions: tuple[str, ...] = (
            "the natural direct and indirect effects are not identifiable: no set "
            "of non-descendants of the treatment satisfies the three adjustment "
            "conditions. This is the cross-world obstruction, typically a "
            "mediator-outcome confounder that is itself affected by the treatment "
            "(a recanting witness); the controlled direct effect may still be "
            "identifiable. Consider a path-specific effect or report the CDE.",
        )
        return MediationID(
            treatment=x,
            outcome=y,
            mediators=m,
            identifiable=False,
            nde_estimand="",
            nie_estimand="",
            cde_identifiable=cde_identifiable,
            cde_estimand=cde_estimand,
            covariates=None,
            assumptions=assumptions,
        )

    if w:
        seq_ignorability = (
            f"sequential ignorability given W = {{{_join(w)}}}: conditional on these "
            "covariates, the treatment is as-if randomized for both the mediator and "
            "the outcome, and the mediator is as-if randomized for the outcome."
        )
    else:
        seq_ignorability = (
            "sequential ignorability with no covariates needed: the treatment is "
            "as-if randomized for both the mediator and the outcome, and the "
            "mediator is as-if randomized for the outcome given the treatment."
        )
    assumptions = (
        seq_ignorability,
        "the cross-world independence Y(x, m) is independent of M(x') given W "
        "(equivalently, no mediator-outcome confounder is affected by the "
        "treatment). This is graphically certified here but is untestable from "
        "data: no observational or experimental design can refute it.",
    )
    return MediationID(
        treatment=x,
        outcome=y,
        mediators=m,
        identifiable=True,
        nde_estimand=_nde_formula(x, y, m, w),
        nie_estimand=_nie_formula(x, y, m, w),
        cde_identifiable=cde_identifiable,
        cde_estimand=cde_estimand,
        covariates=w,
        assumptions=assumptions,
    )

instruments(graph, treatment, outcome)

Find conditional instruments for treatment -> outcome.

Returns (instrument, conditioning set) pairs. Candidates and conditioning nodes are observed non-descendants of the treatment and the outcome; the search tries the empty set and single nodes as conditioning sets.

Source code in src/ergodic/identification/_criteria.py

def instruments(
    graph: MixedGraph, treatment: str, outcome: str
) -> list[tuple[str, frozenset[str]]]:
    """Find conditional instruments for ``treatment -> outcome``.

    Returns ``(instrument, conditioning set)`` pairs. Candidates and conditioning
    nodes are observed non-descendants of the treatment and the outcome; the
    search tries the empty set and single nodes as conditioning sets.
    """
    x = frozenset({treatment})
    y = frozenset({outcome})
    _validate(graph, x, y)
    _check_single_graph(graph, "instrument identification")
    usable = graph.nodes - graph.latent - _descendants(graph, x) - x - y
    non_descendant = sorted(usable - _descendants(graph, y))
    found: list[tuple[str, frozenset[str]]] = []
    for candidate in sorted(usable):
        condition_sets: list[frozenset[str]] = [frozenset()]
        condition_sets += [frozenset({w}) for w in non_descendant if w != candidate]
        for w in condition_sets:
            if is_instrument(graph, candidate, treatment, outcome, w):
                found.append((candidate, frozenset(w)))
                break
    return found

is_frontdoor_set(graph, treatment, outcome, mediators)

Return True if mediators satisfies the front-door criterion.

A mediator set holding a latent node is never valid, since the front-door formula conditions on the mediators.

Source code in src/ergodic/identification/_criteria.py

def is_frontdoor_set(
    graph: MixedGraph,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    mediators: str | Iterable[str],
) -> bool:
    """Return ``True`` if ``mediators`` satisfies the front-door criterion.

    A mediator set holding a latent node is never valid, since the front-door
    formula conditions on the mediators.
    """
    x = _as_set(treatment)
    y = _as_set(outcome)
    m = _as_set(mediators)
    _validate(graph, x, y, m)
    _check_single_graph(graph, "the front-door criterion")
    if not m or (x & m) or (y & m) or (m & graph.latent):
        return False
    # (1) every directed path from x to y passes through m
    without_m = graph.induced_subgraph(graph.nodes - m)
    if _descendants(without_m, x) & y:
        return False
    # (2) no open back-door path from x to m
    if not _cut_outgoing(graph, x).m_separated(x, m):
        return False
    # (3) treatment blocks every back-door path from m to y
    return _cut_outgoing(graph, m).m_separated(m, y, x)

is_instrument(graph, instrument, treatment, outcome, given=frozenset())

Return True if instrument is a conditional instrument given given.

The instrument and the conditioning set are observed non-descendants of the treatment and the outcome, the instrument is associated with the treatment, and the instrument reaches the outcome only through the treatment (it separates from the outcome once the treatment's outgoing edges are cut).

Descendants of the treatment are rejected outright: cutting the treatment's outgoing edges would sever the very edge that makes them endogenous, so the exclusion check alone would pass them.

Source code in src/ergodic/identification/_criteria.py

def is_instrument(
    graph: MixedGraph,
    instrument: str,
    treatment: str,
    outcome: str,
    given: str | Iterable[str] = frozenset(),
) -> bool:
    """Return ``True`` if ``instrument`` is a conditional instrument given ``given``.

    The instrument and the conditioning set are observed non-descendants of the
    treatment and the outcome, the instrument is associated with the treatment,
    and the instrument reaches the outcome only through the treatment (it
    separates from the outcome once the treatment's outgoing edges are cut).

    Descendants of the treatment are rejected outright: cutting the treatment's
    outgoing edges would sever the very edge that makes them endogenous, so the
    exclusion check alone would pass them.
    """
    i = frozenset({instrument})
    x = frozenset({treatment})
    y = frozenset({outcome})
    w = _as_set(given)
    _validate(graph, x, y, i | w)
    _check_single_graph(graph, "instrument identification")
    if (i & x) or (i & y) or (w & (i | x | y)):
        return False
    if (i | w) & graph.latent:
        return False
    if (i | w) & _descendants(graph, x):
        return False
    if w & (_descendants(graph, y) | y):
        return False
    if graph.m_separated(i, x, w):  # relevance: must be associated
        return False
    return _cut_outgoing(graph, x).m_separated(i, y, w)  # exclusion

is_valid_adjustment_set(graph, treatment, outcome, adjustment)

Return True if adjustment is a valid adjustment set for the effect.

On a DAG or ADMG, the complete adjustment criterion: the set avoids the forbidden nodes (descendants of the causal-path nodes) and blocks every proper non-causal path from treatment to outcome. A set holding a latent node is never valid, since a latent node cannot be conditioned on.

On a CPDAG, MAG, or PAG, the generalized criterion (Perkovic et al., 2018): the class must be amenable, the set avoids the possible descendants of the possibly-causal-path nodes, and it blocks every proper definite-status non-causal path. True then means the set is valid in every graph the class represents.

Source code in src/ergodic/identification/_criteria.py

def is_valid_adjustment_set(
    graph: MixedGraph,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    adjustment: str | Iterable[str],
) -> bool:
    """Return ``True`` if ``adjustment`` is a valid adjustment set for the effect.

    On a DAG or ADMG, the complete adjustment criterion: the set avoids the
    forbidden nodes (descendants of the causal-path nodes) and blocks every
    proper non-causal path from treatment to outcome. A set holding a latent
    node is never valid, since a latent node cannot be conditioned on.

    On a CPDAG, MAG, or PAG, the generalized criterion (Perkovic et al.,
    2018): the class must be amenable, the set avoids the possible
    descendants of the possibly-causal-path nodes, and it blocks every
    proper definite-status non-causal path. ``True`` then means the set is
    valid in *every* graph the class represents.
    """
    x = _as_set(treatment)
    y = _as_set(outcome)
    z = _as_set(adjustment)
    _validate(graph, x, y, z)
    if _is_class_kind(graph):
        if (x & z) or (y & z):
            return False
        cn = _possible_causal_nodes(graph, x, y)
        if _amenability_violations(graph, x, cn):
            return False
        if z & _possible_forbidden(graph, x, y):
            return False
        return _proper_backdoor_graph(graph, x, cn).m_separated(x, y, z)
    if (x & z) or (y & z) or (z & graph.latent):
        return False
    cn = _causal_nodes(graph, x, y)
    if z & _forbidden(graph, x, y):
        return False
    proper = _proper_backdoor_graph(graph, x, cn)
    return proper.m_separated(x, y, z)

minimal_adjustment_set(graph, treatment, outcome)

Return one minimal valid adjustment set, or None if none exists.

Source code in src/ergodic/identification/_criteria.py

def minimal_adjustment_set(
    graph: MixedGraph, treatment: str | Iterable[str], outcome: str | Iterable[str]
) -> frozenset[str] | None:
    """Return one minimal valid adjustment set, or ``None`` if none exists."""
    _validate(graph, _as_set(treatment), _as_set(outcome))
    start = adjustment_set(graph, treatment, outcome)
    if start is None:
        return None
    keep = set(start)
    changed = True
    while changed:
        changed = False
        for node in sorted(keep):
            if is_valid_adjustment_set(graph, treatment, outcome, keep - {node}):
                keep.discard(node)
                changed = True
    return frozenset(keep)

Causal discovery

Structure learning: the algorithms, the CI tests and scores they consume, and the discover facade.

ergodic.discovery

Causal discovery: learn causal structure from data.

Two ingredients drive the searches. A conditional-independence test (:class:CITest) is a class constructed with data, which makes it the stateful oracle the constraint algorithms query: :class:Linear (partial correlation), :class:Discrete (G squared, natively categorical), :class:Mixed (degenerate Gaussian for mixed columns), :class:GCM (any regression :class:~ergodic.inference.Learner becomes a test), and :class:DSeparationOracle as ground truth. A decomposable :class:Score answers local(node, parents) for the score-based searches: :class:BIC (Gaussian), :class:DiscreteBIC, :class:BDeu, and :class:MixedBIC (degenerate Gaussian). Algorithms are thin consumers: :class:PC wraps a test and returns a CPDAG, :class:FCI wraps a test and returns a PAG (latent confounders allowed), :class:MMPC wraps a test and returns an undirected skeleton screen, :class:GES wraps a score and returns a CPDAG, :class:AStar wraps a score and returns the exact optimum, and :class:HillClimb is the fast greedy baseline. discover is the one-line facade, and :func:bootstrap reruns any method on resamples and returns edge frequencies (:class:BootstrapResult), the first soft output. Background knowledge (including a skeleton restriction through restrict_to) enters through :class:ergodic.DomainKnowledge.

Time series and panels get their own methods: :class:PCMCI reads a :class:~ergodic.data.TimeSeriesData or :class:~ergodic.data.PanelData through a lagged design matrix (:func:lagged_tabular) and returns a window graph over nodes like "X[t-1]" (:class:TimeSeriesDiscoveryResult), and :func:granger is the predictive baseline, a pair-level table on purpose rather than a graph.

AStar

Exact structure search: the highest-BIC DAG, found with A*.

Parameters:

Name	Type	Description	Default
score	Score | str	A :class:~ergodic.discovery.Score instance (bound or unbound) or a name: "auto" (pick by the data schema), "bic", "bic_discrete", "bdeu", "bic_mixed".	'auto'
max_parents	int | None	An in-degree cap; None enumerates every parent set.	None
node_limit	int	Refuse problems with more variables than this; exact search is exponential in the node count.	20

Source code in src/ergodic/discovery/algorithms/_astar.py

class AStar:
    """Exact structure search: the highest-BIC DAG, found with A*.

    Args:
        score: A :class:`~ergodic.discovery.Score` instance (bound or
            unbound) or a name: ``"auto"`` (pick by the data schema),
            ``"bic"``, ``"bic_discrete"``, ``"bdeu"``, ``"bic_mixed"``.
        max_parents: An in-degree cap; ``None`` enumerates every parent set.
        node_limit: Refuse problems with more variables than this; exact
            search is exponential in the node count.
    """

    ingredient: ClassVar[str] = "score"

    def __init__(
        self,
        score: Score | str = "auto",
        *,
        max_parents: int | None = None,
        node_limit: int = 20,
    ) -> None:
        if not isinstance(score, Score | str):
            raise TypeError(f"score must be a Score or a string, got {type(score).__name__}")
        if max_parents is not None and max_parents < 0:
            raise ValueError(f"max_parents must be non-negative, got {max_parents}")
        if node_limit < 1:
            raise ValueError(f"node_limit must be at least 1, got {node_limit}")
        self._score = score
        self._max_parents = max_parents
        self._node_limit = node_limit

    # -- the public entry -------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Search for a highest-scoring DAG and return it with its provenance."""
        score = resolve_score(self._score, data)
        variables = list(score.variables)
        if len(variables) > self._node_limit:
            raise ValueError(
                f"exact search over {len(variables)} variables exceeds "
                f"node_limit={self._node_limit}; cap the in-degree with max_parents, "
                "shrink the problem with DomainKnowledge.restrict_to, or raise "
                "node_limit if you accept the cost"
            )
        tables = {v: self._node_table(score, v, variables, knowledge) for v in variables}
        parents_of, expanded = self._search(variables, tables)
        edges = [f"{u} -> {v}" for v in variables for u in sorted(parents_of[v])]
        graph = dag(edges, nodes=variables)
        if knowledge is not None:
            assert knowledge.is_consistent(graph), (
                "the searched DAG violates the domain knowledge; constraints a DAG "
                "cannot express (confounded pairs) or required edges over variables "
                "missing from the data can cause this"
            )
        diagnostics = dict(score.diagnostics())
        diagnostics["expanded_states"] = float(expanded)
        return DiscoveryResult(
            graph=graph,
            method="astar",
            score=score.name,
            alpha=None,
            n=score.n,
            diagnostics=diagnostics,
        )

    # -- per-node tables ------------------------------------------------------------

    def _node_table(
        self,
        score: Score,
        node: str,
        variables: list[str],
        knowledge: DomainKnowledge | None,
    ) -> _NodeTable:
        """Score every admissible parent set of ``node`` and sweep the best-set DP."""
        others = [u for u in variables if u != node]
        required: frozenset[str] = frozenset()
        if knowledge is None:
            candidates = others
        else:
            candidates = [u for u in others if knowledge.permits_edge(u, node)]
            required = frozenset(u for u in others if knowledge.requires_edge(u, node))
            stranded = sorted(required - set(candidates))
            if stranded:
                raise ValueError(
                    f"required parent(s) {stranded} of {node!r} are not permitted "
                    "by the same knowledge"
                )
        if self._max_parents is not None and len(required) > self._max_parents:
            raise ValueError(
                f"{node!r} has {len(required)} required parents, more than "
                f"max_parents={self._max_parents}; raise the cap or relax the knowledge"
            )
        position = {name: i for i, name in enumerate(candidates)}
        scored = score.enumerate_locals(node, candidates, self._max_parents)
        size = 1 << len(candidates)
        exact = np.full(size, -math.inf, dtype=np.float64)
        for parent_set, value in scored.items():
            if required <= parent_set:
                mask = 0
                for parent in parent_set:
                    mask |= 1 << position[parent]
                exact[mask] = value
        # bp[S] = max(exact[S], max over single-element removals of bp[S \ {s}]),
        # swept in increasing mask order so every submask is already final.
        bp = exact.copy()
        choice = np.where(np.isfinite(exact), np.arange(size, dtype=np.int64), np.int64(-1))
        for mask in range(size):
            best = float(bp[mask])
            pick = int(choice[mask])
            remaining = mask
            while remaining:
                low = remaining & -remaining
                value = float(bp[mask ^ low])
                if value > best:
                    best = value
                    pick = int(choice[mask ^ low])
                remaining ^= low
            bp[mask] = best
            choice[mask] = pick
        return _NodeTable(candidates=tuple(candidates), bp=bp, choice=choice)

    # -- the search --------------------------------------------------------------

    def _search(
        self, variables: list[str], tables: dict[str, _NodeTable]
    ) -> tuple[dict[str, frozenset[str]], int]:
        """A* from the empty set to the full set; returns the parent map and a counter."""
        bit_of = {name: 1 << i for i, name in enumerate(variables)}
        projections = {
            v: tuple((bit_of[u], 1 << j) for j, u in enumerate(tables[v].candidates))
            for v in variables
        }

        def project(state: int, node: str) -> int:
            """Restrict a placed-set state to ``node``'s candidate pool bitmask."""
            out = 0
            for variable_bit, candidate_bit in projections[node]:
                if state & variable_bit:
                    out |= candidate_bit
            return out

        # The heuristic term per unplaced node: its best local over ALL candidates.
        # bp is monotone in the subset, so h never overestimates and is consistent.
        h_terms = {v: -float(tables[v].bp[-1]) for v in variables}
        full = (1 << len(variables)) - 1
        g_best: dict[int, float] = {0: 0.0}
        came: dict[int, tuple[int, str]] = {}
        closed: set[int] = set()
        # Entries are (f, state); equal f breaks ties on the state integer, so
        # the expansion order is deterministic.
        frontier: list[tuple[float, int]] = [(sum(h_terms.values()), 0)]
        expanded = 0
        while frontier:
            f_value, state = heapq.heappop(frontier)
            if state in closed:
                continue
            closed.add(state)
            expanded += 1
            if state == full:
                return self._reconstruct(state, came, tables, project), expanded
            g_state = g_best[state]
            h_state = f_value - g_state
            for node in variables:
                bit = bit_of[node]
                if state & bit:
                    continue
                local_best = float(tables[node].bp[project(state, node)])
                if math.isinf(local_best):
                    continue  # required parents not yet placed
                child = state | bit
                if child in closed:
                    continue
                g_child = g_state - local_best
                if g_child < g_best.get(child, math.inf):
                    g_best[child] = g_child
                    came[child] = (state, node)
                    f_child = g_child + h_state - h_terms[node]
                    heapq.heappush(frontier, (f_child, child))
        raise ValueError(
            "the constraints admit no DAG over these variables; "
            "check the knowledge for required-edge cycles"
        )

    @staticmethod
    def _reconstruct(
        goal: int,
        came: dict[int, tuple[int, str]],
        tables: dict[str, _NodeTable],
        project: _Projector,
    ) -> dict[str, frozenset[str]]:
        """Walk the order back from the goal and read each node's parents off its DP."""
        parents_of: dict[str, frozenset[str]] = {}
        state = goal
        while state:
            previous, placed = came[state]
            table = tables[placed]
            pick = int(table.choice[project(previous, placed)])
            if pick < 0:  # pragma: no cover - the search never relaxes through -inf
                raise RuntimeError(f"no admissible parent set for {placed!r}; this is a bug")
            parents_of[placed] = frozenset(
                name for j, name in enumerate(table.candidates) if pick >> j & 1
            )
            state = previous
        return parents_of

fit(data=None, knowledge=None)

Search for a highest-scoring DAG and return it with its provenance.

Source code in src/ergodic/discovery/algorithms/_astar.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Search for a highest-scoring DAG and return it with its provenance."""
    score = resolve_score(self._score, data)
    variables = list(score.variables)
    if len(variables) > self._node_limit:
        raise ValueError(
            f"exact search over {len(variables)} variables exceeds "
            f"node_limit={self._node_limit}; cap the in-degree with max_parents, "
            "shrink the problem with DomainKnowledge.restrict_to, or raise "
            "node_limit if you accept the cost"
        )
    tables = {v: self._node_table(score, v, variables, knowledge) for v in variables}
    parents_of, expanded = self._search(variables, tables)
    edges = [f"{u} -> {v}" for v in variables for u in sorted(parents_of[v])]
    graph = dag(edges, nodes=variables)
    if knowledge is not None:
        assert knowledge.is_consistent(graph), (
            "the searched DAG violates the domain knowledge; constraints a DAG "
            "cannot express (confounded pairs) or required edges over variables "
            "missing from the data can cause this"
        )
    diagnostics = dict(score.diagnostics())
    diagnostics["expanded_states"] = float(expanded)
    return DiscoveryResult(
        graph=graph,
        method="astar",
        score=score.name,
        alpha=None,
        n=score.n,
        diagnostics=diagnostics,
    )

BDeu

Bases: Score

The Bayesian Dirichlet equivalent uniform score, decomposable.

Every column must be discrete (categorical, ordinal, or binary): the score runs on integer codes, so binding to data with a continuous column raises TypeError. Score-equivalent for any fixed ess, but the value of ess shapes the result; see the module docstring.

Parameters:

Name	Type	Description	Default
data	TabularData | None	Bind at construction, or later through bind.	None
ess	float	The equivalent sample size, the total prior mass spread over each node's table. Must be positive. Larger values smooth the cell estimates and favor denser graphs; 10.0 is a common default, and a sweep is worth it when the result matters.	10.0

Source code in src/ergodic/discovery/scores/_bdeu.py

class BDeu(Score):
    """The Bayesian Dirichlet equivalent uniform score, decomposable.

    Every column must be discrete (categorical, ordinal, or binary): the
    score runs on integer codes, so binding to data with a continuous column
    raises ``TypeError``. Score-equivalent for any fixed ``ess``, but the
    value of ``ess`` shapes the result; see the module docstring.

    Args:
        data: Bind at construction, or later through ``bind``.
        ess: The equivalent sample size, the total prior mass spread over
            each node's table. Must be positive. Larger values smooth the
            cell estimates and favor denser graphs; ``10.0`` is a common
            default, and a sweep is worth it when the result matters.
    """

    name: ClassVar[str] = "bdeu"
    needs: ClassVar[str] = "codes"
    decomposable: ClassVar[bool] = True
    score_equivalent: ClassVar[bool] = True

    def __init__(self, data: TabularData | None = None, *, ess: float = 10.0) -> None:
        if ess <= 0.0:
            raise ValueError(f"ess must be positive, got {ess}")
        self._ess = float(ess)
        super().__init__(data)

    @property
    def ess(self) -> float:
        """The equivalent sample size of the Dirichlet prior."""
        return self._ess

    def local_from_codes(
        self,
        y: npt.NDArray[np.int64],
        y_card: int,
        zs: tuple[npt.NDArray[np.int64], ...],
        z_cards: tuple[int, ...],
    ) -> float:
        """The Dirichlet log marginal likelihood from integer codes."""
        n = int(y.shape[0])
        if n == 0:
            raise ValueError("BDeu needs at least one row")
        r = int(y_card)
        q = 1
        for card in z_cards:
            q *= int(card)
        flat = np.ravel_multi_index((*zs, y), (*z_cards, r))
        counts = np.bincount(flat, minlength=q * r).reshape(q, r)
        totals = counts.sum(axis=1)
        a_j = self._ess / q
        a_jk = self._ess / (q * r)
        lgamma_a_j = math.lgamma(a_j)
        lgamma_a_jk = math.lgamma(a_jk)
        local = 0.0
        for j in np.flatnonzero(totals):
            local += lgamma_a_j - math.lgamma(a_j + float(totals[j]))
            for count in counts[j]:
                if count:
                    local += math.lgamma(a_jk + float(count)) - lgamma_a_jk
        return local

ess property

The equivalent sample size of the Dirichlet prior.

local_from_codes(y, y_card, zs, z_cards)

The Dirichlet log marginal likelihood from integer codes.

Source code in src/ergodic/discovery/scores/_bdeu.py

def local_from_codes(
    self,
    y: npt.NDArray[np.int64],
    y_card: int,
    zs: tuple[npt.NDArray[np.int64], ...],
    z_cards: tuple[int, ...],
) -> float:
    """The Dirichlet log marginal likelihood from integer codes."""
    n = int(y.shape[0])
    if n == 0:
        raise ValueError("BDeu needs at least one row")
    r = int(y_card)
    q = 1
    for card in z_cards:
        q *= int(card)
    flat = np.ravel_multi_index((*zs, y), (*z_cards, r))
    counts = np.bincount(flat, minlength=q * r).reshape(q, r)
    totals = counts.sum(axis=1)
    a_j = self._ess / q
    a_jk = self._ess / (q * r)
    lgamma_a_j = math.lgamma(a_j)
    lgamma_a_jk = math.lgamma(a_jk)
    local = 0.0
    for j in np.flatnonzero(totals):
        local += lgamma_a_j - math.lgamma(a_j + float(totals[j]))
        for count in counts[j]:
            if count:
                local += math.lgamma(a_jk + float(count)) - lgamma_a_jk
    return local

BIC

Bases: Score

The Gaussian BIC, decomposable and score-equivalent.

Nodes must be single numeric columns (continuous, datetime, or binary); scoring a wider categorical block raises TypeError lazily, at the first query against that node. Categorical parents are fine: their one-hot columns enter the regression and each column counts toward k.

Parameters:

Name	Type	Description	Default
data	TabularData | None	Bind at construction, or later through bind.	None
penalty	float	A multiplier on the complexity term ((k + 1)/2) log n. 1.0 is the textbook BIC; smaller values prune less.	1.0

Source code in src/ergodic/discovery/scores/_bic.py

class BIC(Score):
    """The Gaussian BIC, decomposable and score-equivalent.

    Nodes must be single numeric columns (continuous, datetime, or binary);
    scoring a wider categorical block raises ``TypeError`` lazily, at the
    first query against that node. Categorical *parents* are fine: their
    one-hot columns enter the regression and each column counts toward ``k``.

    Args:
        data: Bind at construction, or later through ``bind``.
        penalty: A multiplier on the complexity term ``((k + 1)/2) log n``.
            ``1.0`` is the textbook BIC; smaller values prune less.
    """

    name: ClassVar[str] = "bic"
    decomposable: ClassVar[bool] = True
    score_equivalent: ClassVar[bool] = True

    def __init__(self, data: TabularData | None = None, *, penalty: float = 1.0) -> None:
        if penalty <= 0.0:
            raise ValueError(f"penalty must be positive, got {penalty}")
        self._penalty = float(penalty)
        self._slices: dict[str, slice] = {}
        self._cov: npt.NDArray[np.float64] = np.empty((0, 0), dtype=np.float64)
        super().__init__(data)

    @property
    def penalty(self) -> float:
        """The multiplier on the complexity term."""
        return self._penalty

    # -- one-time sufficient statistics ---------------------------------------

    def _prepare(self) -> None:
        if self._n == 0:
            raise ValueError("BIC needs at least one complete row; the data has none")
        slices: dict[str, slice] = {}
        parts: list[npt.NDArray[np.float64]] = []
        start = 0
        for name in self.variables:
            block = self._blocks[name]
            stop = start + block.shape[1]
            slices[name] = slice(start, stop)
            parts.append(block)
            start = stop
        design = np.hstack(parts) if parts else np.empty((self._n, 0), dtype=np.float64)
        centered = design - design.mean(axis=0, keepdims=True)
        self._slices = slices
        self._cov = np.asarray(centered.T @ centered, dtype=np.float64) / float(self._n)

    # -- the covariance path ---------------------------------------------------

    def _local_value(self, node: str, parents: frozenset[str]) -> float:
        """The local by Schur complement on the precomputed covariance."""
        node_slice = self._slices[node]
        width = node_slice.stop - node_slice.start
        if width != 1:
            raise TypeError(
                f"BIC scores single-column nodes; {node!r} encodes to {width} columns. "
                "The mixed score lands in D2."
            )
        target = node_slice.start
        columns = [
            column
            for parent in sorted(parents)
            for column in range(self._slices[parent].start, self._slices[parent].stop)
        ]
        syy = float(self._cov[target, target])
        if columns:
            spp = self._cov[np.ix_(columns, columns)]
            spy = self._cov[columns, target]
            sigma2 = syy - float(spy @ np.linalg.pinv(spp, hermitian=True) @ spy)
        else:
            sigma2 = syy
        return self._bic_value(sigma2, k=len(columns), n=self._n)

    # -- the array path (user contract, and the cross-check) -------------------

    def local_from_arrays(self, y: npt.NDArray[np.float64], X: npt.NDArray[np.float64]) -> float:
        """The same local by direct least squares on ``(n, 1)`` ``y`` and ``(n, k)`` ``X``.

        Agrees with the covariance path to numerical precision; the tests hold
        the two to ``1e-8``.
        """
        if y.ndim != 2 or y.shape[1] != 1:
            raise TypeError(
                "BIC scores single-column nodes; y must have shape (n, 1). "
                "The mixed score lands in D2."
            )
        n = int(y.shape[0])
        if n == 0:
            raise ValueError("BIC needs at least one row")
        centered_y = y - y.mean(axis=0, keepdims=True)
        if X.shape[1]:
            centered_x = X - X.mean(axis=0, keepdims=True)
            beta, *_ = np.linalg.lstsq(centered_x, centered_y, rcond=None)
            residual = centered_y - centered_x @ beta
        else:
            residual = centered_y
        sigma2 = float((residual * residual).sum()) / n
        return self._bic_value(sigma2, k=int(X.shape[1]), n=n)

    # -- the shared tail --------------------------------------------------------

    def _bic_value(self, sigma2: float, *, k: int, n: int) -> float:
        bounded = max(sigma2, _VARIANCE_FLOOR)
        size = float(n)
        fit = -(size / 2.0) * math.log(bounded)
        complexity = self._penalty * ((k + 1) / 2.0) * math.log(size)
        return fit - complexity

penalty property

The multiplier on the complexity term.

local_from_arrays(y, X)

The same local by direct least squares on (n, 1) y and (n, k) X.

Agrees with the covariance path to numerical precision; the tests hold the two to 1e-8.

Source code in src/ergodic/discovery/scores/_bic.py

def local_from_arrays(self, y: npt.NDArray[np.float64], X: npt.NDArray[np.float64]) -> float:
    """The same local by direct least squares on ``(n, 1)`` ``y`` and ``(n, k)`` ``X``.

    Agrees with the covariance path to numerical precision; the tests hold
    the two to ``1e-8``.
    """
    if y.ndim != 2 or y.shape[1] != 1:
        raise TypeError(
            "BIC scores single-column nodes; y must have shape (n, 1). "
            "The mixed score lands in D2."
        )
    n = int(y.shape[0])
    if n == 0:
        raise ValueError("BIC needs at least one row")
    centered_y = y - y.mean(axis=0, keepdims=True)
    if X.shape[1]:
        centered_x = X - X.mean(axis=0, keepdims=True)
        beta, *_ = np.linalg.lstsq(centered_x, centered_y, rcond=None)
        residual = centered_y - centered_x @ beta
    else:
        residual = centered_y
    sigma2 = float((residual * residual).sum()) / n
    return self._bic_value(sigma2, k=int(X.shape[1]), n=n)

BOSS

Best order score search over a decomposable score.

Parameters:

Name	Type	Description	Default
score	Score | str	A :class:~ergodic.discovery.Score instance (bound or unbound) or a name: "auto" (pick by the data schema), "bic", "bic_discrete", "bdeu", "bic_mixed".	'auto'
use_bes	bool	Finish with the backward equivalence step (GES's delete phase) on the CPDAG of the projected DAG. The paper's correctness proposition needs it, so it defaults on; the paper's experiments ran without it.	True
restarts	int	Extra runs from shuffled starting permutations, on top of the sorted start. The best final score across starts wins, with the earlier start keeping ties.	0
seed	int | None	The seed for the restart shuffles. None means a fixed default of 0, so a run is reproducible without choosing one; unused when restarts is 0.	None

fit(data, knowledge) returns a :class:DiscoveryResult whose graph carries the CPDAG face when the final PDAG is a valid CPDAG and the mixed face otherwise (forced knowledge orientations can leave the class; diagnostics["cpdag_valid"] says which happened).

Source code in src/ergodic/discovery/algorithms/_boss.py

class BOSS:
    """Best order score search over a decomposable score.

    Args:
        score: A :class:`~ergodic.discovery.Score` instance (bound or
            unbound) or a name: ``"auto"`` (pick by the data schema),
            ``"bic"``, ``"bic_discrete"``, ``"bdeu"``, ``"bic_mixed"``.
        use_bes: Finish with the backward equivalence step (GES's delete
            phase) on the CPDAG of the projected DAG. The paper's
            correctness proposition needs it, so it defaults on; the
            paper's experiments ran without it.
        restarts: Extra runs from shuffled starting permutations, on top
            of the sorted start. The best final score across starts wins,
            with the earlier start keeping ties.
        seed: The seed for the restart shuffles. ``None`` means a fixed
            default of 0, so a run is reproducible without choosing one;
            unused when ``restarts`` is 0.

    ``fit(data, knowledge)`` returns a :class:`DiscoveryResult` whose graph
    carries the CPDAG face when the final PDAG is a valid CPDAG and the
    ``mixed`` face otherwise (forced knowledge orientations can leave the
    class; ``diagnostics["cpdag_valid"]`` says which happened).
    """

    ingredient: ClassVar[str] = "score"

    def __init__(
        self,
        score: Score | str = "auto",
        *,
        use_bes: bool = True,
        restarts: int = 0,
        seed: int | None = None,
    ) -> None:
        if not isinstance(score, Score | str):
            raise TypeError(f"score must be a Score or a string, got {type(score).__name__}")
        if restarts < 0:
            raise ValueError(f"restarts must be non-negative, got {restarts}")
        self._score = score
        self._use_bes = use_bes
        self._restarts = restarts
        self._seed = seed

    # -- the public entry -------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Sweep permutations to convergence and return the class found."""
        score = resolve_score(self._score, data)
        variables = list(score.variables)
        if not type(score).score_equivalent:
            warnings.warn(
                f"{type(score).__name__} declares score_equivalent=False; BOSS scores "
                "orderings through projected DAGs and reports an equivalence class, "
                "so the search may not be optimal",
                UserWarning,
                stacklevel=2,
            )
        req_parents, req_children = self._required_maps(variables, knowledge)
        starts = [_required_consistent(variables, req_parents)]
        if self._restarts > 0:
            rng = np.random.default_rng(0 if self._seed is None else self._seed)
            for _ in range(self._restarts):
                shuffled = [variables[int(i)] for i in rng.permutation(len(variables))]
                starts.append(_required_consistent(shuffled, req_parents))
        best_order: list[str] = []
        best_total = -np.inf
        n_sweeps = 0
        n_moves = 0
        for start in starts:
            order, total, sweeps, moves = self._search(
                score, start, knowledge, req_parents, req_children
            )
            n_sweeps += sweeps
            n_moves += moves
            if total > best_total:
                best_order, best_total = order, total
        parents_of, _ = _project(score, best_order, knowledge)
        edges = [f"{u} -> {v}" for v in variables for u in sorted(parents_of[v])]
        state = cpdag_of_dag(dag(edges, nodes=variables))
        n_bes_deletes = 0
        if self._use_bes:
            while (move := GES._best_delete(score, state, variables, knowledge)) is not None:
                state = GES._apply_delete(state, move)
                n_bes_deletes += 1
        final, cpdag_valid = GES._face(GES._knowledge_tail(state, knowledge))
        if knowledge is not None and not knowledge.is_consistent(final):
            raise ValueError(
                "BOSS produced a graph that violates the domain knowledge; the data and "
                "the constraints disagree (knowledge.violations(graph) lists how)"
            )
        diagnostics = dict(score.diagnostics())
        diagnostics["n_sweeps"] = float(n_sweeps)
        diagnostics["n_moves"] = float(n_moves)
        diagnostics["n_bes_deletes"] = float(n_bes_deletes)
        diagnostics["cpdag_valid"] = cpdag_valid
        return DiscoveryResult(
            graph=final,
            method="boss",
            score=score.name,
            alpha=None,
            n=score.n,
            diagnostics=diagnostics,
        )

    # -- the sweep ----------------------------------------------------------------

    @staticmethod
    def _search(
        score: Score,
        start: list[str],
        knowledge: DomainKnowledge | None,
        req_parents: dict[str, frozenset[str]],
        req_children: dict[str, frozenset[str]],
    ) -> tuple[list[str], float, int, int]:
        """Best-move sweeps from ``start`` until one passes with no move.

        A sweep visits every variable in sorted name order, scores the
        permutation with the variable inserted at every admissible position
        (required parents stay before it, required children after), and
        moves it only when the best position beats the current one by more
        than the slack; the earliest best position wins ties. Returns the
        converged order, its total, and the sweep and move counts.
        """
        order = list(start)
        n_sweeps = 0
        n_moves = 0
        moved = True
        while moved:
            moved = False
            n_sweeps += 1
            for node in sorted(order):
                rest = [u for u in order if u != node]
                current = order.index(node)
                lo = 0
                hi = len(rest)
                for i, u in enumerate(rest):
                    if u in req_parents[node]:
                        lo = i + 1
                    elif u in req_children[node]:
                        hi = min(hi, i)
                _, baseline = _project(score, order, knowledge)
                best_total = baseline
                best_pos = current
                for pos in range(lo, hi + 1):
                    if pos == current:
                        continue
                    candidate = [*rest[:pos], node, *rest[pos:]]
                    _, total = _project(score, candidate, knowledge)
                    if total > best_total:
                        best_total = total
                        best_pos = pos
                if best_pos != current and best_total > baseline + _SLACK:
                    order = [*rest[:best_pos], node, *rest[best_pos:]]
                    n_moves += 1
                    moved = True
        _, total = _project(score, order, knowledge)
        return order, total, n_sweeps, n_moves

    # -- required-edge structure -----------------------------------------------------

    @staticmethod
    def _required_maps(
        variables: list[str], knowledge: DomainKnowledge | None
    ) -> tuple[dict[str, frozenset[str]], dict[str, frozenset[str]]]:
        """Per-node required parents and children, restricted to ``variables``."""
        req_parents: dict[str, set[str]] = {v: set() for v in variables}
        req_children: dict[str, set[str]] = {v: set() for v in variables}
        if knowledge is not None:
            names = set(variables)
            for u, v in sorted(knowledge.required):
                if u in names and v in names:
                    req_parents[v].add(u)
                    req_children[u].add(v)
        return (
            {v: frozenset(ps) for v, ps in req_parents.items()},
            {v: frozenset(cs) for v, cs in req_children.items()},
        )

fit(data=None, knowledge=None)

Sweep permutations to convergence and return the class found.

Source code in src/ergodic/discovery/algorithms/_boss.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Sweep permutations to convergence and return the class found."""
    score = resolve_score(self._score, data)
    variables = list(score.variables)
    if not type(score).score_equivalent:
        warnings.warn(
            f"{type(score).__name__} declares score_equivalent=False; BOSS scores "
            "orderings through projected DAGs and reports an equivalence class, "
            "so the search may not be optimal",
            UserWarning,
            stacklevel=2,
        )
    req_parents, req_children = self._required_maps(variables, knowledge)
    starts = [_required_consistent(variables, req_parents)]
    if self._restarts > 0:
        rng = np.random.default_rng(0 if self._seed is None else self._seed)
        for _ in range(self._restarts):
            shuffled = [variables[int(i)] for i in rng.permutation(len(variables))]
            starts.append(_required_consistent(shuffled, req_parents))
    best_order: list[str] = []
    best_total = -np.inf
    n_sweeps = 0
    n_moves = 0
    for start in starts:
        order, total, sweeps, moves = self._search(
            score, start, knowledge, req_parents, req_children
        )
        n_sweeps += sweeps
        n_moves += moves
        if total > best_total:
            best_order, best_total = order, total
    parents_of, _ = _project(score, best_order, knowledge)
    edges = [f"{u} -> {v}" for v in variables for u in sorted(parents_of[v])]
    state = cpdag_of_dag(dag(edges, nodes=variables))
    n_bes_deletes = 0
    if self._use_bes:
        while (move := GES._best_delete(score, state, variables, knowledge)) is not None:
            state = GES._apply_delete(state, move)
            n_bes_deletes += 1
    final, cpdag_valid = GES._face(GES._knowledge_tail(state, knowledge))
    if knowledge is not None and not knowledge.is_consistent(final):
        raise ValueError(
            "BOSS produced a graph that violates the domain knowledge; the data and "
            "the constraints disagree (knowledge.violations(graph) lists how)"
        )
    diagnostics = dict(score.diagnostics())
    diagnostics["n_sweeps"] = float(n_sweeps)
    diagnostics["n_moves"] = float(n_moves)
    diagnostics["n_bes_deletes"] = float(n_bes_deletes)
    diagnostics["cpdag_valid"] = cpdag_valid
    return DiscoveryResult(
        graph=final,
        method="boss",
        score=score.name,
        alpha=None,
        n=score.n,
        diagnostics=diagnostics,
    )

BootstrapResult dataclass

Edge frequencies from rerunning one discovery method on bootstrap resamples.

edges maps an unordered node pair to the fraction of successful runs in which the pair was adjacent, whatever the marks. marks maps a canonical edge glyph (:meth:~ergodic.graph.Edge.glyph, for example "X --> Y" or "X <-> Y") to its frequency, so one pair can spread its mass over several glyphs; the glyph frequencies for a pair sum to its adjacency frequency. Both mappings hold only pairs and glyphs that appeared at least once, with keys in sorted order. Frequencies are over the n_runs successful replicates, not over what was attempted. A temporal method keys them on the window nodes, "X[t-1] --> Y[t]".

Source code in src/ergodic/discovery/_bootstrap.py

@dataclass(frozen=True)
class BootstrapResult:
    """Edge frequencies from rerunning one discovery method on bootstrap resamples.

    ``edges`` maps an unordered node pair to the fraction of successful runs
    in which the pair was adjacent, whatever the marks. ``marks`` maps a
    canonical edge glyph (:meth:`~ergodic.graph.Edge.glyph`, for example
    ``"X --> Y"`` or ``"X <-> Y"``) to its frequency, so one pair can spread
    its mass over several glyphs; the glyph frequencies for a pair sum to its
    adjacency frequency. Both mappings hold only pairs and glyphs that
    appeared at least once, with keys in sorted order. Frequencies are over
    the ``n_runs`` successful replicates, not over what was attempted. A
    temporal method keys them on the window nodes, ``"X[t-1] --> Y[t]"``.
    """

    method: str
    n_runs: int
    n_failed: int
    edges: Mapping[frozenset[str], float]
    marks: Mapping[str, float]
    variables: tuple[str, ...] = field(repr=False)

    def frequency(self, a: str, b: str) -> float:
        """The adjacency frequency of the pair, ``0.0`` when it never appeared."""
        return self.edges.get(frozenset((a, b)), 0.0)

    def summary(self, threshold: float = 0.5) -> MixedGraph:
        """The stability graph: pairs at or above ``threshold``, one edge each.

        A qualifying pair gets its modal glyph, the orientation it wore most
        often (ties broken by the sorted glyph string, so the result is
        deterministic). The graph keeps the ``mixed`` kind: a vote over graphs
        carries no equivalence-class guarantee, even when every replicate was
        a CPDAG. Nodes are all the data's variables, so a variable that never
        earned a stable edge shows up isolated. ``threshold`` must be in
        ``(0, 1]``.
        """
        if not 0.0 < threshold <= 1.0:
            raise ValueError(f"threshold must be in (0, 1], got {threshold}")
        votes: dict[frozenset[str], list[tuple[float, str]]] = {}
        for glyph, freq in self.marks.items():
            parts = glyph.split()
            votes.setdefault(frozenset((parts[0], parts[-1])), []).append((freq, glyph))
        winners: list[str] = []
        for pair, adjacency in self.edges.items():
            if adjacency >= threshold:
                # highest frequency wins; on a tie, the smaller glyph string
                best = min(votes[pair], key=lambda vote: (-vote[0], vote[1]))
                winners.append(best[1])
        return mixed(winners, nodes=self.variables)

    def __repr__(self) -> str:
        stable = sum(1 for freq in self.edges.values() if freq >= 0.5)
        return (
            f"BootstrapResult({self.method}: {self.n_runs} runs, {self.n_failed} failed, "
            f"{stable}/{len(self.edges)} edges stable at 0.5)"
        )

frequency(a, b)

The adjacency frequency of the pair, 0.0 when it never appeared.

Source code in src/ergodic/discovery/_bootstrap.py

def frequency(self, a: str, b: str) -> float:
    """The adjacency frequency of the pair, ``0.0`` when it never appeared."""
    return self.edges.get(frozenset((a, b)), 0.0)

summary(threshold=0.5)

The stability graph: pairs at or above threshold, one edge each.

A qualifying pair gets its modal glyph, the orientation it wore most often (ties broken by the sorted glyph string, so the result is deterministic). The graph keeps the mixed kind: a vote over graphs carries no equivalence-class guarantee, even when every replicate was a CPDAG. Nodes are all the data's variables, so a variable that never earned a stable edge shows up isolated. threshold must be in (0, 1].

Source code in src/ergodic/discovery/_bootstrap.py

def summary(self, threshold: float = 0.5) -> MixedGraph:
    """The stability graph: pairs at or above ``threshold``, one edge each.

    A qualifying pair gets its modal glyph, the orientation it wore most
    often (ties broken by the sorted glyph string, so the result is
    deterministic). The graph keeps the ``mixed`` kind: a vote over graphs
    carries no equivalence-class guarantee, even when every replicate was
    a CPDAG. Nodes are all the data's variables, so a variable that never
    earned a stable edge shows up isolated. ``threshold`` must be in
    ``(0, 1]``.
    """
    if not 0.0 < threshold <= 1.0:
        raise ValueError(f"threshold must be in (0, 1], got {threshold}")
    votes: dict[frozenset[str], list[tuple[float, str]]] = {}
    for glyph, freq in self.marks.items():
        parts = glyph.split()
        votes.setdefault(frozenset((parts[0], parts[-1])), []).append((freq, glyph))
    winners: list[str] = []
    for pair, adjacency in self.edges.items():
        if adjacency >= threshold:
            # highest frequency wins; on a tie, the smaller glyph string
            best = min(votes[pair], key=lambda vote: (-vote[0], vote[1]))
            winners.append(best[1])
    return mixed(winners, nodes=self.variables)

CIResult dataclass

One conditional-independence query: the statistic, its p-value, and dof.

Source code in src/ergodic/discovery/_ci.py

@dataclass(frozen=True)
class CIResult:
    """One conditional-independence query: the statistic, its p-value, and dof."""

    statistic: float
    p_value: float
    dof: float | None = None

CITest

Base for conditional-independence tests; see the module docstring.

Subclasses set name and needs, may take options in __init__ (set them before calling super().__init__(data), since binding may read them), and implement statistic. The optional _prepare hook runs once after binding, for sufficient statistics.

Source code in src/ergodic/discovery/_ci.py

class CITest:
    """Base for conditional-independence tests; see the module docstring.

    Subclasses set ``name`` and ``needs``, may take options in ``__init__``
    (set them before calling ``super().__init__(data)``, since binding may
    read them), and implement ``statistic``. The optional ``_prepare`` hook
    runs once after binding, for sufficient statistics.
    """

    name: ClassVar[str] = "ci_test"
    needs: ClassVar[str] = "numeric"  # "numeric" or "codes"

    def __init__(self, data: TabularData | None = None) -> None:
        self._data: TabularData | None = None
        self._n: int = 0
        self._names: frozenset[str] = frozenset()
        self._blocks: dict[str, npt.NDArray[np.float64]] = {}
        self._codes: dict[str, npt.NDArray[np.int64]] = {}
        self._cards: dict[str, int] = {}
        self._cache: dict[tuple[str, str, frozenset[str]], CIResult] = {}
        self._n_tests: dict[int, int] = {}
        self._cache_hits: int = 0
        if data is not None:
            self._attach(data)

    # -- binding ---------------------------------------------------------

    @property
    def is_bound(self) -> bool:
        return self._data is not None

    @property
    def data(self) -> TabularData | None:
        """The dataset this test is bound to, after dropping incomplete rows."""
        return self._data

    @property
    def n(self) -> int:
        """Rows the oracle answers from (after dropping incomplete rows)."""
        return self._n

    @property
    def variables(self) -> tuple[str, ...]:
        """The variables this oracle answers about, sorted."""
        return tuple(sorted(self._names))

    def bind(self, data: TabularData) -> Self:
        """A new instance of this test bound to ``data``; never mutates ``self``."""
        fresh = copy.copy(self)
        fresh._attach(data)
        return fresh

    def _attach(self, data: TabularData) -> None:
        complete = data.dropna()
        self._data = complete
        self._n = complete.n_rows
        self._names = frozenset(complete.variables)
        self._cache = {}
        self._n_tests = {}
        self._cache_hits = 0
        if self.needs == "numeric":
            self._blocks = _numeric_blocks(complete)
        elif self.needs == "codes":
            self._codes, self._cards = _codes_columns(complete)
        else:
            raise ValueError(f"{type(self).__name__}.needs must be 'numeric' or 'codes'")
        self._prepare()

    def _prepare(self) -> None:
        """Hook for one-time sufficient statistics after binding."""

    # -- the query ---------------------------------------------------------

    def test(self, x: str, y: str, given: Iterable[str] = ()) -> CIResult:
        """The cached conditional-independence query, symmetric in ``x`` and ``y``."""
        self._require_bound()
        conditioning = frozenset(given)
        self._validate(x, y, conditioning)
        a, b = sorted((x, y))
        key = (a, b, conditioning)
        hit = self._cache.get(key)
        if hit is not None:
            self._cache_hits += 1
            return hit
        result = self._dispatch(a, b, conditioning)
        self._cache[key] = result
        size = len(conditioning)
        self._n_tests[size] = self._n_tests.get(size, 0) + 1
        return result

    def independent(self, x: str, y: str, given: Iterable[str] = (), *, alpha: float) -> bool:
        """The decision the algorithms read: fail to reject at ``alpha``."""
        if not 0.0 < alpha < 1.0:
            raise ValueError(f"alpha must be in (0, 1), got {alpha}")
        return self.test(x, y, given).p_value > alpha

    def diagnostics(self) -> dict[str, float]:
        """Run counters: total queries, largest conditioning set, cache hits."""
        total = float(sum(self._n_tests.values()))
        largest = float(max(self._n_tests, default=0))
        return {
            "n_ci_tests": total,
            "max_cond_size": largest,
            "cache_hits": float(self._cache_hits),
        }

    # -- subclass surface ---------------------------------------------------

    def statistic(self, *args: Any) -> tuple[float, float, float | None]:
        """Return ``(statistic, p_value, dof)``; signature depends on ``needs``.

        ``numeric``: ``statistic(x, y, z)`` with ``x`` ``(n, dx)``, ``y``
        ``(n, dy)``, ``z`` ``(n, dz)`` float blocks. ``codes``:
        ``statistic(x, y, z, cards)`` with ``x``, ``y`` ``(n,)`` integer
        codes, ``z`` a tuple of ``(n,)`` code columns, and ``cards`` the
        matching ``(card_x, card_y, cards_z)``.
        """
        raise NotImplementedError(f"{type(self).__name__} does not implement statistic")

    # -- internals -----------------------------------------------------------

    def _require_bound(self) -> None:
        if not self.is_bound:
            raise RuntimeError(
                f"{type(self).__name__} is not bound to data; construct it with data "
                "or call bind(data)"
            )

    def _validate(self, x: str, y: str, given: frozenset[str]) -> None:
        unknown = sorted(({x, y} | given) - self._names)
        if unknown:
            raise ValueError(f"unknown variable(s) {unknown}; the data has {sorted(self._names)}")
        if x == y:
            raise ValueError(f"x and y must differ; both are {x!r}")
        if x in given or y in given:
            raise ValueError("the conditioning set must not contain x or y")

    def _dispatch(self, x: str, y: str, given: frozenset[str]) -> CIResult:
        ordered = sorted(given)
        if self.needs == "numeric":
            xb = self._blocks[x]
            yb = self._blocks[y]
            zb = (
                np.hstack([self._blocks[w] for w in ordered])
                if ordered
                else np.empty((self._n, 0), dtype=np.float64)
            )
            stat, p, dof = self.statistic(xb, yb, zb)
        else:
            zc = tuple(self._codes[w] for w in ordered)
            cards = (self._cards[x], self._cards[y], tuple(self._cards[w] for w in ordered))
            stat, p, dof = self.statistic(self._codes[x], self._codes[y], zc, cards)
        return CIResult(statistic=float(stat), p_value=float(p), dof=dof)

data property

The dataset this test is bound to, after dropping incomplete rows.

n property

Rows the oracle answers from (after dropping incomplete rows).

variables property

The variables this oracle answers about, sorted.

bind(data)

A new instance of this test bound to data; never mutates self.

Source code in src/ergodic/discovery/_ci.py

def bind(self, data: TabularData) -> Self:
    """A new instance of this test bound to ``data``; never mutates ``self``."""
    fresh = copy.copy(self)
    fresh._attach(data)
    return fresh

diagnostics()

Run counters: total queries, largest conditioning set, cache hits.

Source code in src/ergodic/discovery/_ci.py

def diagnostics(self) -> dict[str, float]:
    """Run counters: total queries, largest conditioning set, cache hits."""
    total = float(sum(self._n_tests.values()))
    largest = float(max(self._n_tests, default=0))
    return {
        "n_ci_tests": total,
        "max_cond_size": largest,
        "cache_hits": float(self._cache_hits),
    }

independent(x, y, given=(), *, alpha)

The decision the algorithms read: fail to reject at alpha.

Source code in src/ergodic/discovery/_ci.py

def independent(self, x: str, y: str, given: Iterable[str] = (), *, alpha: float) -> bool:
    """The decision the algorithms read: fail to reject at ``alpha``."""
    if not 0.0 < alpha < 1.0:
        raise ValueError(f"alpha must be in (0, 1), got {alpha}")
    return self.test(x, y, given).p_value > alpha

statistic(*args)

Return (statistic, p_value, dof); signature depends on needs.

numeric: statistic(x, y, z) with x (n, dx), y (n, dy), z (n, dz) float blocks. codes: statistic(x, y, z, cards) with x, y (n,) integer codes, z a tuple of (n,) code columns, and cards the matching (card_x, card_y, cards_z).

Source code in src/ergodic/discovery/_ci.py

def statistic(self, *args: Any) -> tuple[float, float, float | None]:
    """Return ``(statistic, p_value, dof)``; signature depends on ``needs``.

    ``numeric``: ``statistic(x, y, z)`` with ``x`` ``(n, dx)``, ``y``
    ``(n, dy)``, ``z`` ``(n, dz)`` float blocks. ``codes``:
    ``statistic(x, y, z, cards)`` with ``x``, ``y`` ``(n,)`` integer
    codes, ``z`` a tuple of ``(n,)`` code columns, and ``cards`` the
    matching ``(card_x, card_y, cards_z)``.
    """
    raise NotImplementedError(f"{type(self).__name__} does not implement statistic")

test(x, y, given=())

The cached conditional-independence query, symmetric in x and y.

Source code in src/ergodic/discovery/_ci.py

def test(self, x: str, y: str, given: Iterable[str] = ()) -> CIResult:
    """The cached conditional-independence query, symmetric in ``x`` and ``y``."""
    self._require_bound()
    conditioning = frozenset(given)
    self._validate(x, y, conditioning)
    a, b = sorted((x, y))
    key = (a, b, conditioning)
    hit = self._cache.get(key)
    if hit is not None:
        self._cache_hits += 1
        return hit
    result = self._dispatch(a, b, conditioning)
    self._cache[key] = result
    size = len(conditioning)
    self._n_tests[size] = self._n_tests.get(size, 0) + 1
    return result

DSeparationOracle

Bases: CITest

The ground-truth oracle: answers queries by m-separation in a known graph.

Quacks like any bound test (p-value 1.0 when separated, 0.0 when connected), so every constraint algorithm can be checked for exactness and run in teaching mode. Needs no data; bind returns self.

Latent and selection nodes are not query-able variables: the oracle exposes the observed margin, which is what a constraint algorithm searching over data would see. Separation is still computed on the full graph, latents included, and every query conditions on the selection nodes, because data sampled under selection only ever shows the distribution given them.

Source code in src/ergodic/discovery/_ci.py

class DSeparationOracle(CITest):
    """The ground-truth oracle: answers queries by m-separation in a known graph.

    Quacks like any bound test (p-value 1.0 when separated, 0.0 when
    connected), so every constraint algorithm can be checked for exactness and
    run in teaching mode. Needs no data; ``bind`` returns ``self``.

    Latent and selection nodes are not query-able variables: the oracle
    exposes the observed margin, which is what a constraint algorithm
    searching over data would see. Separation is still computed on the full
    graph, latents included, and every query conditions on the selection
    nodes, because data sampled under selection only ever shows the
    distribution given them.
    """

    name: ClassVar[str] = "d_separation"

    def __init__(self, graph: MixedGraph) -> None:
        super().__init__(None)
        self._graph = graph
        self._names = frozenset(graph.nodes) - graph.latent - graph.selection

    @property
    def is_bound(self) -> bool:
        return True

    @property
    def graph(self) -> MixedGraph:
        return self._graph

    def bind(self, data: TabularData) -> Self:
        return self

    def _require_bound(self) -> None:
        return

    def _dispatch(self, x: str, y: str, given: frozenset[str]) -> CIResult:
        separated = self._graph.m_separated(x, y, given | self._graph.selection)
        return CIResult(statistic=0.0, p_value=1.0 if separated else 0.0, dof=None)

DirectLiNGAM

DirectLiNGAM over continuous tabular data.

Parameters:

Name	Type	Description	Default
prune	str	"bic" (the default) backward-eliminates each variable's predecessors by Gaussian BIC on the original data; "none" keeps every predecessor, the full lower-triangular DAG of the found order.	'bic'

fit(data, knowledge) returns a :class:DiscoveryResult whose graph is a DAG: with independent non-Gaussian noises the whole orientation is identified, so there is no equivalence class to report. The class consumes neither a CI test nor a decomposable score (ingredient is None); the data is its only ingredient.

Source code in src/ergodic/discovery/algorithms/_lingam.py

class DirectLiNGAM:
    """DirectLiNGAM over continuous tabular data.

    Args:
        prune: ``"bic"`` (the default) backward-eliminates each variable's
            predecessors by Gaussian BIC on the original data; ``"none"``
            keeps every predecessor, the full lower-triangular DAG of the
            found order.

    ``fit(data, knowledge)`` returns a :class:`DiscoveryResult` whose graph
    is a DAG: with independent non-Gaussian noises the whole orientation is
    identified, so there is no equivalence class to report. The class
    consumes neither a CI test nor a decomposable score (``ingredient`` is
    ``None``); the data is its only ingredient.
    """

    ingredient: ClassVar[str | None] = None

    def __init__(self, prune: str = "bic") -> None:
        if prune not in ("bic", "none"):
            raise ValueError(f"prune must be 'bic' or 'none', got {prune!r}")
        self.prune = prune

    # -- the run ---------------------------------------------------------

    def fit(self, data: TabularData, knowledge: DomainKnowledge | None = None) -> DiscoveryResult:
        """Find the causal order, prune the parents, and return the DAG."""
        if not isinstance(data, TabularData):
            raise TypeError(
                f"DirectLiNGAM reads TabularData, got {type(data).__name__}; temporal "
                "shapes have their own methods (method='pcmci' for time series and panels)"
            )
        complete = data.dropna()
        names = sorted(complete.variables)
        wrong = [v for v in names if not complete.schema[v].is_continuous]
        if wrong:
            raise ValueError(
                f"column(s) {wrong} are not continuous; the LiNGAM model is linear in "
                "continuous variables, so retype them or mark them excluded"
            )
        if complete.n_rows < 2:
            raise ValueError(f"DirectLiNGAM needs at least 2 complete rows, got {complete.n_rows}")
        matrix = complete.matrix(names)
        order = _causal_order(matrix, names, knowledge)
        _check_required_order(order, knowledge)
        parents, n_pruned, n_knowledge_pruned = self._parents_of(matrix, names, order, knowledge)
        graph = dag(
            [(p, child) for child in order for p in parents[child]],
            nodes=names,
        )
        if knowledge is not None and not knowledge.is_consistent(graph):
            raise ValueError(
                "DirectLiNGAM produced a graph that violates the domain knowledge; the "
                "data and the constraints disagree (knowledge.violations(graph) lists how)"
            )
        diagnostics = {
            "n_root_picks": float(len(order)),
            "n_pruned_edges": float(n_pruned),
            "n_knowledge_pruned": float(n_knowledge_pruned),
        }
        return DiscoveryResult(
            graph=graph,
            method="lingam",
            test=None,
            score=None,
            alpha=None,
            n=complete.n_rows,
            diagnostics=diagnostics,
        )

    # -- pruning -----------------------------------------------------------

    def _parents_of(
        self,
        matrix: npt.NDArray[np.float64],
        names: list[str],
        order: list[str],
        knowledge: DomainKnowledge | None,
    ) -> tuple[dict[str, list[str]], int, int]:
        """Each variable's parents among its order predecessors, with counters.

        Knowledge filters the predecessor set first: a forbidden ``a -> b``
        never enters ``b``'s regression (counted in ``n_knowledge_pruned``)
        and a required edge can never leave it. ``prune="bic"`` then
        backward-eliminates; ``prune="none"`` keeps every survivor.
        """
        column = {v: k for k, v in enumerate(names)}
        parents: dict[str, list[str]] = {}
        n_pruned = 0
        n_knowledge_pruned = 0
        for t, target in enumerate(order):
            kept: list[str] = []
            required: set[str] = set()
            for p in order[:t]:
                if knowledge is not None and knowledge.requires_edge(p, target):
                    kept.append(p)
                    required.add(p)
                elif knowledge is None or knowledge.permits_edge(p, target):
                    kept.append(p)
                else:
                    n_knowledge_pruned += 1
            if self.prune == "bic":
                kept, dropped = _backward_bic(matrix, column, target, kept, frozenset(required))
                n_pruned += dropped
            parents[target] = sorted(kept)
        return parents, n_pruned, n_knowledge_pruned

fit(data, knowledge=None)

Find the causal order, prune the parents, and return the DAG.

Source code in src/ergodic/discovery/algorithms/_lingam.py

def fit(self, data: TabularData, knowledge: DomainKnowledge | None = None) -> DiscoveryResult:
    """Find the causal order, prune the parents, and return the DAG."""
    if not isinstance(data, TabularData):
        raise TypeError(
            f"DirectLiNGAM reads TabularData, got {type(data).__name__}; temporal "
            "shapes have their own methods (method='pcmci' for time series and panels)"
        )
    complete = data.dropna()
    names = sorted(complete.variables)
    wrong = [v for v in names if not complete.schema[v].is_continuous]
    if wrong:
        raise ValueError(
            f"column(s) {wrong} are not continuous; the LiNGAM model is linear in "
            "continuous variables, so retype them or mark them excluded"
        )
    if complete.n_rows < 2:
        raise ValueError(f"DirectLiNGAM needs at least 2 complete rows, got {complete.n_rows}")
    matrix = complete.matrix(names)
    order = _causal_order(matrix, names, knowledge)
    _check_required_order(order, knowledge)
    parents, n_pruned, n_knowledge_pruned = self._parents_of(matrix, names, order, knowledge)
    graph = dag(
        [(p, child) for child in order for p in parents[child]],
        nodes=names,
    )
    if knowledge is not None and not knowledge.is_consistent(graph):
        raise ValueError(
            "DirectLiNGAM produced a graph that violates the domain knowledge; the "
            "data and the constraints disagree (knowledge.violations(graph) lists how)"
        )
    diagnostics = {
        "n_root_picks": float(len(order)),
        "n_pruned_edges": float(n_pruned),
        "n_knowledge_pruned": float(n_knowledge_pruned),
    }
    return DiscoveryResult(
        graph=graph,
        method="lingam",
        test=None,
        score=None,
        alpha=None,
        n=complete.n_rows,
        diagnostics=diagnostics,
    )

DiscoveryResult dataclass

What a discovery run found, and how.

graph is the headline. test or score names the ingredient the algorithm consumed (one of them is None). diagnostics carries the run counters (tests by conditioning size, cache hits, score evaluations). sepsets maps an unordered removed pair to the set that separated it, when a constraint method produced one.

Source code in src/ergodic/discovery/_result.py

@dataclass(frozen=True)
class DiscoveryResult:
    """What a discovery run found, and how.

    ``graph`` is the headline. ``test`` or ``score`` names the ingredient the
    algorithm consumed (one of them is ``None``). ``diagnostics`` carries the
    run counters (tests by conditioning size, cache hits, score evaluations).
    ``sepsets`` maps an unordered removed pair to the set that separated it,
    when a constraint method produced one.
    """

    graph: MixedGraph
    method: str
    test: str | None = None
    score: str | None = None
    alpha: float | None = None
    n: int = 0
    diagnostics: Mapping[str, float] = field(default_factory=dict)
    sepsets: Mapping[frozenset[str], frozenset[str]] | None = field(default=None, repr=False)

    def __repr__(self) -> str:
        ingredient = self.test if self.test is not None else self.score
        body = f"{self.method}/{ingredient}: {self.graph.kind}"
        body += f", {len(self.graph.nodes)} nodes, {len(self.graph.edges)} edges"
        if self.alpha is not None:
            body += f", alpha={self.alpha:g}"
        return f"DiscoveryResult({body}, n={self.n})"

Discrete

Bases: CITest

The G squared conditional-independence test on contingency tables.

Degrees of freedom are adjusted for sparsity, bnlearn-style: each nonempty stratum of the conditioning configuration contributes (card_x - 1) * (card_y - 1), and configurations with no observed rows contribute nothing. The unadjusted product over the full configuration space would overstate the dof badly on sparse tables and make the test blind. When the dof comes out zero (a single-level endpoint, or no rows), the query is vacuous and the result is (0.0, 1.0, 0.0).

Source code in src/ergodic/discovery/ci/_discrete.py

class Discrete(CITest):
    """The G squared conditional-independence test on contingency tables.

    Degrees of freedom are adjusted for sparsity, bnlearn-style: each
    **nonempty** stratum of the conditioning configuration contributes
    ``(card_x - 1) * (card_y - 1)``, and configurations with no observed rows
    contribute nothing. The unadjusted product over the full configuration
    space would overstate the dof badly on sparse tables and make the test
    blind. When the dof comes out zero (a single-level endpoint, or no rows),
    the query is vacuous and the result is ``(0.0, 1.0, 0.0)``.
    """

    name: ClassVar[str] = "discrete"
    needs: ClassVar[str] = "codes"

    def statistic(
        self,
        x: npt.NDArray[np.int64],
        y: npt.NDArray[np.int64],
        z: tuple[npt.NDArray[np.int64], ...],
        cards: tuple[int, int, tuple[int, ...]],
    ) -> tuple[float, float, float | None]:
        card_x, card_y, _ = cards
        n = x.shape[0]
        if z:
            _, strata = np.unique(np.stack(z, axis=1), axis=0, return_inverse=True)
            strata = strata.reshape(-1)
            n_strata = int(strata.max()) + 1 if strata.size else 0
        else:
            strata = np.zeros(n, dtype=np.int64)
            n_strata = 1 if n else 0
        per_stratum = (card_x - 1) * (card_y - 1)
        g2 = 0.0
        dof = 0.0
        for s in range(n_strata):
            mask = strata == s
            counts = np.bincount(x[mask] * card_y + y[mask], minlength=card_x * card_y)
            table = counts.reshape(card_x, card_y).astype(np.float64)
            expected = np.outer(table.sum(axis=1), table.sum(axis=0)) / table.sum()
            seen = table > 0
            g2 += 2.0 * float(np.sum(table[seen] * np.log(table[seen] / expected[seen])))
            dof += per_stratum
        if dof == 0.0:
            return (0.0, 1.0, 0.0)
        return g2, _chi2_sf(g2, dof), dof

DiscreteBIC

Bases: Score

The multinomial BIC, decomposable and score-equivalent.

Every column must be discrete (categorical, ordinal, or binary): the score runs on integer codes, so binding to data with a continuous column raises TypeError. Parameter counts agree across a Markov equivalence class, so Markov-equivalent DAGs get equal totals, exactly as with the Gaussian BIC.

Parameters:

Name	Type	Description	Default
data	TabularData | None	Bind at construction, or later through bind.	None
penalty	float	A multiplier on the complexity term ((r - 1) q / 2) log n. 1.0 is the textbook BIC; smaller values prune less.	1.0

Source code in src/ergodic/discovery/scores/_discrete_bic.py

class DiscreteBIC(Score):
    """The multinomial BIC, decomposable and score-equivalent.

    Every column must be discrete (categorical, ordinal, or binary): the
    score runs on integer codes, so binding to data with a continuous column
    raises ``TypeError``. Parameter counts agree across a Markov equivalence
    class, so Markov-equivalent DAGs get equal totals, exactly as with the
    Gaussian ``BIC``.

    Args:
        data: Bind at construction, or later through ``bind``.
        penalty: A multiplier on the complexity term ``((r - 1) q / 2) log n``.
            ``1.0`` is the textbook BIC; smaller values prune less.
    """

    name: ClassVar[str] = "bic_discrete"
    needs: ClassVar[str] = "codes"
    decomposable: ClassVar[bool] = True
    score_equivalent: ClassVar[bool] = True

    def __init__(self, data: TabularData | None = None, *, penalty: float = 1.0) -> None:
        if penalty <= 0.0:
            raise ValueError(f"penalty must be positive, got {penalty}")
        self._penalty = float(penalty)
        super().__init__(data)

    @property
    def penalty(self) -> float:
        """The multiplier on the complexity term."""
        return self._penalty

    def local_from_codes(
        self,
        y: npt.NDArray[np.int64],
        y_card: int,
        zs: tuple[npt.NDArray[np.int64], ...],
        z_cards: tuple[int, ...],
    ) -> float:
        """The multinomial local from integer codes; see the module docstring."""
        n = int(y.shape[0])
        if n == 0:
            raise ValueError("DiscreteBIC needs at least one row")
        r = int(y_card)
        q = 1
        for card in z_cards:
            q *= int(card)
        flat = np.ravel_multi_index((*zs, y), (*z_cards, r))
        counts = np.bincount(flat, minlength=q * r).reshape(q, r)
        totals = np.broadcast_to(counts.sum(axis=1, keepdims=True), counts.shape)
        observed = counts > 0
        cells = counts[observed].astype(np.float64)
        rows = totals[observed].astype(np.float64)
        loglik = float(np.sum(cells * (np.log(cells) - np.log(rows))))
        complexity = self._penalty * ((r - 1) * q / 2.0) * math.log(n)
        return loglik - complexity

penalty property

The multiplier on the complexity term.

local_from_codes(y, y_card, zs, z_cards)

The multinomial local from integer codes; see the module docstring.

Source code in src/ergodic/discovery/scores/_discrete_bic.py

def local_from_codes(
    self,
    y: npt.NDArray[np.int64],
    y_card: int,
    zs: tuple[npt.NDArray[np.int64], ...],
    z_cards: tuple[int, ...],
) -> float:
    """The multinomial local from integer codes; see the module docstring."""
    n = int(y.shape[0])
    if n == 0:
        raise ValueError("DiscreteBIC needs at least one row")
    r = int(y_card)
    q = 1
    for card in z_cards:
        q *= int(card)
    flat = np.ravel_multi_index((*zs, y), (*z_cards, r))
    counts = np.bincount(flat, minlength=q * r).reshape(q, r)
    totals = np.broadcast_to(counts.sum(axis=1, keepdims=True), counts.shape)
    observed = counts > 0
    cells = counts[observed].astype(np.float64)
    rows = totals[observed].astype(np.float64)
    loglik = float(np.sum(cells * (np.log(cells) - np.log(rows))))
    complexity = self._penalty * ((r - 1) * q / 2.0) * math.log(n)
    return loglik - complexity

FCI

FCI over a conditional-independence test, returning a PAG.

Parameters:

Name	Type	Description	Default
test	CITest | str	A :class:CITest instance (bound or unbound, including :class:DSeparationOracle), or a name: "auto" (pick by the data schema), "linear"/"fisher_z", "discrete"/ "g_sq", "mixed"/"dg", "gcm".	'auto'
alpha	float	Significance level for the independence decisions.	0.05
max_cond_size	int | None	Cap on the conditioning-set size, applied to the skeleton phase and to the possible-d-sep subsets alike; None for no cap. Possible-d-sep sets can span most of the graph, so the cap is the honest scaling knob.	None

fit(data, knowledge) returns a :class:DiscoveryResult whose graph carries the PAG face when the result is a valid complete PAG and the mixed face otherwise (noisy tests or knowledge-seeded marks). With a perfect oracle and no knowledge the output is exactly the canonical PAG: pag_of_mag(mag_of_dag(truth)).

Source code in src/ergodic/discovery/algorithms/_fci.py

class FCI:
    """FCI over a conditional-independence test, returning a PAG.

    Args:
        test: A :class:`CITest` instance (bound or unbound, including
            :class:`DSeparationOracle`), or a name: ``"auto"`` (pick by the
            data schema), ``"linear"``/``"fisher_z"``, ``"discrete"``/
            ``"g_sq"``, ``"mixed"``/``"dg"``, ``"gcm"``.
        alpha: Significance level for the independence decisions.
        max_cond_size: Cap on the conditioning-set size, applied to the
            skeleton phase and to the possible-d-sep subsets alike; ``None``
            for no cap. Possible-d-sep sets can span most of the graph, so
            the cap is the honest scaling knob.

    ``fit(data, knowledge)`` returns a :class:`DiscoveryResult` whose graph
    carries the PAG face when the result is a valid complete PAG and the
    ``mixed`` face otherwise (noisy tests or knowledge-seeded marks). With a
    perfect oracle and no knowledge the output is exactly the canonical PAG:
    ``pag_of_mag(mag_of_dag(truth))``.
    """

    ingredient: ClassVar[str] = "test"

    def __init__(
        self,
        test: CITest | str = "auto",
        *,
        alpha: float = 0.05,
        max_cond_size: int | None = None,
    ) -> None:
        if not 0.0 < alpha < 1.0:
            raise ValueError(f"alpha must be in (0, 1), got {alpha}")
        if max_cond_size is not None and max_cond_size < 0:
            raise ValueError(f"max_cond_size must be None or >= 0, got {max_cond_size}")
        self.test = test
        self.alpha = alpha
        self.max_cond_size = max_cond_size

    # -- the run ---------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Run FCI and return the discovered graph with its provenance."""
        oracle = resolve_test(self.test, data)
        nodes = sorted(oracle.variables)
        adj, sepsets = stable_skeleton(
            oracle, nodes, knowledge, alpha=self.alpha, max_cond_size=self.max_cond_size
        )
        protected = _protected_pairs(knowledge)

        marks = _circle_marks(adj)
        _orient_colliders(marks, adj, sepsets)
        n_pds_removed = self._pds_prune(oracle, nodes, adj, marks, sepsets, protected)

        marks = _circle_marks(adj)
        n_v_structures = _orient_colliders(marks, adj, sepsets)
        _seed_knowledge(marks, knowledge)

        def from_sepsets(d: str, a: str, b: str, c: str) -> bool | None:
            separator = sepsets.get(frozenset((d, c)))
            if separator is None:
                return None  # the pair was never tested (knowledge-pruned)
            return b not in separator

        closed = fci_closure(_to_graph(nodes, marks), discriminating=from_sepsets)
        try:
            final = closed.as_kind(GraphKind.PAG)
            pag_valid = 1.0
        except InvalidGraphError:
            final = closed.as_kind(GraphKind.MIXED)
            pag_valid = 0.0
        if knowledge is not None and not knowledge.is_consistent(final):
            raise ValueError(
                "FCI produced a graph that violates the domain knowledge; the data and "
                "the constraints disagree (knowledge.violations(graph) lists how)"
            )
        diagnostics: dict[str, float] = dict(oracle.diagnostics())
        diagnostics["n_edges"] = float(len(final.edges))
        diagnostics["n_v_structures"] = float(n_v_structures)
        diagnostics["n_pds_removed"] = float(n_pds_removed)
        for label, mark in (
            ("n_arrowheads", Mark.ARROW),
            ("n_tails", Mark.TAIL),
            ("n_circles", Mark.CIRCLE),
        ):
            diagnostics[label] = float(
                sum((e.u_mark is mark) + (e.v_mark is mark) for e in final.edges)
            )
        diagnostics["pag_valid"] = pag_valid
        return DiscoveryResult(
            graph=final,
            method="fci",
            test=oracle.name,
            alpha=self.alpha,
            n=oracle.n,
            diagnostics=diagnostics,
            sepsets=sepsets,
        )

    def _pds_prune(
        self,
        oracle: CITest,
        nodes: list[str],
        adj: _Adjacency,
        marks: _Marks,
        sepsets: _Sepsets,
        protected: frozenset[frozenset[str]],
    ) -> int:
        """The FCI extension of the skeleton: re-test against Possible-D-Sep sets.

        The sets are computed once, on the adjacencies and v-structure marks
        as they stand after the first orientation, then every surviving pair
        is re-tested against subsets of pds(x) from each side in turn.
        Removals record their separating sets like any skeleton removal.
        """
        pds = {x: _possible_d_sep(x, adj, marks) for x in nodes}
        removed = 0
        for x in nodes:
            for y in sorted(adj[x]):
                if y not in adj[x]:
                    continue  # removed from the other side earlier
                if frozenset((x, y)) in protected:
                    continue
                candidates = sorted(pds[x] - {x, y})
                cap = len(candidates) if self.max_cond_size is None else self.max_cond_size
                found = False
                for size in range(min(cap, len(candidates)) + 1):
                    for subset in combinations(candidates, size):
                        if oracle.independent(x, y, subset, alpha=self.alpha):
                            adj[x].discard(y)
                            adj[y].discard(x)
                            sepsets[frozenset((x, y))] = frozenset(subset)
                            removed += 1
                            found = True
                            break
                    if found:
                        break
        return removed

fit(data=None, knowledge=None)

Run FCI and return the discovered graph with its provenance.

Source code in src/ergodic/discovery/algorithms/_fci.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Run FCI and return the discovered graph with its provenance."""
    oracle = resolve_test(self.test, data)
    nodes = sorted(oracle.variables)
    adj, sepsets = stable_skeleton(
        oracle, nodes, knowledge, alpha=self.alpha, max_cond_size=self.max_cond_size
    )
    protected = _protected_pairs(knowledge)

    marks = _circle_marks(adj)
    _orient_colliders(marks, adj, sepsets)
    n_pds_removed = self._pds_prune(oracle, nodes, adj, marks, sepsets, protected)

    marks = _circle_marks(adj)
    n_v_structures = _orient_colliders(marks, adj, sepsets)
    _seed_knowledge(marks, knowledge)

    def from_sepsets(d: str, a: str, b: str, c: str) -> bool | None:
        separator = sepsets.get(frozenset((d, c)))
        if separator is None:
            return None  # the pair was never tested (knowledge-pruned)
        return b not in separator

    closed = fci_closure(_to_graph(nodes, marks), discriminating=from_sepsets)
    try:
        final = closed.as_kind(GraphKind.PAG)
        pag_valid = 1.0
    except InvalidGraphError:
        final = closed.as_kind(GraphKind.MIXED)
        pag_valid = 0.0
    if knowledge is not None and not knowledge.is_consistent(final):
        raise ValueError(
            "FCI produced a graph that violates the domain knowledge; the data and "
            "the constraints disagree (knowledge.violations(graph) lists how)"
        )
    diagnostics: dict[str, float] = dict(oracle.diagnostics())
    diagnostics["n_edges"] = float(len(final.edges))
    diagnostics["n_v_structures"] = float(n_v_structures)
    diagnostics["n_pds_removed"] = float(n_pds_removed)
    for label, mark in (
        ("n_arrowheads", Mark.ARROW),
        ("n_tails", Mark.TAIL),
        ("n_circles", Mark.CIRCLE),
    ):
        diagnostics[label] = float(
            sum((e.u_mark is mark) + (e.v_mark is mark) for e in final.edges)
        )
    diagnostics["pag_valid"] = pag_valid
    return DiscoveryResult(
        graph=final,
        method="fci",
        test=oracle.name,
        alpha=self.alpha,
        n=oracle.n,
        diagnostics=diagnostics,
        sepsets=sepsets,
    )

GCM

Bases: CITest

The generalized covariance measure over a regression learner.

Regresses x and y on z with a fresh clone of the learner per endpoint, forms the elementwise residual product, and reports sqrt(n) * mean(R) / sd(R) against the standard normal, two-sided. Residuals are in-sample, as in the paper's basic form. With an empty conditioning set the endpoints are centered instead, so the statistic is a plain covariance test.

Parameters:

Name	Type	Description	Default
data	TabularData | None	The dataset to bind, or None for a reusable config.	None
learner	Any	The regression slot, accepting exactly what an estimator's learner slot accepts: any sklearn-style regressor or an alias string like "ols" or "knn". None means the native :class:~ergodic.inference.KNNRegressor, constructed lazily at first use.	None

Source code in src/ergodic/discovery/ci/_gcm.py

class GCM(CITest):
    """The generalized covariance measure over a regression learner.

    Regresses x and y on z with a fresh clone of the learner per endpoint,
    forms the elementwise residual product, and reports
    ``sqrt(n) * mean(R) / sd(R)`` against the standard normal, two-sided.
    Residuals are in-sample, as in the paper's basic form. With an empty
    conditioning set the endpoints are centered instead, so the statistic is
    a plain covariance test.

    Args:
        data: The dataset to bind, or ``None`` for a reusable config.
        learner: The regression slot, accepting exactly what an estimator's
            learner slot accepts: any sklearn-style regressor or an alias
            string like ``"ols"`` or ``"knn"``. ``None`` means the native
            :class:`~ergodic.inference.KNNRegressor`, constructed lazily at
            first use.
    """

    name: ClassVar[str] = "gcm"
    needs: ClassVar[str] = "numeric"

    def __init__(self, data: TabularData | None = None, *, learner: Any = None) -> None:
        self.learner: Any = learner
        self._template: Learner | None = None
        super().__init__(data)

    def statistic(
        self,
        x: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        z: npt.NDArray[np.float64],
    ) -> tuple[float, float, float | None]:
        if x.shape[1] != 1 or y.shape[1] != 1:
            raise TypeError(
                "GCM needs scalar endpoints; a categorical endpoint expands to a "
                "one-hot block, use Mixed for categorical endpoints"
            )
        n = x.shape[0]
        if z.shape[1] == 0:
            rx = x[:, 0] - float(np.mean(x[:, 0]))
            ry = y[:, 0] - float(np.mean(y[:, 0]))
        else:
            rx = self._residual(z, x[:, 0])
            ry = self._residual(z, y[:, 0])
        product = rx * ry
        mean_r = float(np.mean(product))
        var_r = max(float(np.var(product)), _VAR_FLOOR)
        stat = math.sqrt(n) * mean_r / math.sqrt(var_r)
        p = 2.0 * (1.0 - _norm_cdf(abs(stat)))
        return stat, p, None

    # -- internals -----------------------------------------------------------

    def _slot(self) -> Learner:
        """The coerced learner template: resolved once, cloned fresh per fit."""
        if self._template is None:
            spec = KNNRegressor() if self.learner is None else self.learner
            self._template = _as_learner(spec)
        return self._template

    def _residual(
        self, z: npt.NDArray[np.float64], target: npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64]:
        """The in-sample residual of ``target`` regressed on ``z``."""
        model = _fresh(self._slot())
        model.fit(z, target.ravel())
        prediction = np.asarray(model.predict(z), dtype=np.float64).reshape(-1)
        return np.asarray(target.ravel() - prediction, dtype=np.float64)

GES

Greedy equivalence search over a decomposable score.

Parameters:

Name	Type	Description	Default
score	Score | str	A :class:~ergodic.discovery.Score instance (bound or unbound) or a name: "auto" (pick by the data schema), "bic", "bic_discrete", "bdeu", "bic_mixed".	'auto'
max_degree	int | None	A cap on the number of edges at any node; an Insert that touches a node already at the cap is never considered. None leaves degrees free.	None

fit(data, knowledge) returns a :class:DiscoveryResult whose graph carries the CPDAG face when the final PDAG is a valid CPDAG and the mixed face otherwise (forced knowledge orientations can leave the class; diagnostics["cpdag_valid"] says which happened).

Source code in src/ergodic/discovery/algorithms/_ges.py

class GES:
    """Greedy equivalence search over a decomposable score.

    Args:
        score: A :class:`~ergodic.discovery.Score` instance (bound or
            unbound) or a name: ``"auto"`` (pick by the data schema),
            ``"bic"``, ``"bic_discrete"``, ``"bdeu"``, ``"bic_mixed"``.
        max_degree: A cap on the number of edges at any node; an Insert
            that touches a node already at the cap is never considered.
            ``None`` leaves degrees free.

    ``fit(data, knowledge)`` returns a :class:`DiscoveryResult` whose graph
    carries the CPDAG face when the final PDAG is a valid CPDAG and the
    ``mixed`` face otherwise (forced knowledge orientations can leave the
    class; ``diagnostics["cpdag_valid"]`` says which happened).
    """

    ingredient: ClassVar[str] = "score"

    def __init__(self, score: Score | str = "auto", *, max_degree: int | None = None) -> None:
        if not isinstance(score, Score | str):
            raise TypeError(f"score must be a Score or a string, got {type(score).__name__}")
        if max_degree is not None and max_degree < 1:
            raise ValueError(f"max_degree must be None or at least 1, got {max_degree}")
        self._score = score
        self._max_degree = max_degree

    # -- the public entry -------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Run the forward and backward phases and return the class found."""
        score = resolve_score(self._score, data)
        variables = list(score.variables)
        if not type(score).score_equivalent:
            warnings.warn(
                f"{type(score).__name__} declares score_equivalent=False; GES assumes "
                "score equivalence when it moves between classes, so the search may "
                "not be optimal",
                UserWarning,
                stacklevel=2,
            )
        current = self._start(variables, knowledge)
        n_inserts = 0
        while (move := self._best_insert(score, current, variables, knowledge)) is not None:
            current = self._apply_insert(current, move)
            n_inserts += 1
        n_deletes = 0
        while (move := self._best_delete(score, current, variables, knowledge)) is not None:
            current = self._apply_delete(current, move)
            n_deletes += 1
        final, cpdag_valid = self._face(self._knowledge_tail(current, knowledge))
        if knowledge is not None and not knowledge.is_consistent(final):
            raise ValueError(
                "GES produced a graph that violates the domain knowledge; the data and "
                "the constraints disagree (knowledge.violations(graph) lists how)"
            )
        diagnostics = dict(score.diagnostics())
        diagnostics["n_inserts"] = float(n_inserts)
        diagnostics["n_deletes"] = float(n_deletes)
        diagnostics["n_edges"] = float(len(final.edges))
        diagnostics["n_directed"] = float(sum(1 for e in final.edges if e.is_directed))
        diagnostics["n_undirected"] = float(sum(1 for e in final.edges if e.is_undirected))
        diagnostics["cpdag_valid"] = cpdag_valid
        return DiscoveryResult(
            graph=final,
            method="ges",
            score=score.name,
            alpha=None,
            n=score.n,
            diagnostics=diagnostics,
        )

    # -- the start state ----------------------------------------------------------

    @staticmethod
    def _start(variables: list[str], knowledge: DomainKnowledge | None) -> MixedGraph:
        """The empty class, with any required edges seeded and completed in."""
        empty = MixedGraph(nodes=frozenset(variables), edges=frozenset(), kind=GraphKind.MIXED)
        if knowledge is None:
            return empty
        names = frozenset(variables)
        seeds = frozenset(
            Edge.make(u, Mark.TAIL, v, Mark.ARROW)
            for u, v in sorted(knowledge.required)
            if u in names and v in names
        )
        if not seeds:
            return empty
        extension = consistent_extension(replace(empty, edges=seeds))
        if extension is None:
            raise ValueError(
                "the required edges admit no DAG over these variables; "
                "check the knowledge for required-edge cycles"
            )
        return cpdag_of_dag(extension)

    # -- the operators --------------------------------------------------------------

    def _best_insert(
        self,
        score: Score,
        graph: MixedGraph,
        nodes: list[str],
        knowledge: DomainKnowledge | None,
    ) -> _Move | None:
        """The best strictly positive Insert(X, Y, T), or ``None`` at convergence.

        The sweep order is deterministic (X ascending, Y ascending, T by
        size then name) and only a strictly larger delta displaces the
        incumbent, so ties break toward the first candidate and runs repeat.
        """
        state = _state_of(graph, nodes)
        best: _Move | None = None
        for x in nodes:
            for y in nodes:
                if x == y or y in state.adjacent[x]:
                    continue
                if knowledge is not None and not knowledge.permits_edge(x, y):
                    continue
                if self._max_degree is not None and (
                    len(state.adjacent[x]) >= self._max_degree
                    or len(state.adjacent[y]) >= self._max_degree
                ):
                    continue
                na_yx = state.undirected[y] & state.adjacent[x]
                pool = sorted(state.undirected[y] - state.adjacent[x])
                for size in range(len(pool) + 1):
                    for combo in combinations(pool, size):
                        cond = na_yx | set(combo)
                        if not _is_clique(cond, state.adjacent):
                            continue
                        if _reaches_semidirected(y, x, cond, state):
                            continue
                        base = frozenset(state.parents[y] | cond)
                        delta = score.local(y, base | {x}) - score.local(y, base)
                        if delta > 0.0 and (best is None or delta > best.delta):
                            best = _Move(delta=delta, x=x, y=y, subset=frozenset(combo))
        return best

    @staticmethod
    def _best_delete(
        score: Score,
        graph: MixedGraph,
        nodes: list[str],
        knowledge: DomainKnowledge | None,
    ) -> _Move | None:
        """The best strictly positive Delete(X, Y, H), or ``None`` at convergence.

        Candidates are the ordered pairs whose edge reads ``x -> y`` or
        ``x -- y`` (never against an arrow); a pair the knowledge requires
        in either direction is protected. The sweep order and tie-breaking
        match :meth:`_best_insert`.
        """
        state = _state_of(graph, nodes)
        best: _Move | None = None
        for x in nodes:
            for y in nodes:
                if x == y:
                    continue
                if y not in state.children[x] and y not in state.undirected[x]:
                    continue
                if knowledge is not None and (
                    knowledge.requires_edge(x, y) or knowledge.requires_edge(y, x)
                ):
                    continue
                na_yx = state.undirected[y] & state.adjacent[x]
                pool = sorted(na_yx)
                for size in range(len(pool) + 1):
                    for combo in combinations(pool, size):
                        kept = na_yx - set(combo)
                        if not _is_clique(kept, state.adjacent):
                            continue
                        base = frozenset(state.parents[y] | kept)
                        delta = score.local(y, base - {x}) - score.local(y, base | {x})
                        if delta > 0.0 and (best is None or delta > best.delta):
                            best = _Move(delta=delta, x=x, y=y, subset=frozenset(combo))
        return best

    @staticmethod
    def _apply_insert(graph: MixedGraph, move: _Move) -> MixedGraph:
        """Add ``x -> y``, orient each ``t -- y`` as ``t -> y``, and complete."""
        working = replace(graph, kind=GraphKind.MIXED)
        working = working.add_edge((move.x, Mark.TAIL, move.y, Mark.ARROW))
        for t in sorted(move.subset):
            working = working.orient(t, move.y)
        return _complete(working)

    @staticmethod
    def _apply_delete(graph: MixedGraph, move: _Move) -> MixedGraph:
        """Drop the x-y edge, orient ``y -> h`` and any undirected x-h as ``x -> h``, complete."""
        working = replace(graph, kind=GraphKind.MIXED).remove_edge(move.x, move.y)
        for h in sorted(move.subset):
            working = working.orient(move.y, h)
            between = working.get_edge(move.x, h)
            if between is not None and between.is_undirected:
                working = working.orient(move.x, h)
        return _complete(working)

    # -- the knowledge tail and the face ------------------------------------------------

    @staticmethod
    def _knowledge_tail(graph: MixedGraph, knowledge: DomainKnowledge | None) -> MixedGraph:
        """PC's orientation tail: compelled directions, then the guarded Meek closure."""
        if knowledge is None:
            return graph
        out = graph
        for u, v in sorted(knowledge.required):
            if out.get_edge(u, v) is not None:
                out = out.orient(u, v)
        for edge in out.edge_list():
            if not edge.is_undirected:
                continue
            forward = knowledge.permits_edge(edge.u, edge.v)
            backward = knowledge.permits_edge(edge.v, edge.u)
            if forward and not backward:
                out = out.orient(edge.u, edge.v)
            elif backward and not forward:
                out = out.orient(edge.v, edge.u)

        def forbids(u: str, v: str) -> bool:
            return not knowledge.permits_edge(u, v)

        return meek_closure(out, forbidden=forbids)

    @staticmethod
    def _face(graph: MixedGraph) -> tuple[MixedGraph, float]:
        """The CPDAG face when the graph earns it, the mixed face otherwise."""
        try:
            return graph.as_kind(GraphKind.CPDAG), 1.0
        except InvalidGraphError:
            return graph.as_kind(GraphKind.MIXED), 0.0

fit(data=None, knowledge=None)

Run the forward and backward phases and return the class found.

Source code in src/ergodic/discovery/algorithms/_ges.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Run the forward and backward phases and return the class found."""
    score = resolve_score(self._score, data)
    variables = list(score.variables)
    if not type(score).score_equivalent:
        warnings.warn(
            f"{type(score).__name__} declares score_equivalent=False; GES assumes "
            "score equivalence when it moves between classes, so the search may "
            "not be optimal",
            UserWarning,
            stacklevel=2,
        )
    current = self._start(variables, knowledge)
    n_inserts = 0
    while (move := self._best_insert(score, current, variables, knowledge)) is not None:
        current = self._apply_insert(current, move)
        n_inserts += 1
    n_deletes = 0
    while (move := self._best_delete(score, current, variables, knowledge)) is not None:
        current = self._apply_delete(current, move)
        n_deletes += 1
    final, cpdag_valid = self._face(self._knowledge_tail(current, knowledge))
    if knowledge is not None and not knowledge.is_consistent(final):
        raise ValueError(
            "GES produced a graph that violates the domain knowledge; the data and "
            "the constraints disagree (knowledge.violations(graph) lists how)"
        )
    diagnostics = dict(score.diagnostics())
    diagnostics["n_inserts"] = float(n_inserts)
    diagnostics["n_deletes"] = float(n_deletes)
    diagnostics["n_edges"] = float(len(final.edges))
    diagnostics["n_directed"] = float(sum(1 for e in final.edges if e.is_directed))
    diagnostics["n_undirected"] = float(sum(1 for e in final.edges if e.is_undirected))
    diagnostics["cpdag_valid"] = cpdag_valid
    return DiscoveryResult(
        graph=final,
        method="ges",
        score=score.name,
        alpha=None,
        n=score.n,
        diagnostics=diagnostics,
    )

GRaSP

Greedy relaxations of the sparsest permutation over a decomposable score.

Parameters:

Name	Type	Description	Default
score	Score | str	A :class:~ergodic.discovery.Score instance (bound or unbound) or a name: "auto" (pick by the data schema), "bic", "bic_discrete", "bdeu", "bic_mixed".	'auto'
depth	int	How many levels of tied covered tucks the walk may chain before backtracking; the paper's experiments use 3.	3

fit(data, knowledge) returns a :class:DiscoveryResult whose graph carries the CPDAG face when the final PDAG is a valid CPDAG and the mixed face otherwise (forced knowledge orientations can leave the class; diagnostics["cpdag_valid"] says which happened).

Source code in src/ergodic/discovery/algorithms/_grasp.py

class GRaSP:
    """Greedy relaxations of the sparsest permutation over a decomposable score.

    Args:
        score: A :class:`~ergodic.discovery.Score` instance (bound or
            unbound) or a name: ``"auto"`` (pick by the data schema),
            ``"bic"``, ``"bic_discrete"``, ``"bdeu"``, ``"bic_mixed"``.
        depth: How many levels of tied covered tucks the walk may chain
            before backtracking; the paper's experiments use 3.

    ``fit(data, knowledge)`` returns a :class:`DiscoveryResult` whose graph
    carries the CPDAG face when the final PDAG is a valid CPDAG and the
    ``mixed`` face otherwise (forced knowledge orientations can leave the
    class; ``diagnostics["cpdag_valid"]`` says which happened).
    """

    ingredient: ClassVar[str] = "score"

    def __init__(self, score: Score | str = "auto", *, depth: int = 3) -> None:
        if not isinstance(score, Score | str):
            raise TypeError(f"score must be a Score or a string, got {type(score).__name__}")
        if depth < 1:
            raise ValueError(f"depth must be at least 1, got {depth}")
        self._score = score
        self._depth = depth

    # -- the public entry -------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Walk tucks to convergence and return the class found."""
        score = resolve_score(self._score, data)
        variables = list(score.variables)
        if not type(score).score_equivalent:
            warnings.warn(
                f"{type(score).__name__} declares score_equivalent=False; GRaSP walks "
                "tied covered tucks as if every member of a class scored the same, "
                "so the search may not be optimal",
                UserWarning,
                stacklevel=2,
            )
        req_parents, _ = BOSS._required_maps(variables, knowledge)
        order = _required_consistent(variables, req_parents)
        parents_of, total = _project(score, order, knowledge)
        n_tucks = 0
        n_visited = 0
        while True:
            history = {tuple(order)}
            found = self._dfs(
                score, order, parents_of, total, knowledge, req_parents, self._depth, history
            )
            n_visited += len(history) - 1
            if found is None:
                break
            order, parents_of, total = found
            n_tucks += 1
        edges = [f"{u} -> {v}" for v in variables for u in sorted(parents_of[v])]
        state = cpdag_of_dag(dag(edges, nodes=variables))
        final, cpdag_valid = GES._face(GES._knowledge_tail(state, knowledge))
        if knowledge is not None and not knowledge.is_consistent(final):
            raise ValueError(
                "GRaSP produced a graph that violates the domain knowledge; the data and "
                "the constraints disagree (knowledge.violations(graph) lists how)"
            )
        diagnostics = dict(score.diagnostics())
        diagnostics["n_tucks"] = float(n_tucks)
        diagnostics["n_visited"] = float(n_visited)
        diagnostics["cpdag_valid"] = cpdag_valid
        return DiscoveryResult(
            graph=final,
            method="grasp",
            score=score.name,
            alpha=None,
            n=score.n,
            diagnostics=diagnostics,
        )

    # -- the walk -------------------------------------------------------------------

    def _dfs(
        self,
        score: Score,
        order: list[str],
        parents_of: _ParentMap,
        total: float,
        knowledge: DomainKnowledge | None,
        req_parents: _ParentMap,
        depth_left: int,
        history: set[tuple[str, ...]],
    ) -> tuple[list[str], _ParentMap, float] | None:
        """The first strictly improving permutation reachable by tucks, or ``None``.

        ``total`` is the score of the permutation the walk started from;
        every state on a tied chain shares it. Edges are swept child-first
        in sorted name order, so the walk is deterministic. Non-covered
        tucks are tried only at the first level (``depth_left`` at its
        starting value); tied covered tucks recurse while depth remains.
        """
        at_root = depth_left == self._depth
        for k in sorted(order):
            for j in sorted(parents_of[k]):
                covered = parents_of[j] == parents_of[k] - {j}
                if not covered and not at_root:
                    continue
                candidate = _tuck(order, j, k, parents_of)
                key = tuple(candidate)
                if key in history or not _respects_required(candidate, req_parents):
                    continue  # seen on this walk, or a required parent would trail its child
                history.add(key)
                new_parents, new_total = _project(score, candidate, knowledge)
                if new_total > total + _SLACK:
                    return candidate, new_parents, new_total
                if covered and depth_left > 1 and abs(new_total - total) <= _SLACK:
                    found = self._dfs(
                        score,
                        candidate,
                        new_parents,
                        total,
                        knowledge,
                        req_parents,
                        depth_left - 1,
                        history,
                    )
                    if found is not None:
                        return found
        return None

fit(data=None, knowledge=None)

Walk tucks to convergence and return the class found.

Source code in src/ergodic/discovery/algorithms/_grasp.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Walk tucks to convergence and return the class found."""
    score = resolve_score(self._score, data)
    variables = list(score.variables)
    if not type(score).score_equivalent:
        warnings.warn(
            f"{type(score).__name__} declares score_equivalent=False; GRaSP walks "
            "tied covered tucks as if every member of a class scored the same, "
            "so the search may not be optimal",
            UserWarning,
            stacklevel=2,
        )
    req_parents, _ = BOSS._required_maps(variables, knowledge)
    order = _required_consistent(variables, req_parents)
    parents_of, total = _project(score, order, knowledge)
    n_tucks = 0
    n_visited = 0
    while True:
        history = {tuple(order)}
        found = self._dfs(
            score, order, parents_of, total, knowledge, req_parents, self._depth, history
        )
        n_visited += len(history) - 1
        if found is None:
            break
        order, parents_of, total = found
        n_tucks += 1
    edges = [f"{u} -> {v}" for v in variables for u in sorted(parents_of[v])]
    state = cpdag_of_dag(dag(edges, nodes=variables))
    final, cpdag_valid = GES._face(GES._knowledge_tail(state, knowledge))
    if knowledge is not None and not knowledge.is_consistent(final):
        raise ValueError(
            "GRaSP produced a graph that violates the domain knowledge; the data and "
            "the constraints disagree (knowledge.violations(graph) lists how)"
        )
    diagnostics = dict(score.diagnostics())
    diagnostics["n_tucks"] = float(n_tucks)
    diagnostics["n_visited"] = float(n_visited)
    diagnostics["cpdag_valid"] = cpdag_valid
    return DiscoveryResult(
        graph=final,
        method="grasp",
        score=score.name,
        alpha=None,
        n=score.n,
        diagnostics=diagnostics,
    )

Glasso

The graphical-lasso screen: an undirected skeleton from a sparse precision.

Like MMPC, this is a cheap screen to run before an expensive search: DomainKnowledge().restrict_to(result.graph) turns the recovered support into a constraint that PC, GES, or AStar then searches inside, with far fewer tests or score calls. The honesty note: a precision zero is a conditional independence under Gaussianity, so the graph is a skeleton source, not causal structure; it says which adjacencies the data supports, never which way they point.

The solve is FHT 2008 block coordinate descent on the correlation matrix (the graph, not the scale, is the deliverable), with the diagonal unpenalized (the scikit-learn convention; the paper's differs, see the module docstring) and the shrunk initialization that keeps every column subproblem positive definite.

Parameters:

Name	Type	Description	Default
lam	float | None	The penalty. None (the default) runs the EBIC path: n_lambdas log-spaced values from max |off-diagonal correlation| (the empty graph) down to one percent of it, warm-started descending. Each support met on the path is scored by EBIC at its restricted maximum likelihood estimate (the Foygel-Drton definition; the shrunk path estimate would bias the pick dense as n grows), and the minimizer wins, ties to the sparser, larger lambda. An explicit value runs single-shot; 0.0 asks for the unpenalized MLE and raises when p >= n.	None
gamma	float	The EBIC sparsity weight (Foygel and Drton 2010), 0.5 by default; 0.0 recovers the plain BIC.	0.5
n_lambdas	int	Path length when lam is None.	100
tol	float	Outer convergence, FHT's criterion: the mean absolute change of the working covariance over a full cycle must drop under tol times the mean absolute off-diagonal correlation.	0.001
max_iter	int	Cap on outer cycles per solve; past it the fit raises RuntimeError rather than return unconverged in silence.	100

Knowledge enters through the penalty, and only its adjacency content binds, because the screen is undirected. A pair forbidden in both directions (an adjacency ban, including pairs outside a restrict_to support) is a hard zero, excluded from the updates entirely. A required pair (either direction) is unpenalized, so it survives every lambda the data gives it any conditional association at. Directional knowledge (tiers, one-way forbids, roots, sinks) does not constrain an undirected adjacency and is left to the search that runs inside the support.

fit(data, knowledge) needs a :class:TabularData whose variables are all continuous (a Gaussian model) and returns a :class:DiscoveryResult whose graph holds one tail-tail edge per nonzero precision pair, on the mixed face. The chosen penalty lives in the diagnostics, with the EBIC value when the path ran.

Source code in src/ergodic/discovery/algorithms/_glasso.py

class Glasso:
    """The graphical-lasso screen: an undirected skeleton from a sparse precision.

    Like MMPC, this is a cheap screen to run before an expensive search:
    ``DomainKnowledge().restrict_to(result.graph)`` turns the recovered
    support into a constraint that PC, GES, or AStar then searches inside,
    with far fewer tests or score calls. The honesty note: a precision zero
    is a conditional independence under Gaussianity, so the graph is a
    skeleton source, not causal structure; it says which adjacencies the
    data supports, never which way they point.

    The solve is FHT 2008 block coordinate descent on the correlation matrix
    (the graph, not the scale, is the deliverable), with the diagonal
    unpenalized (the scikit-learn convention; the paper's differs, see the
    module docstring) and the shrunk initialization that keeps every column
    subproblem positive definite.

    Args:
        lam: The penalty. ``None`` (the default) runs the EBIC path:
            ``n_lambdas`` log-spaced values from ``max |off-diagonal
            correlation|`` (the empty graph) down to one percent of it,
            warm-started descending. Each support met on the path is scored
            by EBIC at its restricted maximum likelihood estimate (the
            Foygel-Drton definition; the shrunk path estimate would bias the
            pick dense as ``n`` grows), and the minimizer wins, ties to the
            sparser, larger lambda. An explicit value runs single-shot;
            ``0.0`` asks for the unpenalized MLE and raises when ``p >= n``.
        gamma: The EBIC sparsity weight (Foygel and Drton 2010), ``0.5`` by
            default; ``0.0`` recovers the plain BIC.
        n_lambdas: Path length when ``lam`` is ``None``.
        tol: Outer convergence, FHT's criterion: the mean absolute change of
            the working covariance over a full cycle must drop under ``tol``
            times the mean absolute off-diagonal correlation.
        max_iter: Cap on outer cycles per solve; past it the fit raises
            ``RuntimeError`` rather than return unconverged in silence.

    Knowledge enters through the penalty, and only its adjacency content
    binds, because the screen is undirected. A pair forbidden in both
    directions (an adjacency ban, including pairs outside a ``restrict_to``
    support) is a hard zero, excluded from the updates entirely. A required
    pair (either direction) is unpenalized, so it survives every lambda the
    data gives it any conditional association at. Directional knowledge
    (tiers, one-way forbids, roots, sinks) does not constrain an undirected
    adjacency and is left to the search that runs inside the support.

    ``fit(data, knowledge)`` needs a :class:`TabularData` whose variables are
    all continuous (a Gaussian model) and returns a :class:`DiscoveryResult`
    whose graph holds one tail-tail edge per nonzero precision pair, on the
    ``mixed`` face. The chosen penalty lives in the diagnostics, with the
    EBIC value when the path ran.
    """

    ingredient: ClassVar[str | None] = None

    def __init__(
        self,
        lam: float | None = None,
        *,
        gamma: float = 0.5,
        n_lambdas: int = 100,
        tol: float = 1e-3,
        max_iter: int = 100,
    ) -> None:
        if lam is not None and lam < 0.0:
            raise ValueError(f"lam must be None or >= 0, got {lam}")
        if gamma < 0.0:
            raise ValueError(f"gamma must be >= 0, got {gamma}")
        if n_lambdas < 1:
            raise ValueError(f"n_lambdas must be >= 1, got {n_lambdas}")
        if tol <= 0.0:
            raise ValueError(f"tol must be > 0, got {tol}")
        if max_iter < 1:
            raise ValueError(f"max_iter must be >= 1, got {max_iter}")
        self.lam = lam
        self.gamma = gamma
        self.n_lambdas = n_lambdas
        self.tol = tol
        self.max_iter = max_iter

    # -- the run ---------------------------------------------------------

    def fit(self, data: TabularData, knowledge: DomainKnowledge | None = None) -> DiscoveryResult:
        """Run the screen and return the undirected skeleton with its provenance."""
        if not isinstance(data, TabularData):
            raise TypeError(
                f"glasso reads TabularData, got {type(data).__name__}; the Gaussian model "
                "assumes exchangeable rows, so temporal shapes do not fit here"
            )
        offenders = [name for name in data.variables if not data.schema[name].is_continuous]
        if offenders:
            raise ValueError(
                "glasso fits a Gaussian model, so every variable must be continuous; "
                f"not continuous: {offenders}"
            )
        names = list(data.variables)
        if len(names) < 2:
            raise ValueError(f"glasso needs at least two continuous variables, got {len(names)}")
        complete = data.dropna()
        n = complete.n_rows
        if n < 2:
            raise ValueError(f"glasso needs at least two complete rows, got {n}")
        x = complete.matrix(names)
        constant = [names[i] for i in np.flatnonzero(x.std(axis=0) == 0.0)]
        if constant:
            raise ValueError(f"constant columns have no correlation to screen: {constant}")
        s = np.asarray(np.corrcoef(x, rowvar=False), dtype=np.float64)
        p = len(names)
        free, factor = self._knowledge_arrays(names, knowledge)
        ebic_value: float | None = None
        if self.lam is not None:
            if self.lam == 0.0 and p >= n:
                raise ValueError(
                    f"lam=0 asks for the unpenalized MLE, which does not exist with p={p} "
                    f"variables and n={n} rows; use a positive lam or the EBIC path"
                )
            theta, _, n_iter = _penalized_precision(
                s, self.lam, factor=factor, free=free, tol=self.tol, max_iter=self.max_iter
            )
            lam, n_tried = self.lam, 1
        else:
            theta, lam, ebic_value, n_iter, n_tried = _ebic_path(
                s,
                n=n,
                factor=factor,
                free=free,
                gamma=self.gamma,
                n_lambdas=self.n_lambdas,
                tol=self.tol,
                max_iter=self.max_iter,
            )
        edges = set()
        for i in range(p):
            for j in range(i + 1, p):
                if abs(float(theta[i, j])) > _EDGE_EPS:
                    a, b = sorted((names[i], names[j]))
                    edges.add(Edge.make(a, Mark.TAIL, b, Mark.TAIL))
        graph = MixedGraph(nodes=frozenset(names), edges=frozenset(edges), kind=GraphKind.MIXED)
        diagnostics: dict[str, float] = {
            "lam": float(lam),
            "n_edges": float(len(edges)),
            "n_iterations": float(n_iter),
            "n_lambdas_tried": float(n_tried),
        }
        if ebic_value is not None:
            diagnostics["ebic"] = float(ebic_value)
        return DiscoveryResult(
            graph=graph,
            method="glasso",
            test=None,
            score=None,
            alpha=None,
            n=n,
            diagnostics=diagnostics,
            sepsets=None,
        )

    @staticmethod
    def _knowledge_arrays(
        names: list[str], knowledge: DomainKnowledge | None
    ) -> tuple[npt.NDArray[np.bool_], npt.NDArray[np.float64]]:
        """The per-entry mask and penalty multipliers the knowledge implies.

        ``free`` is False where a pair is banned in both directions (the
        coordinate never updates, so the entry is a hard zero); ``factor``
        is 0 on required pairs (unpenalized) and 1 elsewhere.
        """
        p = len(names)
        free = _default_free(p)
        factor = _default_factor(p)
        if knowledge is None:
            return free, factor
        for i in range(p):
            for j in range(i + 1, p):
                a, b = names[i], names[j]
                if not knowledge.permits_edge(a, b) and not knowledge.permits_edge(b, a):
                    free[i, j] = free[j, i] = False
        index = {name: i for i, name in enumerate(names)}
        for u, v in knowledge.required:
            if u in index and v in index:
                i, j = index[u], index[v]
                factor[i, j] = factor[j, i] = 0.0
        return free, factor

fit(data, knowledge=None)

Run the screen and return the undirected skeleton with its provenance.

Source code in src/ergodic/discovery/algorithms/_glasso.py

def fit(self, data: TabularData, knowledge: DomainKnowledge | None = None) -> DiscoveryResult:
    """Run the screen and return the undirected skeleton with its provenance."""
    if not isinstance(data, TabularData):
        raise TypeError(
            f"glasso reads TabularData, got {type(data).__name__}; the Gaussian model "
            "assumes exchangeable rows, so temporal shapes do not fit here"
        )
    offenders = [name for name in data.variables if not data.schema[name].is_continuous]
    if offenders:
        raise ValueError(
            "glasso fits a Gaussian model, so every variable must be continuous; "
            f"not continuous: {offenders}"
        )
    names = list(data.variables)
    if len(names) < 2:
        raise ValueError(f"glasso needs at least two continuous variables, got {len(names)}")
    complete = data.dropna()
    n = complete.n_rows
    if n < 2:
        raise ValueError(f"glasso needs at least two complete rows, got {n}")
    x = complete.matrix(names)
    constant = [names[i] for i in np.flatnonzero(x.std(axis=0) == 0.0)]
    if constant:
        raise ValueError(f"constant columns have no correlation to screen: {constant}")
    s = np.asarray(np.corrcoef(x, rowvar=False), dtype=np.float64)
    p = len(names)
    free, factor = self._knowledge_arrays(names, knowledge)
    ebic_value: float | None = None
    if self.lam is not None:
        if self.lam == 0.0 and p >= n:
            raise ValueError(
                f"lam=0 asks for the unpenalized MLE, which does not exist with p={p} "
                f"variables and n={n} rows; use a positive lam or the EBIC path"
            )
        theta, _, n_iter = _penalized_precision(
            s, self.lam, factor=factor, free=free, tol=self.tol, max_iter=self.max_iter
        )
        lam, n_tried = self.lam, 1
    else:
        theta, lam, ebic_value, n_iter, n_tried = _ebic_path(
            s,
            n=n,
            factor=factor,
            free=free,
            gamma=self.gamma,
            n_lambdas=self.n_lambdas,
            tol=self.tol,
            max_iter=self.max_iter,
        )
    edges = set()
    for i in range(p):
        for j in range(i + 1, p):
            if abs(float(theta[i, j])) > _EDGE_EPS:
                a, b = sorted((names[i], names[j]))
                edges.add(Edge.make(a, Mark.TAIL, b, Mark.TAIL))
    graph = MixedGraph(nodes=frozenset(names), edges=frozenset(edges), kind=GraphKind.MIXED)
    diagnostics: dict[str, float] = {
        "lam": float(lam),
        "n_edges": float(len(edges)),
        "n_iterations": float(n_iter),
        "n_lambdas_tried": float(n_tried),
    }
    if ebic_value is not None:
        diagnostics["ebic"] = float(ebic_value)
    return DiscoveryResult(
        graph=graph,
        method="glasso",
        test=None,
        score=None,
        alpha=None,
        n=n,
        diagnostics=diagnostics,
        sepsets=None,
    )

GrangerResult dataclass

Granger tests for every ordered pair: a predictive table, not a graph.

pvalues maps (cause, effect) to the likelihood-ratio p-value; significant holds the pairs with p <= alpha. n counts the stacked window rows the regressions used.

Source code in src/ergodic/discovery/_granger.py

@dataclass(frozen=True)
class GrangerResult:
    """Granger tests for every ordered pair: a predictive table, not a graph.

    ``pvalues`` maps ``(cause, effect)`` to the likelihood-ratio p-value;
    ``significant`` holds the pairs with ``p <= alpha``. ``n`` counts the
    stacked window rows the regressions used.
    """

    pvalues: Mapping[tuple[str, str], float]
    significant: frozenset[tuple[str, str]]
    max_lag: int
    alpha: float
    n: int

HillClimb

Greedy DAG search: add, remove, and reverse moves with a tabu ring.

Steepest ascent from the empty (or knowledge-seeded) graph, returning the best state seen. n_iterations in the diagnostics counts applied moves, so it always equals n_adds + n_removes + n_reverses.

Parameters:

Name	Type	Description	Default
score	Score | str	A :class:~ergodic.discovery.Score instance (bound or unbound) or a name: "auto" (pick by the data schema), "bic", "bic_discrete", "bdeu", "bic_mixed".	'auto'
tabu	int	The length of the tabu ring (inverses of recent moves that may not be replayed) and the bound on consecutive non-improving escape steps; 0 disables both, so the search stops at the first local optimum.	10
max_iter	int	A hard cap on the number of applied moves.	10000
max_parents	int | None	An in-degree cap enforced on adds and reversals; None leaves the in-degree free. Required edges are seeded whatever the cap.	None

Source code in src/ergodic/discovery/algorithms/_hc.py

class HillClimb:
    """Greedy DAG search: add, remove, and reverse moves with a tabu ring.

    Steepest ascent from the empty (or knowledge-seeded) graph, returning
    the best state seen. ``n_iterations`` in the diagnostics counts applied
    moves, so it always equals ``n_adds + n_removes + n_reverses``.

    Args:
        score: A :class:`~ergodic.discovery.Score` instance (bound or
            unbound) or a name: ``"auto"`` (pick by the data schema),
            ``"bic"``, ``"bic_discrete"``, ``"bdeu"``, ``"bic_mixed"``.
        tabu: The length of the tabu ring (inverses of recent moves that may
            not be replayed) and the bound on consecutive non-improving
            escape steps; ``0`` disables both, so the search stops at the
            first local optimum.
        max_iter: A hard cap on the number of applied moves.
        max_parents: An in-degree cap enforced on adds and reversals;
            ``None`` leaves the in-degree free. Required edges are seeded
            whatever the cap.
    """

    ingredient: ClassVar[str] = "score"

    def __init__(
        self,
        score: Score | str = "auto",
        *,
        tabu: int = 10,
        max_iter: int = 10_000,
        max_parents: int | None = None,
    ) -> None:
        if not isinstance(score, Score | str):
            raise TypeError(f"score must be a Score or a string, got {type(score).__name__}")
        if tabu < 0:
            raise ValueError(f"tabu must be non-negative, got {tabu}")
        if max_iter < 1:
            raise ValueError(f"max_iter must be at least 1, got {max_iter}")
        if max_parents is not None and max_parents < 0:
            raise ValueError(f"max_parents must be non-negative, got {max_parents}")
        self._score = score
        self._tabu = tabu
        self._max_iter = max_iter
        self._max_parents = max_parents

    # -- the public entry -------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Climb to a high-scoring DAG and return it with its provenance."""
        score = resolve_score(self._score, data)
        variables = sorted(score.variables)
        parents_of: dict[str, set[str]] = {v: set() for v in variables}
        required: frozenset[tuple[str, str]] = frozenset()
        if knowledge is not None:
            known = set(variables)
            required = frozenset(
                pair for pair in knowledge.required if pair[0] in known and pair[1] in known
            )
            for u, v in required:
                parents_of[v].add(u)
            if _has_cycle(parents_of):
                raise ValueError(
                    "the required edges form a cycle, so no DAG can satisfy the "
                    "knowledge; drop one of the required edges"
                )
        total = float(sum(score.local(node, parents_of[node]) for node in variables))
        best_total = total
        best = {node: frozenset(ps) for node, ps in parents_of.items()}
        ring: deque[_Move] | None = deque(maxlen=self._tabu) if self._tabu > 0 else None
        iterations = 0
        stalls = 0
        applied = {"add": 0, "remove": 0, "reverse": 0}
        while iterations < self._max_iter:
            found = self._best_move(score, variables, parents_of, required, knowledge, ring)
            if found is None:
                break
            move, delta = found
            if delta > _EPS:
                stalls = 0
            elif ring is not None and stalls < self._tabu:
                stalls += 1  # a bounded escape: take the least bad move anyway
            else:
                break
            kind, u, v = move
            if kind == "add":
                parents_of[v].add(u)
            elif kind == "remove":
                parents_of[v].remove(u)
            else:
                parents_of[v].remove(u)
                parents_of[u].add(v)
            total += delta
            iterations += 1
            applied[kind] += 1
            if ring is not None:
                ring.append(_inverse(move))
            if total > best_total:
                best_total = total
                best = {node: frozenset(ps) for node, ps in parents_of.items()}
        edges = [f"{u} -> {v}" for v in variables for u in sorted(best[v])]
        graph = dag(edges, nodes=variables)
        if knowledge is not None:
            problems = knowledge.violations(graph)
            if problems:
                raise ValueError(
                    "the climbed DAG violates the domain knowledge; constraints a DAG "
                    "cannot express (confounded pairs) or required edges over variables "
                    "missing from the data can cause this: " + "; ".join(problems)
                )
        diagnostics = dict(score.diagnostics())
        diagnostics["n_iterations"] = float(iterations)
        diagnostics["n_adds"] = float(applied["add"])
        diagnostics["n_removes"] = float(applied["remove"])
        diagnostics["n_reverses"] = float(applied["reverse"])
        diagnostics["n_edges"] = float(sum(len(ps) for ps in best.values()))
        return DiscoveryResult(
            graph=graph,
            method="hc",
            score=score.name,
            alpha=None,
            n=score.n,
            diagnostics=diagnostics,
        )

    # -- move selection -----------------------------------------------------------

    def _best_move(
        self,
        score: Score,
        variables: list[str],
        parents_of: dict[str, set[str]],
        required: frozenset[tuple[str, str]],
        knowledge: DomainKnowledge | None,
        ring: deque[_Move] | None,
    ) -> tuple[_Move, float] | None:
        """The best admissible move off the tabu ring, or ``None`` when nothing is legal.

        Moves are enumerated in a fixed order (adds, then removals, then
        reversals, each over sorted node pairs) and a tie keeps the first,
        so a run is deterministic for a given score and knowledge. Required
        pairs are never removed or reversed, adds and reversals respect
        ``permits_edge``, the in-degree cap, and acyclicity by reachability.
        """
        children: dict[str, set[str]] = {node: set() for node in variables}
        for node, ps in parents_of.items():
            for parent in ps:
                children[parent].add(node)
        cap = self._max_parents
        best: tuple[_Move, float] | None = None

        def consider(move: _Move, delta: float) -> None:
            nonlocal best
            if ring is not None and move in ring:
                return
            if best is None or delta > best[1]:
                best = (move, delta)

        for u in variables:
            for v in variables:
                if (
                    u != v
                    and u not in parents_of[v]
                    and v not in parents_of[u]
                    and (knowledge is None or knowledge.permits_edge(u, v))
                    and (cap is None or len(parents_of[v]) < cap)
                    and not _reaches(children, v, u)
                ):
                    consider(("add", u, v), score.local_diff(v, parents_of[v], u))
        for u in variables:
            for v in variables:
                if (
                    u != v
                    and u in parents_of[v]
                    and (u, v) not in required
                    and (v, u) not in required
                ):
                    removal = score.local(v, parents_of[v] - {u}) - score.local(v, parents_of[v])
                    consider(("remove", u, v), removal)
        for u in variables:
            for v in variables:
                if (
                    u != v
                    and u in parents_of[v]
                    and (u, v) not in required
                    and (v, u) not in required
                    and (knowledge is None or knowledge.permits_edge(v, u))
                    and (cap is None or len(parents_of[u]) < cap)
                    and not _reaches(children, u, v, drop=(u, v))
                ):
                    removal = score.local(v, parents_of[v] - {u}) - score.local(v, parents_of[v])
                    gain = removal + score.local_diff(u, parents_of[u], v)
                    consider(("reverse", u, v), gain)
        return best

fit(data=None, knowledge=None)

Climb to a high-scoring DAG and return it with its provenance.

Source code in src/ergodic/discovery/algorithms/_hc.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Climb to a high-scoring DAG and return it with its provenance."""
    score = resolve_score(self._score, data)
    variables = sorted(score.variables)
    parents_of: dict[str, set[str]] = {v: set() for v in variables}
    required: frozenset[tuple[str, str]] = frozenset()
    if knowledge is not None:
        known = set(variables)
        required = frozenset(
            pair for pair in knowledge.required if pair[0] in known and pair[1] in known
        )
        for u, v in required:
            parents_of[v].add(u)
        if _has_cycle(parents_of):
            raise ValueError(
                "the required edges form a cycle, so no DAG can satisfy the "
                "knowledge; drop one of the required edges"
            )
    total = float(sum(score.local(node, parents_of[node]) for node in variables))
    best_total = total
    best = {node: frozenset(ps) for node, ps in parents_of.items()}
    ring: deque[_Move] | None = deque(maxlen=self._tabu) if self._tabu > 0 else None
    iterations = 0
    stalls = 0
    applied = {"add": 0, "remove": 0, "reverse": 0}
    while iterations < self._max_iter:
        found = self._best_move(score, variables, parents_of, required, knowledge, ring)
        if found is None:
            break
        move, delta = found
        if delta > _EPS:
            stalls = 0
        elif ring is not None and stalls < self._tabu:
            stalls += 1  # a bounded escape: take the least bad move anyway
        else:
            break
        kind, u, v = move
        if kind == "add":
            parents_of[v].add(u)
        elif kind == "remove":
            parents_of[v].remove(u)
        else:
            parents_of[v].remove(u)
            parents_of[u].add(v)
        total += delta
        iterations += 1
        applied[kind] += 1
        if ring is not None:
            ring.append(_inverse(move))
        if total > best_total:
            best_total = total
            best = {node: frozenset(ps) for node, ps in parents_of.items()}
    edges = [f"{u} -> {v}" for v in variables for u in sorted(best[v])]
    graph = dag(edges, nodes=variables)
    if knowledge is not None:
        problems = knowledge.violations(graph)
        if problems:
            raise ValueError(
                "the climbed DAG violates the domain knowledge; constraints a DAG "
                "cannot express (confounded pairs) or required edges over variables "
                "missing from the data can cause this: " + "; ".join(problems)
            )
    diagnostics = dict(score.diagnostics())
    diagnostics["n_iterations"] = float(iterations)
    diagnostics["n_adds"] = float(applied["add"])
    diagnostics["n_removes"] = float(applied["remove"])
    diagnostics["n_reverses"] = float(applied["reverse"])
    diagnostics["n_edges"] = float(sum(len(ps) for ps in best.values()))
    return DiscoveryResult(
        graph=graph,
        method="hc",
        score=score.name,
        alpha=None,
        n=score.n,
        diagnostics=diagnostics,
    )

Linear

Bases: CITest

Partial correlation through the Fisher z transform.

Residualizes x and y on [1, z], computes the correlation r of the residuals, and reports sqrt(n - dz - 3) * atanh(r) against the standard normal, two-sided. The sign of r survives in the statistic; the p-value does not depend on it.

Source code in src/ergodic/discovery/ci/_linear.py

class Linear(CITest):
    """Partial correlation through the Fisher z transform.

    Residualizes x and y on ``[1, z]``, computes the correlation ``r`` of the
    residuals, and reports ``sqrt(n - dz - 3) * atanh(r)`` against the
    standard normal, two-sided. The sign of ``r`` survives in the statistic;
    the p-value does not depend on it.
    """

    name: ClassVar[str] = "linear"
    needs: ClassVar[str] = "numeric"

    def statistic(
        self,
        x: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        z: npt.NDArray[np.float64],
    ) -> tuple[float, float, float | None]:
        n, dz = x.shape[0], z.shape[1]
        if x.shape[1] != 1 or y.shape[1] != 1:
            raise TypeError(
                "Linear needs scalar endpoints; a categorical endpoint expands to a "
                "one-hot block, use Mixed for categorical endpoints"
            )
        if n - dz - 3 <= 0:
            raise ValueError(
                "too few rows for the conditioning size: the Fisher z transform needs "
                f"n - dz - 3 > 0, got n={n} with dz={dz} conditioning columns"
            )
        rx = _residualize(x, z)[:, 0]
        ry = _residualize(y, z)[:, 0]
        denom = math.sqrt(float(rx @ rx) * float(ry @ ry))
        r = float(rx @ ry) / denom if denom > 0.0 else 0.0
        r = min(_CLIP, max(-_CLIP, r))
        stat = math.sqrt(n - dz - 3) * math.atanh(r)
        p = 2.0 * (1.0 - _norm_cdf(abs(stat)))
        return stat, p, None

MMPC

The max-min parents-and-children screen over a conditional-independence test.

MMPC is a cheap skeleton screen to run before an expensive search: it returns the undirected adjacencies the data supports, and DomainKnowledge().restrict_to(result.graph) turns them into a support that PC, GES, or AStar then searches inside, with far fewer tests or score calls. Adjacencies are kept under the AND rule (each endpoint must sit in the other's parents-children set), the stricter of the two symmetry choices, so the screen errs toward sparsity.

Parameters:

Name	Type	Description	Default
test	CITest | str	A :class:CITest instance (bound or unbound, including :class:DSeparationOracle), or a name: "auto" (pick by the data schema), "linear"/"fisher_z", "discrete"/ "g_sq", "mixed"/"dg", "gcm".	'auto'
alpha	float	Significance level for the independence decisions.	0.05
max_cond_size	int | None	Cap on the conditioning-set size, None for no cap. Subsets of a growing CPC multiply combinatorially, so the default keeps the standard practical cap of 3.	3

Knowledge enters in three modes. Prune: a pair forbidden in both directions is never a candidate. Protect: required and confounded pairs enter every CPC untested, survive the backward sweep, and pass the AND rule unconditionally. Validate: the finished skeleton is checked at the adjacency level only, because a skeleton has no directions, so is_consistent (which demands required edges be directed) cannot apply; orientation constraints (tiers, single-direction forbids, roots, sinks) do not constrain an undirected adjacency and are left to the search that runs inside the support.

fit(data, knowledge) returns a :class:DiscoveryResult whose graph holds one tail-tail edge per surviving pair, with the mixed face.

Source code in src/ergodic/discovery/algorithms/_mmpc.py

class MMPC:
    """The max-min parents-and-children screen over a conditional-independence test.

    MMPC is a cheap skeleton screen to run before an expensive search: it
    returns the undirected adjacencies the data supports, and
    ``DomainKnowledge().restrict_to(result.graph)`` turns them into a support
    that PC, GES, or AStar then searches inside, with far fewer tests or
    score calls. Adjacencies are kept under the AND rule (each endpoint must
    sit in the other's parents-children set), the stricter of the two
    symmetry choices, so the screen errs toward sparsity.

    Args:
        test: A :class:`CITest` instance (bound or unbound, including
            :class:`DSeparationOracle`), or a name: ``"auto"`` (pick by the
            data schema), ``"linear"``/``"fisher_z"``, ``"discrete"``/
            ``"g_sq"``, ``"mixed"``/``"dg"``, ``"gcm"``.
        alpha: Significance level for the independence decisions.
        max_cond_size: Cap on the conditioning-set size, ``None`` for no
            cap. Subsets of a growing CPC multiply combinatorially, so the
            default keeps the standard practical cap of 3.

    Knowledge enters in three modes. Prune: a pair forbidden in both
    directions is never a candidate. Protect: required and confounded pairs
    enter every CPC untested, survive the backward sweep, and pass the AND
    rule unconditionally. Validate: the finished skeleton is checked at the
    adjacency level only, because a skeleton has no directions, so
    ``is_consistent`` (which demands required edges be directed) cannot
    apply; orientation constraints (tiers, single-direction forbids, roots,
    sinks) do not constrain an undirected adjacency and are left to the
    search that runs inside the support.

    ``fit(data, knowledge)`` returns a :class:`DiscoveryResult` whose graph
    holds one tail-tail edge per surviving pair, with the ``mixed`` face.
    """

    ingredient: ClassVar[str] = "test"

    def __init__(
        self,
        test: CITest | str = "auto",
        *,
        alpha: float = 0.05,
        max_cond_size: int | None = 3,
    ) -> None:
        if not 0.0 < alpha < 1.0:
            raise ValueError(f"alpha must be in (0, 1), got {alpha}")
        if max_cond_size is not None and max_cond_size < 0:
            raise ValueError(f"max_cond_size must be None or >= 0, got {max_cond_size}")
        self.test = test
        self.alpha = alpha
        self.max_cond_size = max_cond_size

    # -- the run ---------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Run MMPC and return the undirected skeleton with its provenance."""
        oracle = resolve_test(self.test, data)
        nodes = sorted(oracle.variables)
        node_set = frozenset(nodes)
        permits = _permits_all if knowledge is None else knowledge.permits_edge
        protected = _protected_pairs(knowledge)
        cpc: dict[str, set[str]] = {}
        n_forward_added = 0
        n_backward_removed = 0
        for target in nodes:
            seeded = (
                frozenset(
                    partner for pair in protected if target in pair for partner in pair - {target}
                )
                & node_set
            )
            candidates = [
                x
                for x in nodes
                if x != target and x not in seeded and (permits(x, target) or permits(target, x))
            ]
            members, added, removed = _parents_children(
                oracle, target, candidates, seeded, alpha=self.alpha, cap=self.max_cond_size
            )
            cpc[target] = members
            n_forward_added += added
            n_backward_removed += removed
        pairs = _and_rule(nodes, cpc, protected)
        edges = set()
        for pair in sorted(pairs, key=sorted):
            a, b = sorted(pair)
            edges.add(Edge.make(a, Mark.TAIL, b, Mark.TAIL))
        graph = MixedGraph(nodes=node_set, edges=frozenset(edges), kind=GraphKind.MIXED)
        if knowledge is not None:
            _validate_skeleton(graph, knowledge)
        diagnostics: dict[str, float] = dict(oracle.diagnostics())
        diagnostics["n_edges"] = float(len(edges))
        diagnostics["n_forward_added"] = float(n_forward_added)
        diagnostics["n_backward_removed"] = float(n_backward_removed)
        return DiscoveryResult(
            graph=graph,
            method="mmpc",
            test=oracle.name,
            alpha=self.alpha,
            n=oracle.n,
            diagnostics=diagnostics,
            sepsets=None,
        )

fit(data=None, knowledge=None)

Run MMPC and return the undirected skeleton with its provenance.

Source code in src/ergodic/discovery/algorithms/_mmpc.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Run MMPC and return the undirected skeleton with its provenance."""
    oracle = resolve_test(self.test, data)
    nodes = sorted(oracle.variables)
    node_set = frozenset(nodes)
    permits = _permits_all if knowledge is None else knowledge.permits_edge
    protected = _protected_pairs(knowledge)
    cpc: dict[str, set[str]] = {}
    n_forward_added = 0
    n_backward_removed = 0
    for target in nodes:
        seeded = (
            frozenset(
                partner for pair in protected if target in pair for partner in pair - {target}
            )
            & node_set
        )
        candidates = [
            x
            for x in nodes
            if x != target and x not in seeded and (permits(x, target) or permits(target, x))
        ]
        members, added, removed = _parents_children(
            oracle, target, candidates, seeded, alpha=self.alpha, cap=self.max_cond_size
        )
        cpc[target] = members
        n_forward_added += added
        n_backward_removed += removed
    pairs = _and_rule(nodes, cpc, protected)
    edges = set()
    for pair in sorted(pairs, key=sorted):
        a, b = sorted(pair)
        edges.add(Edge.make(a, Mark.TAIL, b, Mark.TAIL))
    graph = MixedGraph(nodes=node_set, edges=frozenset(edges), kind=GraphKind.MIXED)
    if knowledge is not None:
        _validate_skeleton(graph, knowledge)
    diagnostics: dict[str, float] = dict(oracle.diagnostics())
    diagnostics["n_edges"] = float(len(edges))
    diagnostics["n_forward_added"] = float(n_forward_added)
    diagnostics["n_backward_removed"] = float(n_backward_removed)
    return DiscoveryResult(
        graph=graph,
        method="mmpc",
        test=oracle.name,
        alpha=self.alpha,
        n=oracle.n,
        diagnostics=diagnostics,
        sepsets=None,
    )

Mixed

Bases: CITest

Partial canonical correlations with Bartlett's chi-square.

Residualizes the x and y blocks on [1, z], takes the reduced QR of each residual block, and reads the canonical correlations rho_i off the singular values of Qx.T @ Qy (clipped below 1). Bartlett's statistic -(n - dz - (dx + dy + 3) / 2) * sum(log(1 - rho_i^2)) is referred to a chi-square with dx * dy degrees of freedom.

With scalar endpoints (dx == dy == 1) the single canonical correlation is the absolute partial correlation, so Mixed agrees with Linear asymptotically. A zero-width block (a single-level categorical) carries no information; that query is vacuous and returns (0.0, 1.0, 0.0).

Source code in src/ergodic/discovery/ci/_mixed.py

class Mixed(CITest):
    """Partial canonical correlations with Bartlett's chi-square.

    Residualizes the x and y blocks on ``[1, z]``, takes the reduced QR of
    each residual block, and reads the canonical correlations ``rho_i`` off
    the singular values of ``Qx.T @ Qy`` (clipped below 1). Bartlett's
    statistic ``-(n - dz - (dx + dy + 3) / 2) * sum(log(1 - rho_i^2))`` is
    referred to a chi-square with ``dx * dy`` degrees of freedom.

    With scalar endpoints (``dx == dy == 1``) the single canonical correlation
    is the absolute partial correlation, so Mixed agrees with Linear
    asymptotically. A zero-width block (a single-level categorical) carries no
    information; that query is vacuous and returns ``(0.0, 1.0, 0.0)``.
    """

    name: ClassVar[str] = "mixed"
    needs: ClassVar[str] = "numeric"

    def statistic(
        self,
        x: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        z: npt.NDArray[np.float64],
    ) -> tuple[float, float, float | None]:
        n = x.shape[0]
        dx, dy, dz = x.shape[1], y.shape[1], z.shape[1]
        if dx == 0 or dy == 0:
            return (0.0, 1.0, 0.0)
        if n - dz - (dx + dy) - 1 <= 0:
            raise ValueError(
                "too few rows for the block sizes: the Bartlett statistic needs "
                f"n - dz - (dx + dy) - 1 > 0, got n={n} with dx={dx}, dy={dy}, dz={dz}"
            )
        qx, _ = np.linalg.qr(_residualize(x, z), mode="reduced")
        qy, _ = np.linalg.qr(_residualize(y, z), mode="reduced")
        rho = np.clip(np.linalg.svd(qx.T @ qy, compute_uv=False), 0.0, _CLIP)
        stat = -(n - dz - (dx + dy + 3) / 2.0) * float(np.sum(np.log1p(-(rho**2))))
        dof = float(dx * dy)
        return stat, _chi2_sf(stat, dof), dof

MixedBIC

Bases: BIC

The degenerate-Gaussian BIC, decomposable and score-equivalent.

Any column kind is a valid node: a continuous node is its single column and a categorical node is its one-hot block (first level dropped), scored jointly. local = -(n/2) log det Sigma_{Y|X} - penalty * p/2 * log n with p = dy dx + dy (dy + 1) / 2 for a dy-column node and dx total parent columns.

Parameters:

Name	Type	Description	Default
data	TabularData | None	Bind at construction, or later through bind.	None
penalty	float	A multiplier on the complexity term, inherited from BIC. 1.0 is the textbook BIC; smaller values prune less.	1.0

Source code in src/ergodic/discovery/scores/_mixed_bic.py

class MixedBIC(BIC):
    """The degenerate-Gaussian BIC, decomposable and score-equivalent.

    Any column kind is a valid node: a continuous node is its single column
    and a categorical node is its one-hot block (first level dropped), scored
    jointly. ``local = -(n/2) log det Sigma_{Y|X} - penalty * p/2 * log n``
    with ``p = dy dx + dy (dy + 1) / 2`` for a ``dy``-column node and ``dx``
    total parent columns.

    Args:
        data: Bind at construction, or later through ``bind``.
        penalty: A multiplier on the complexity term, inherited from ``BIC``.
            ``1.0`` is the textbook BIC; smaller values prune less.
    """

    name: ClassVar[str] = "bic_mixed"

    # -- the covariance path ----------------------------------------------------

    def _local_value(self, node: str, parents: frozenset[str]) -> float:
        """The local by block Schur complement on the precomputed covariance."""
        node_slice = self._slices[node]
        y_columns = list(range(node_slice.start, node_slice.stop))
        x_columns = [
            column
            for parent in sorted(parents)
            for column in range(self._slices[parent].start, self._slices[parent].stop)
        ]
        s_yy = self._cov[np.ix_(y_columns, y_columns)]
        if x_columns:
            s_xx = self._cov[np.ix_(x_columns, x_columns)]
            s_xy = self._cov[np.ix_(x_columns, y_columns)]
            conditional = s_yy - s_xy.T @ np.linalg.pinv(s_xx, hermitian=True) @ s_xy
        else:
            conditional = s_yy
        return self._mixed_value(
            np.asarray(conditional, dtype=np.float64),
            dy=len(y_columns),
            dx=len(x_columns),
            n=self._n,
        )

    # -- the array path (user contract, and the cross-check) ---------------------

    def local_from_arrays(self, y: npt.NDArray[np.float64], X: npt.NDArray[np.float64]) -> float:
        """The same local by direct least squares on ``(n, dy)`` ``y`` and ``(n, k)`` ``X``.

        Agrees with the covariance path to numerical precision; the tests hold
        the two to ``1e-8``.
        """
        if y.ndim != 2 or y.shape[1] < 1:
            raise TypeError("MixedBIC scores column blocks; y must have shape (n, dy), dy >= 1")
        n = int(y.shape[0])
        if n == 0:
            raise ValueError("MixedBIC needs at least one row")
        centered_y = y - y.mean(axis=0, keepdims=True)
        if X.shape[1]:
            centered_x = X - X.mean(axis=0, keepdims=True)
            beta, *_ = np.linalg.lstsq(centered_x, centered_y, rcond=None)
            residual = centered_y - centered_x @ beta
        else:
            residual = centered_y
        conditional = np.asarray(residual.T @ residual, dtype=np.float64) / float(n)
        return self._mixed_value(conditional, dy=int(y.shape[1]), dx=int(X.shape[1]), n=n)

    # -- the shared tail -----------------------------------------------------------

    def _mixed_value(
        self, conditional: npt.NDArray[np.float64], *, dy: int, dx: int, n: int
    ) -> float:
        eigenvalues = np.linalg.eigvalsh(conditional)
        logdet = float(np.log(np.maximum(eigenvalues, _EIGENVALUE_FLOOR)).sum())
        size = float(n)
        fit = -(size / 2.0) * logdet
        parameters = dy * dx + dy * (dy + 1) / 2.0
        complexity = self._penalty * (parameters / 2.0) * math.log(size)
        return fit - complexity

local_from_arrays(y, X)

The same local by direct least squares on (n, dy) y and (n, k) X.

Agrees with the covariance path to numerical precision; the tests hold the two to 1e-8.

Source code in src/ergodic/discovery/scores/_mixed_bic.py

def local_from_arrays(self, y: npt.NDArray[np.float64], X: npt.NDArray[np.float64]) -> float:
    """The same local by direct least squares on ``(n, dy)`` ``y`` and ``(n, k)`` ``X``.

    Agrees with the covariance path to numerical precision; the tests hold
    the two to ``1e-8``.
    """
    if y.ndim != 2 or y.shape[1] < 1:
        raise TypeError("MixedBIC scores column blocks; y must have shape (n, dy), dy >= 1")
    n = int(y.shape[0])
    if n == 0:
        raise ValueError("MixedBIC needs at least one row")
    centered_y = y - y.mean(axis=0, keepdims=True)
    if X.shape[1]:
        centered_x = X - X.mean(axis=0, keepdims=True)
        beta, *_ = np.linalg.lstsq(centered_x, centered_y, rcond=None)
        residual = centered_y - centered_x @ beta
    else:
        residual = centered_y
    conditional = np.asarray(residual.T @ residual, dtype=np.float64) / float(n)
    return self._mixed_value(conditional, dy=int(y.shape[1]), dx=int(X.shape[1]), n=n)

PC

PC-stable over a conditional-independence test.

Parameters:

Name	Type	Description	Default
test	CITest | str	A :class:CITest instance (bound or unbound, including :class:DSeparationOracle), or a name: "auto" (pick by the data schema), "linear"/"fisher_z", "discrete"/ "g_sq", "mixed"/"dg", "gcm".	'auto'
alpha	float	Significance level for the independence decisions.	0.05
max_cond_size	int | None	Cap on the conditioning-set size, None for no cap.	None
colliders	str	The v-structure rule: "standard" reads the recorded sepsets, "conservative" and "majority" re-test every unshielded triple and leave the ambiguous ones unoriented (see the module docstring).	'standard'

fit(data, knowledge) returns a :class:DiscoveryResult whose graph carries the CPDAG face when the oriented PDAG is a valid CPDAG and the mixed face otherwise.

Source code in src/ergodic/discovery/algorithms/_pc.py

class PC:
    """PC-stable over a conditional-independence test.

    Args:
        test: A :class:`CITest` instance (bound or unbound, including
            :class:`DSeparationOracle`), or a name: ``"auto"`` (pick by the
            data schema), ``"linear"``/``"fisher_z"``, ``"discrete"``/
            ``"g_sq"``, ``"mixed"``/``"dg"``, ``"gcm"``.
        alpha: Significance level for the independence decisions.
        max_cond_size: Cap on the conditioning-set size, ``None`` for no cap.
        colliders: The v-structure rule: ``"standard"`` reads the recorded
            sepsets, ``"conservative"`` and ``"majority"`` re-test every
            unshielded triple and leave the ambiguous ones unoriented (see
            the module docstring).

    ``fit(data, knowledge)`` returns a :class:`DiscoveryResult` whose graph
    carries the CPDAG face when the oriented PDAG is a valid CPDAG and the
    ``mixed`` face otherwise.
    """

    ingredient: ClassVar[str] = "test"

    def __init__(
        self,
        test: CITest | str = "auto",
        *,
        alpha: float = 0.05,
        max_cond_size: int | None = None,
        colliders: str = "standard",
    ) -> None:
        if not 0.0 < alpha < 1.0:
            raise ValueError(f"alpha must be in (0, 1), got {alpha}")
        if max_cond_size is not None and max_cond_size < 0:
            raise ValueError(f"max_cond_size must be None or >= 0, got {max_cond_size}")
        if colliders not in ("standard", "conservative", "majority"):
            raise ValueError(
                f"colliders must be 'standard', 'conservative', or 'majority', got {colliders!r}"
            )
        self.test = test
        self.alpha = alpha
        self.max_cond_size = max_cond_size
        self.colliders = colliders

    # -- the run ---------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Run PC-stable and return the discovered graph with its provenance."""
        oracle = resolve_test(self.test, data)
        nodes = sorted(oracle.variables)
        adj, sepsets = stable_skeleton(
            oracle, nodes, knowledge, alpha=self.alpha, max_cond_size=self.max_cond_size
        )
        collider_triples: frozenset[tuple[str, str, str]] | None = None
        ambiguous_triples: frozenset[tuple[str, str, str]] = frozenset()
        if self.colliders != "standard":
            collider_triples, ambiguous_triples = _vote_triples(
                oracle,
                nodes,
                adj,
                sepsets,
                rule=self.colliders,
                alpha=self.alpha,
                max_cond_size=self.max_cond_size,
            )
        graph, n_v_structures = _orient(
            nodes, adj, sepsets, knowledge, colliders=collider_triples, ambiguous=ambiguous_triples
        )
        try:
            final = graph.as_kind(GraphKind.CPDAG)
            cpdag_valid = 1.0
        except InvalidGraphError:
            final = graph.as_kind(GraphKind.MIXED)
            cpdag_valid = 0.0
        if knowledge is not None and not knowledge.is_consistent(final):
            raise ValueError(
                "PC produced a graph that violates the domain knowledge; the data and "
                "the constraints disagree (knowledge.violations(graph) lists how)"
            )
        diagnostics: dict[str, float] = dict(oracle.diagnostics())
        diagnostics["n_edges"] = float(len(final.edges))
        diagnostics["n_directed"] = float(sum(1 for e in final.edges if e.is_directed))
        diagnostics["n_undirected"] = float(sum(1 for e in final.edges if e.is_undirected))
        diagnostics["n_v_structures"] = float(n_v_structures)
        diagnostics["n_ambiguous_triples"] = float(len(ambiguous_triples))
        diagnostics["cpdag_valid"] = cpdag_valid
        return DiscoveryResult(
            graph=final,
            method="pc",
            test=oracle.name,
            alpha=self.alpha,
            n=oracle.n,
            diagnostics=diagnostics,
            sepsets=sepsets,
        )

fit(data=None, knowledge=None)

Run PC-stable and return the discovered graph with its provenance.

Source code in src/ergodic/discovery/algorithms/_pc.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Run PC-stable and return the discovered graph with its provenance."""
    oracle = resolve_test(self.test, data)
    nodes = sorted(oracle.variables)
    adj, sepsets = stable_skeleton(
        oracle, nodes, knowledge, alpha=self.alpha, max_cond_size=self.max_cond_size
    )
    collider_triples: frozenset[tuple[str, str, str]] | None = None
    ambiguous_triples: frozenset[tuple[str, str, str]] = frozenset()
    if self.colliders != "standard":
        collider_triples, ambiguous_triples = _vote_triples(
            oracle,
            nodes,
            adj,
            sepsets,
            rule=self.colliders,
            alpha=self.alpha,
            max_cond_size=self.max_cond_size,
        )
    graph, n_v_structures = _orient(
        nodes, adj, sepsets, knowledge, colliders=collider_triples, ambiguous=ambiguous_triples
    )
    try:
        final = graph.as_kind(GraphKind.CPDAG)
        cpdag_valid = 1.0
    except InvalidGraphError:
        final = graph.as_kind(GraphKind.MIXED)
        cpdag_valid = 0.0
    if knowledge is not None and not knowledge.is_consistent(final):
        raise ValueError(
            "PC produced a graph that violates the domain knowledge; the data and "
            "the constraints disagree (knowledge.violations(graph) lists how)"
        )
    diagnostics: dict[str, float] = dict(oracle.diagnostics())
    diagnostics["n_edges"] = float(len(final.edges))
    diagnostics["n_directed"] = float(sum(1 for e in final.edges if e.is_directed))
    diagnostics["n_undirected"] = float(sum(1 for e in final.edges if e.is_undirected))
    diagnostics["n_v_structures"] = float(n_v_structures)
    diagnostics["n_ambiguous_triples"] = float(len(ambiguous_triples))
    diagnostics["cpdag_valid"] = cpdag_valid
    return DiscoveryResult(
        graph=final,
        method="pc",
        test=oracle.name,
        alpha=self.alpha,
        n=oracle.n,
        diagnostics=diagnostics,
        sepsets=sepsets,
    )

PCMCI

PCMCI+ over a conditional-independence test, for time-series and panel data.

Parameters:

Name	Type	Description	Default
test	CITest | str	A :class:CITest instance or a name; "auto" reads the lagged frame's schema (all-continuous data gets the partial-correlation test).	'auto'
alpha	float	Significance level for every independence decision.	0.05
max_lag	int	Largest lag a reported link can have, at least one.	1
max_cond_size	int | None	Cap on the screen round and on the contemporaneous subset size, None for no cap.	None

fit(data) takes :class:~ergodic.data.TimeSeriesData or :class:~ergodic.data.PanelData (panel windows stack per entity) and returns a :class:TimeSeriesDiscoveryResult whose graph lives on the window nodes v[t-k].

fit(data, knowledge) reads a :class:~ergodic.knowledge.DomainKnowledge over the process names, not the window nodes, so a constraint on (x, y) applies at every lag, the contemporaneous slice included:

• A directional forbid x -> y means the link (x, y, lag) is never screened or MCI-tested at any lag and the lag-zero orientation x -> y is banned; the contemporaneous edge itself may survive when y -> x stays permitted (only the direction dies). A pair forbidden in both directions is an adjacency ban: the contemporaneous pair is removed without testing (and records no sepset) and both lagged directions are pruned. An explicit forbid(x, x) removes the self-memory link x[t-1] -> x[t].
• Tiers and order: a later variable never causes an earlier one at any lag, pruned exactly as a forbid; additionally a surviving contemporaneous edge between ordered processes is seeded earlier -> later before the collider phase, and a data-driven arrowhead that disagrees follows the orientation-conflict convention already in this file (the edge stays undirected and is counted in diagnostics["orientation_conflicts"]).
• Roots take no incoming cross-variable link at any lag and sinks no outgoing one; self-memory x[t-1] -> x[t] survives both, since a process can carry its own past without being caused by anything else. A skeleton restriction (restrict_to) bans pairs outside its support the same way, self-memory exempt.
• required (x, y) names no lag, so it cannot be forced; it is validated after the run instead, and fit raises ValueError when no compatible link (a lagged x -> y, an oriented lag-zero x -> y, or a still-undirected lag-zero edge) survived.
• confound has no analogue in a window graph (it carries no bidirected edges) and raises NotImplementedError.

diagnostics["n_knowledge_pruned"] counts what the knowledge took off the table: lagged links never screened, contemporaneous pairs removed outright, and lag-zero direction bans, one count each.

Source code in src/ergodic/discovery/algorithms/_pcmci.py

class PCMCI:
    """PCMCI+ over a conditional-independence test, for time-series and panel data.

    Args:
        test: A :class:`CITest` instance or a name; ``"auto"`` reads the
            lagged frame's schema (all-continuous data gets the
            partial-correlation test).
        alpha: Significance level for every independence decision.
        max_lag: Largest lag a reported link can have, at least one.
        max_cond_size: Cap on the screen round and on the contemporaneous
            subset size, ``None`` for no cap.

    ``fit(data)`` takes :class:`~ergodic.data.TimeSeriesData` or
    :class:`~ergodic.data.PanelData` (panel windows stack per entity) and
    returns a :class:`TimeSeriesDiscoveryResult` whose graph lives on the
    window nodes ``v[t-k]``.

    ``fit(data, knowledge)`` reads a
    :class:`~ergodic.knowledge.DomainKnowledge` over the process names, not
    the window nodes, so a constraint on ``(x, y)`` applies at every lag,
    the contemporaneous slice included:

    - A directional forbid ``x -> y`` means the link ``(x, y, lag)`` is
      never screened or MCI-tested at any lag and the lag-zero orientation
      ``x -> y`` is banned; the contemporaneous edge itself may survive when
      ``y -> x`` stays permitted (only the direction dies). A pair forbidden
      in both directions is an adjacency ban: the contemporaneous pair is
      removed without testing (and records no sepset) and both lagged
      directions are pruned. An explicit ``forbid(x, x)`` removes the
      self-memory link ``x[t-1] -> x[t]``.
    - Tiers and ``order``: a later variable never causes an earlier one at
      any lag, pruned exactly as a forbid; additionally a surviving
      contemporaneous edge between ordered processes is seeded
      earlier ``->`` later before the collider phase, and a data-driven
      arrowhead that disagrees follows the orientation-conflict convention
      already in this file (the edge stays undirected and is counted in
      ``diagnostics["orientation_conflicts"]``).
    - Roots take no incoming cross-variable link at any lag and sinks no
      outgoing one; self-memory ``x[t-1] -> x[t]`` survives both, since a
      process can carry its own past without being caused by anything else.
      A skeleton restriction (``restrict_to``) bans pairs outside its
      support the same way, self-memory exempt.
    - ``required (x, y)`` names no lag, so it cannot be forced; it is
      validated after the run instead, and ``fit`` raises ``ValueError``
      when no compatible link (a lagged ``x -> y``, an oriented lag-zero
      ``x -> y``, or a still-undirected lag-zero edge) survived.
    - ``confound`` has no analogue in a window graph (it carries no
      bidirected edges) and raises ``NotImplementedError``.

    ``diagnostics["n_knowledge_pruned"]`` counts what the knowledge took
    off the table: lagged links never screened, contemporaneous pairs
    removed outright, and lag-zero direction bans, one count each.
    """

    ingredient: ClassVar[str] = "test"

    def __init__(
        self,
        test: CITest | str = "auto",
        *,
        alpha: float = 0.05,
        max_lag: int = 1,
        max_cond_size: int | None = None,
    ) -> None:
        if not 0.0 < alpha < 1.0:
            raise ValueError(f"alpha must be in (0, 1), got {alpha}")
        if max_lag < 1:
            raise ValueError(f"max_lag must be >= 1, got {max_lag}")
        if max_cond_size is not None and max_cond_size < 0:
            raise ValueError(f"max_cond_size must be None or >= 0, got {max_cond_size}")
        self.test = test
        self.alpha = alpha
        self.max_lag = max_lag
        self.max_cond_size = max_cond_size

    # -- the run ---------------------------------------------------------

    def fit(
        self,
        data: TimeSeriesData | PanelData | None = None,
        knowledge: DomainKnowledge | None = None,
    ) -> TimeSeriesDiscoveryResult:
        """Run PCMCI+ and return the window graph with its provenance.

        ``knowledge`` constrains processes, so every constraint applies at
        every lag; the class docstring lists the mappings. ``confound``
        raises here (a window graph has no bidirected edges), and a
        ``required`` edge that survived at no lag raises after the run.
        """
        if not isinstance(data, TimeSeriesData | PanelData):
            raise TypeError(
                f"PCMCI needs TimeSeriesData or PanelData, got {type(data).__name__}; "
                "for cross-sectional rows use a tabular method (PC, GES, ...) via discover()"
            )
        if knowledge is not None and knowledge.confounded:
            raise NotImplementedError(
                "confounded pairs have no analogue in a window graph (it carries no "
                "bidirected edges); drop knowledge.confound for PCMCI"
            )
        flags = {"mci_truncated": 0.0, "acyclicity_fallback": 0.0}
        shallow = lagged_tabular(data, self.max_lag)  # validates the rows, voices the warnings
        deep_lag = 2 * self.max_lag
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")  # the shallow build already warned once
            try:
                frame = lagged_tabular(data, deep_lag)
            except ValueError:  # too few rows for the deep frame: condition at max_lag
                frame = shallow
                deep_lag = self.max_lag
                flags["mci_truncated"] = 1.0
        oracle = resolve_test(self.test, frame)
        names = sorted(data.variables)
        banned, removed, seeds, n_pruned = _knowledge_maps(knowledge, names, self.max_lag)

        screen = {target: self._pc1(oracle, names, target, banned) for target in names}
        mci = {target: self._mci(oracle, screen, target, deep_lag, banned) for target in names}
        links: dict[tuple[str, str, int], float] = {}
        for target in names:
            for (source, tau), p in sorted(mci[target].items()):
                links[(source, target, tau)] = p
        n_lagged = len(links)

        mci_names = {
            target: frozenset(lagged_node(source, tau) for source, tau in mci[target])
            for target in names
        }
        adj, sepsets, last_p = self._contemporaneous(oracle, names, mci_names, removed)
        present_seeds: dict[frozenset[str], tuple[str, str]] = {}
        for pair, orientation in seeds.items():
            a, b = sorted(pair)
            if b in adj[a]:  # seed only edges that survived the skeleton
                present_seeds[pair] = orientation
        directed, conflicts = _orient_slice(names, adj, sepsets, seeds=present_seeds, banned=banned)

        graph, n_contemporaneous, n_oriented = self._assemble(
            names, adj, directed, last_p, links, flags
        )
        if knowledge is not None:
            _check_required(knowledge, names, links, graph, self.max_lag)
        diagnostics: dict[str, float] = dict(oracle.diagnostics())
        diagnostics.update(flags)
        diagnostics["n_lagged_links"] = float(n_lagged)
        diagnostics["n_contemporaneous"] = float(n_contemporaneous)
        diagnostics["n_oriented"] = float(n_oriented)
        diagnostics["orientation_conflicts"] = float(conflicts)
        diagnostics["n_knowledge_pruned"] = float(n_pruned)
        out_sepsets = {
            frozenset(lagged_node(v, 0) for v in pair): frozenset(lagged_node(s, 0) for s in sep)
            for pair, sep in sepsets.items()
        }
        return TimeSeriesDiscoveryResult(
            graph=graph,
            method="pcmci",
            test=oracle.name,
            alpha=self.alpha,
            n=oracle.n,
            diagnostics=diagnostics,
            sepsets=out_sepsets,
            max_lag=self.max_lag,
            links=links,
        )

    # -- step 1: the PC1 lagged screen ---------------------------------------

    def _pc1(
        self,
        oracle: CITest,
        names: list[str],
        target: str,
        banned: frozenset[tuple[str, str]],
    ) -> dict[_Lagged, float]:
        """Surviving lagged parent candidates of one target, with their last p.

        A ``banned`` pair (knowledge says the source never causes the target)
        is skipped at every lag, never tested.
        """
        y = lagged_node(target, 0)
        pvals: dict[_Lagged, float] = {}
        for tau in range(1, self.max_lag + 1):
            for source in names:
                if (source, target) in banned:
                    continue  # knowledge: source never causes target, at any lag
                p = oracle.test(lagged_node(source, tau), y).p_value
                if p <= self.alpha:
                    pvals[(source, tau)] = p
        q = 1
        while q <= len(pvals) - 1 and (self.max_cond_size is None or q <= self.max_cond_size):
            ranked = [c for c, _ in sorted(pvals.items(), key=lambda item: (item[1], item[0]))]
            for candidate in ranked:
                if candidate not in pvals:
                    continue  # removed earlier in this round
                strongest = [c for c in ranked if c != candidate and c in pvals][:q]
                if len(strongest) < q:
                    continue
                conditioning = [lagged_node(source, tau) for source, tau in strongest]
                x = lagged_node(candidate[0], candidate[1])
                p = oracle.test(x, y, conditioning).p_value
                if p > self.alpha:
                    del pvals[candidate]
                else:
                    pvals[candidate] = p
            q += 1
        return pvals

    # -- step 2: MCI for the lagged links -------------------------------------

    def _mci(
        self,
        oracle: CITest,
        screen: dict[str, dict[_Lagged, float]],
        target: str,
        deep_lag: int,
        banned: frozenset[tuple[str, str]],
    ) -> dict[_Lagged, float]:
        """Re-test each screened lagged parent under the MCI conditioning set."""
        y = lagged_node(target, 0)
        parents = screen[target]
        kept: dict[_Lagged, float] = {}
        for source, tau in sorted(parents):
            if (source, target) in banned:  # pragma: no cover - the screen already skipped it
                continue
            conditioning = {
                lagged_node(other, k) for other, k in parents if (other, k) != (source, tau)
            }
            for other, k in screen[source]:
                if k + tau <= deep_lag:  # parents shifted past the deep frame are dropped
                    conditioning.add(lagged_node(other, k + tau))
            p = oracle.test(lagged_node(source, tau), y, sorted(conditioning)).p_value
            if p <= self.alpha:
                kept[(source, tau)] = p
        return kept

    # -- step 3: the contemporaneous skeleton ---------------------------------

    def _contemporaneous(
        self,
        oracle: CITest,
        names: list[str],
        mci_names: dict[str, frozenset[str]],
        removed: frozenset[frozenset[str]],
    ) -> tuple[
        dict[str, set[str]],
        dict[frozenset[str], frozenset[str]],
        dict[frozenset[str], float],
    ]:
        """PC-stable over the lag-zero pairs, always conditioning on MCI parents.

        ``removed`` pairs (knowledge bans in both directions) never enter the
        adjacency, are never tested, and record no sepset, exactly like a
        knowledge-pruned pair in the iid skeleton.
        """
        adj: dict[str, set[str]] = {v: set(names) - {v} for v in names}
        for pair in removed:
            a, b = sorted(pair)
            adj[a].discard(b)
            adj[b].discard(a)
        sepsets: dict[frozenset[str], frozenset[str]] = {}
        last_p: dict[frozenset[str], float] = {}
        q = 0
        while self.max_cond_size is None or q <= self.max_cond_size:
            snapshot = {v: frozenset(adj[v]) for v in names}
            if all(len(snapshot[v]) <= q for v in names):
                break
            for x in names:
                for y in sorted(snapshot[x]):
                    if y not in adj[x]:
                        continue  # removed earlier at this level
                    candidates = sorted(snapshot[x] - {y})
                    if len(candidates) < q:
                        continue
                    base = mci_names[x] | mci_names[y]
                    for subset in combinations(candidates, q):
                        conditioning = base | {lagged_node(s, 0) for s in subset}
                        result = oracle.test(lagged_node(x, 0), lagged_node(y, 0), conditioning)
                        if result.p_value > self.alpha:
                            adj[x].discard(y)
                            adj[y].discard(x)
                            sepsets[frozenset((x, y))] = frozenset(subset)
                            break
                        last_p[frozenset((x, y))] = result.p_value
            q += 1
        return adj, sepsets, last_p

    # -- step 5: the window graph ----------------------------------------------

    def _assemble(
        self,
        names: list[str],
        adj: dict[str, set[str]],
        directed: dict[frozenset[str], tuple[str, str]],
        last_p: dict[frozenset[str], float],
        links: dict[tuple[str, str, int], float],
        flags: dict[str, float],
    ) -> tuple[MixedGraph, int, int]:
        """Build the window graph and fold the contemporaneous edges into ``links``."""
        edges = [
            f"{lagged_node(source, tau)} -> {lagged_node(target, 0)}"
            for source, target, tau in sorted(links)
        ]
        n_contemporaneous = 0
        n_oriented = 0
        has_undirected = False
        for x, y in combinations(names, 2):
            if y not in adj[x]:
                continue
            pair = frozenset((x, y))
            n_contemporaneous += 1
            orientation = directed.get(pair)
            if orientation is None:
                has_undirected = True
                links[(x, y, 0)] = last_p[pair]
                edges.append(f"{lagged_node(x, 0)} --- {lagged_node(y, 0)}")
            else:
                tail, head = orientation
                n_oriented += 1
                links[(tail, head, 0)] = last_p[pair]
                edges.append(f"{lagged_node(tail, 0)} -> {lagged_node(head, 0)}")
        nodes = [lagged_node(v, k) for v in names for k in range(self.max_lag + 1)]
        if has_undirected:
            graph = mixed(edges, nodes=nodes)
        else:
            try:
                graph = dag(edges, nodes=nodes)
            except InvalidGraphError:
                graph = mixed(edges, nodes=nodes)
                flags["acyclicity_fallback"] = 1.0
        return graph, n_contemporaneous, n_oriented

fit(data=None, knowledge=None)

Run PCMCI+ and return the window graph with its provenance.

knowledge constrains processes, so every constraint applies at every lag; the class docstring lists the mappings. confound raises here (a window graph has no bidirected edges), and a required edge that survived at no lag raises after the run.

Source code in src/ergodic/discovery/algorithms/_pcmci.py

def fit(
    self,
    data: TimeSeriesData | PanelData | None = None,
    knowledge: DomainKnowledge | None = None,
) -> TimeSeriesDiscoveryResult:
    """Run PCMCI+ and return the window graph with its provenance.

    ``knowledge`` constrains processes, so every constraint applies at
    every lag; the class docstring lists the mappings. ``confound``
    raises here (a window graph has no bidirected edges), and a
    ``required`` edge that survived at no lag raises after the run.
    """
    if not isinstance(data, TimeSeriesData | PanelData):
        raise TypeError(
            f"PCMCI needs TimeSeriesData or PanelData, got {type(data).__name__}; "
            "for cross-sectional rows use a tabular method (PC, GES, ...) via discover()"
        )
    if knowledge is not None and knowledge.confounded:
        raise NotImplementedError(
            "confounded pairs have no analogue in a window graph (it carries no "
            "bidirected edges); drop knowledge.confound for PCMCI"
        )
    flags = {"mci_truncated": 0.0, "acyclicity_fallback": 0.0}
    shallow = lagged_tabular(data, self.max_lag)  # validates the rows, voices the warnings
    deep_lag = 2 * self.max_lag
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")  # the shallow build already warned once
        try:
            frame = lagged_tabular(data, deep_lag)
        except ValueError:  # too few rows for the deep frame: condition at max_lag
            frame = shallow
            deep_lag = self.max_lag
            flags["mci_truncated"] = 1.0
    oracle = resolve_test(self.test, frame)
    names = sorted(data.variables)
    banned, removed, seeds, n_pruned = _knowledge_maps(knowledge, names, self.max_lag)

    screen = {target: self._pc1(oracle, names, target, banned) for target in names}
    mci = {target: self._mci(oracle, screen, target, deep_lag, banned) for target in names}
    links: dict[tuple[str, str, int], float] = {}
    for target in names:
        for (source, tau), p in sorted(mci[target].items()):
            links[(source, target, tau)] = p
    n_lagged = len(links)

    mci_names = {
        target: frozenset(lagged_node(source, tau) for source, tau in mci[target])
        for target in names
    }
    adj, sepsets, last_p = self._contemporaneous(oracle, names, mci_names, removed)
    present_seeds: dict[frozenset[str], tuple[str, str]] = {}
    for pair, orientation in seeds.items():
        a, b = sorted(pair)
        if b in adj[a]:  # seed only edges that survived the skeleton
            present_seeds[pair] = orientation
    directed, conflicts = _orient_slice(names, adj, sepsets, seeds=present_seeds, banned=banned)

    graph, n_contemporaneous, n_oriented = self._assemble(
        names, adj, directed, last_p, links, flags
    )
    if knowledge is not None:
        _check_required(knowledge, names, links, graph, self.max_lag)
    diagnostics: dict[str, float] = dict(oracle.diagnostics())
    diagnostics.update(flags)
    diagnostics["n_lagged_links"] = float(n_lagged)
    diagnostics["n_contemporaneous"] = float(n_contemporaneous)
    diagnostics["n_oriented"] = float(n_oriented)
    diagnostics["orientation_conflicts"] = float(conflicts)
    diagnostics["n_knowledge_pruned"] = float(n_pruned)
    out_sepsets = {
        frozenset(lagged_node(v, 0) for v in pair): frozenset(lagged_node(s, 0) for s in sep)
        for pair, sep in sepsets.items()
    }
    return TimeSeriesDiscoveryResult(
        graph=graph,
        method="pcmci",
        test=oracle.name,
        alpha=self.alpha,
        n=oracle.n,
        diagnostics=diagnostics,
        sepsets=out_sepsets,
        max_lag=self.max_lag,
        links=links,
    )

RFCI

RFCI over a conditional-independence test: FCI's marks at PC's price.

Parameters:

Name	Type	Description	Default
test	CITest | str	A :class:CITest instance (bound or unbound, including :class:DSeparationOracle), or a name: "auto" (pick by the data schema), "linear"/"fisher_z", "discrete"/ "g_sq", "mixed"/"dg", "gcm".	'auto'
alpha	float	Significance level for the independence decisions.	0.05
max_cond_size	int | None	Cap on the conditioning-set size in the skeleton phase; None for no cap. The later stages only ever condition on subsets of recorded separating sets, which the cap already bounds, so unlike FCI there is no second phase the cap must reach.	None

fit(data, knowledge) returns a :class:DiscoveryResult whose graph carries the PAG face when the output happens to be a valid complete PAG and the mixed face with diagnostics["pag_valid"] = 0.0 otherwise. Unlike FCI, the mixed face is an expected honest outcome here even under a perfect oracle: the RFCI skeleton can keep edges only the possible-d-sep phase would remove, and the result then stands for more than one Markov equivalence class. Domain knowledge enters exactly as in FCI: prune (a pair forbidden in both directions is never tested), protect (required and confounded pairs are never tested or removed), orient through ancestry (the same seeding helpers), validate at the end.

Source code in src/ergodic/discovery/algorithms/_rfci.py

class RFCI:
    """RFCI over a conditional-independence test: FCI's marks at PC's price.

    Args:
        test: A :class:`CITest` instance (bound or unbound, including
            :class:`DSeparationOracle`), or a name: ``"auto"`` (pick by the
            data schema), ``"linear"``/``"fisher_z"``, ``"discrete"``/
            ``"g_sq"``, ``"mixed"``/``"dg"``, ``"gcm"``.
        alpha: Significance level for the independence decisions.
        max_cond_size: Cap on the conditioning-set size in the skeleton
            phase; ``None`` for no cap. The later stages only ever condition
            on subsets of recorded separating sets, which the cap already
            bounds, so unlike FCI there is no second phase the cap must reach.

    ``fit(data, knowledge)`` returns a :class:`DiscoveryResult` whose graph
    carries the PAG face when the output happens to be a valid complete PAG
    and the ``mixed`` face with ``diagnostics["pag_valid"] = 0.0`` otherwise.
    Unlike FCI, the mixed face is an expected honest outcome here even under
    a perfect oracle: the RFCI skeleton can keep edges only the
    possible-d-sep phase would remove, and the result then stands for more
    than one Markov equivalence class. Domain knowledge enters exactly as in
    FCI: prune (a pair forbidden in both directions is never tested), protect
    (required and confounded pairs are never tested or removed), orient
    through ancestry (the same seeding helpers), validate at the end.
    """

    ingredient: ClassVar[str] = "test"

    def __init__(
        self,
        test: CITest | str = "auto",
        *,
        alpha: float = 0.05,
        max_cond_size: int | None = None,
    ) -> None:
        if not 0.0 < alpha < 1.0:
            raise ValueError(f"alpha must be in (0, 1), got {alpha}")
        if max_cond_size is not None and max_cond_size < 0:
            raise ValueError(f"max_cond_size must be None or >= 0, got {max_cond_size}")
        self.test = test
        self.alpha = alpha
        self.max_cond_size = max_cond_size

    # -- the run ---------------------------------------------------------

    def fit(
        self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
    ) -> DiscoveryResult:
        """Run RFCI and return the discovered graph with its provenance."""
        oracle = resolve_test(self.test, data)
        nodes = sorted(oracle.variables)
        adj, sepsets = stable_skeleton(
            oracle, nodes, knowledge, alpha=self.alpha, max_cond_size=self.max_cond_size
        )
        protected = _protected_pairs(knowledge)
        counts = _Counts()
        triples = self._triple_stage(oracle, adj, sepsets, protected, counts)
        closed, n_v_structures = self._orientation_stage(
            oracle, nodes, adj, sepsets, protected, triples, knowledge, counts
        )
        try:
            final = closed.as_kind(GraphKind.PAG)
            pag_valid = 1.0
        except InvalidGraphError:
            final = closed.as_kind(GraphKind.MIXED)
            pag_valid = 0.0
        if knowledge is not None and not knowledge.is_consistent(final):
            raise ValueError(
                "RFCI produced a graph that violates the domain knowledge; the data and "
                "the constraints disagree (knowledge.violations(graph) lists how)"
            )
        diagnostics: dict[str, float] = dict(oracle.diagnostics())
        diagnostics["n_edges"] = float(len(final.edges))
        diagnostics["n_v_structures"] = float(n_v_structures)
        diagnostics["n_stage2_removed"] = float(counts.stage2_removed)
        diagnostics["n_stage3_removed"] = float(counts.stage3_removed)
        diagnostics["n_triple_tests"] = float(counts.triple_tests)
        diagnostics["n_path_tests"] = float(counts.path_tests)
        for label, mark in (
            ("n_arrowheads", Mark.ARROW),
            ("n_tails", Mark.TAIL),
            ("n_circles", Mark.CIRCLE),
        ):
            diagnostics[label] = float(
                sum((e.u_mark is mark) + (e.v_mark is mark) for e in final.edges)
            )
        diagnostics["pag_valid"] = pag_valid
        return DiscoveryResult(
            graph=final,
            method="rfci",
            test=oracle.name,
            alpha=self.alpha,
            n=oracle.n,
            diagnostics=diagnostics,
            sepsets=sepsets,
        )

    # -- stage two: the unshielded-triple rule -----------------------------

    def _triple_stage(
        self,
        oracle: CITest,
        adj: _Adjacency,
        sepsets: _Sepsets,
        protected: frozenset[frozenset[str]],
        counts: _Counts,
    ) -> list[_Triple]:
        """The unshielded-triple rule (Lemma 3.1): trust no triple untested.

        Each unshielded triple ``(x, m, y)`` re-tests its two edges given the
        recorded separating set of ``(x, y)`` minus ``m``. Both dependent:
        the triple is kept for orientation once the worklist empties
        (deferring the arrowheads is what makes the stage order independent,
        the paper's two-list design). Any independence: the edge comes off, a
        minimal separating set is recorded by re-testing subsets of the same
        set in ascending size (the set that exposed the independence need not
        be minimal, and Lemma 3.1's condition needs a minimal one), triples
        the removal destroyed leave both lists, and triples it unshielded
        join the worklist. A triple whose ends were never tested (a
        knowledge-pruned pair records no separating set) is skipped and stays
        unoriented, the same convention FCI's collider phase uses.
        """
        queue: deque[_Triple] = deque(_unshielded_triples(adj))
        queued: set[_Triple] = set(queue)
        kept: list[_Triple] = []
        while queue:
            triple = queue.popleft()
            queued.discard(triple)
            x, m, y = triple
            separator = sepsets.get(frozenset((x, y)))
            if separator is None:
                continue
            conditioning = sorted(separator - {m})
            independent: list[tuple[str, str]] = []
            for u, v in ((x, m), (m, y)):
                if frozenset((u, v)) in protected:
                    continue  # protect: required and confounded pairs are never tested
                counts.triple_tests += 1
                if oracle.independent(u, v, conditioning, alpha=self.alpha):
                    independent.append((u, v))
            if not independent:
                kept.append(triple)
                continue
            for u, v in independent:
                minimal = self._minimal_separator(oracle, u, v, conditioning, counts)
                _delete_edge(adj, sepsets, u, v, minimal, queue, queued, kept)
                counts.stage2_removed += 1
        return kept

    def _minimal_separator(
        self, oracle: CITest, u: str, v: str, base: list[str], counts: _Counts
    ) -> frozenset[str]:
        """The first separating subset of ``base``, swept in ascending size.

        Ascending order makes the winner minimum-cardinality, so minimal,
        which is what the soundness lemmas consume. The full set separated
        the pair a moment ago, so the sweep always lands.
        """
        for size in range(len(base) + 1):
            for subset in combinations(base, size):
                counts.triple_tests += 1
                if oracle.independent(u, v, subset, alpha=self.alpha):
                    return frozenset(subset)
        return frozenset(base)  # unreachable: the full set is the sweep's last subset

    # -- stage three: closure with the guarded discriminating-path rule -----

    def _orientation_stage(
        self,
        oracle: CITest,
        nodes: list[str],
        adj: _Adjacency,
        sepsets: _Sepsets,
        protected: frozenset[frozenset[str]],
        triples: list[_Triple],
        knowledge: DomainKnowledge | None,
        counts: _Counts,
    ) -> tuple[MixedGraph, int]:
        """Stage three: Zhang's rules R1 to R10 with R4 guarded by Lemma 3.2.

        Marks start as circles plus the stage-two v-structures and any
        knowledge seeds, then close under ``fci_closure``. The discriminating
        callback reads a snapshot of the marks the closure started from, so a
        path that only exists through marks set later in the same pass comes
        back unfound; a pass that changed marks but deleted nothing therefore
        reruns from the closed graph with a fresh snapshot, until a pass
        changes nothing, at which point the callback saw exactly the marks
        the closure saw. A pass that found independencies instead fills the
        pending box: those edges come off with their minimal separating sets
        recorded, stage two reruns in full on the smaller adjacency (the
        simplest correct restart; the oracle's cache makes the re-walk
        cheap), marks reset, and the closure starts over. The loop ends when
        a closure pass ends with no pending deletions. The paper accepts that
        this stage is order dependent; every walk here is sorted, so the run
        is deterministic.
        """
        while True:
            marks = _circle_marks(adj)
            n_v_structures = _orient_triples(marks, sepsets, triples)
            _seed_knowledge(marks, knowledge)
            graph = _to_graph(nodes, marks)
            closed, pending = self._close_once(oracle, graph, sepsets, protected, counts)
            while not pending and closed.edges != graph.edges:
                graph = closed
                closed, pending = self._close_once(oracle, graph, sepsets, protected, counts)
            if not pending:
                return closed, n_v_structures
            for pair in sorted(pending, key=sorted):
                u, v = sorted(pair)
                adj[u].discard(v)
                adj[v].discard(u)
                sepsets[pair] = pending[pair]
                counts.stage3_removed += 1
            triples = self._triple_stage(oracle, adj, sepsets, protected, counts)

    def _close_once(
        self,
        oracle: CITest,
        graph: MixedGraph,
        sepsets: _Sepsets,
        protected: frozenset[frozenset[str]],
        counts: _Counts,
    ) -> tuple[MixedGraph, _Pending]:
        """One closure pass; the pending box collects Lemma 3.2 removals."""
        pending: _Pending = {}
        closed = fci_closure(
            graph,
            discriminating=self._lemma_3_2(oracle, graph, sepsets, protected, pending, counts),
        )
        return closed, pending

    def _lemma_3_2(
        self,
        oracle: CITest,
        graph: MixedGraph,
        sepsets: _Sepsets,
        protected: frozenset[frozenset[str]],
        pending: _Pending,
        counts: _Counts,
    ) -> Callable[[str, str, str, str], bool | None]:
        """R4's collider question, guarded by Lemma 3.2's extra tests.

        The closure hands over the ends ``(d, a, b, c)`` it found; the path
        itself is recovered here over the snapshot marks. Every consecutive
        pair on the recovered path must stay dependent given every subset of
        the recorded separating set of ``(d, c)`` minus the pair (the paper
        is explicit that the single full set is not sufficient). The first
        independence found goes into the pending box with its minimal
        separating set, and the path abstains; a path through an edge already
        condemned this pass abstains too, since the pass restarts anyway.
        With every pair dependent the answer is FCI's: ``b`` is a collider
        exactly when it is outside the recorded separating set. Verdicts are
        memoized per pass; the snapshot, the oracle's cache, and the recorded
        sets make the callback a pure function of its arguments within one.
        """
        marks = _mark_maps(graph)
        adj = {n: sorted(marks[n]) for n in marks}
        memo: dict[tuple[str, str, str, str], bool | None] = {}

        def judge(d: str, a: str, b: str, c: str) -> bool | None:
            separator = sepsets.get(frozenset((d, c)))
            if separator is None:
                return None  # the pair was never tested (knowledge-pruned)
            path = _minimal_discriminating_path(marks, adj, d, a, b, c)
            if path is None:
                return None  # not in this snapshot; the refresh loop will re-ask
            pairs = list(pairwise(path))
            if any(frozenset(pair) in pending for pair in pairs):
                return None  # the path runs through an edge already condemned
            for r, q in pairs:
                pair = frozenset((r, q))
                if pair in protected:
                    continue  # protect: never tested or removed
                base = sorted(separator - {r, q})
                for size in range(len(base) + 1):
                    for subset in combinations(base, size):
                        counts.path_tests += 1
                        if oracle.independent(r, q, subset, alpha=self.alpha):
                            pending[pair] = frozenset(subset)
                            return None
            return b not in separator

        def callback(d: str, a: str, b: str, c: str) -> bool | None:
            key = (d, a, b, c)
            if key not in memo:
                memo[key] = judge(d, a, b, c)
            return memo[key]

        return callback

fit(data=None, knowledge=None)

Run RFCI and return the discovered graph with its provenance.

Source code in src/ergodic/discovery/algorithms/_rfci.py

def fit(
    self, data: TabularData | None = None, knowledge: DomainKnowledge | None = None
) -> DiscoveryResult:
    """Run RFCI and return the discovered graph with its provenance."""
    oracle = resolve_test(self.test, data)
    nodes = sorted(oracle.variables)
    adj, sepsets = stable_skeleton(
        oracle, nodes, knowledge, alpha=self.alpha, max_cond_size=self.max_cond_size
    )
    protected = _protected_pairs(knowledge)
    counts = _Counts()
    triples = self._triple_stage(oracle, adj, sepsets, protected, counts)
    closed, n_v_structures = self._orientation_stage(
        oracle, nodes, adj, sepsets, protected, triples, knowledge, counts
    )
    try:
        final = closed.as_kind(GraphKind.PAG)
        pag_valid = 1.0
    except InvalidGraphError:
        final = closed.as_kind(GraphKind.MIXED)
        pag_valid = 0.0
    if knowledge is not None and not knowledge.is_consistent(final):
        raise ValueError(
            "RFCI produced a graph that violates the domain knowledge; the data and "
            "the constraints disagree (knowledge.violations(graph) lists how)"
        )
    diagnostics: dict[str, float] = dict(oracle.diagnostics())
    diagnostics["n_edges"] = float(len(final.edges))
    diagnostics["n_v_structures"] = float(n_v_structures)
    diagnostics["n_stage2_removed"] = float(counts.stage2_removed)
    diagnostics["n_stage3_removed"] = float(counts.stage3_removed)
    diagnostics["n_triple_tests"] = float(counts.triple_tests)
    diagnostics["n_path_tests"] = float(counts.path_tests)
    for label, mark in (
        ("n_arrowheads", Mark.ARROW),
        ("n_tails", Mark.TAIL),
        ("n_circles", Mark.CIRCLE),
    ):
        diagnostics[label] = float(
            sum((e.u_mark is mark) + (e.v_mark is mark) for e in final.edges)
        )
    diagnostics["pag_valid"] = pag_valid
    return DiscoveryResult(
        graph=final,
        method="rfci",
        test=oracle.name,
        alpha=self.alpha,
        n=oracle.n,
        diagnostics=diagnostics,
        sepsets=sepsets,
    )

Score

Base for decomposable scores; see the module docstring.

Subclasses set name, needs, and the property flags, may take options in __init__ (set them before calling super().__init__(data)), and implement the method matching needs: local_from_arrays for "numeric", local_from_codes for "codes". _prepare runs once after binding for sufficient statistics; a native score may also override _local_value to skip array assembly entirely (the covariance path in BIC).

Source code in src/ergodic/discovery/_score.py

class Score:
    """Base for decomposable scores; see the module docstring.

    Subclasses set ``name``, ``needs``, and the property flags, may take
    options in ``__init__`` (set them before calling ``super().__init__(data)``),
    and implement the method matching ``needs``: ``local_from_arrays`` for
    ``"numeric"``, ``local_from_codes`` for ``"codes"``. ``_prepare`` runs once
    after binding for sufficient statistics; a native score may also override
    ``_local_value`` to skip array assembly entirely (the covariance path in
    BIC).
    """

    name: ClassVar[str] = "score"
    needs: ClassVar[str] = "numeric"
    decomposable: ClassVar[bool] = True
    score_equivalent: ClassVar[bool] = True

    def __init__(self, data: TabularData | None = None) -> None:
        self._data: TabularData | None = None
        self._n: int = 0
        self._names: frozenset[str] = frozenset()
        self._blocks: dict[str, npt.NDArray[np.float64]] = {}
        self._codes: dict[str, npt.NDArray[np.int64]] = {}
        self._cards: dict[str, int] = {}
        self._cache: dict[tuple[str, frozenset[str]], float] = {}
        self._n_evals: int = 0
        self._cache_hits: int = 0
        if data is not None:
            self._attach(data)

    # -- binding ---------------------------------------------------------

    @property
    def is_bound(self) -> bool:
        return self._data is not None

    @property
    def data(self) -> TabularData | None:
        """The dataset this score is bound to, after dropping incomplete rows."""
        return self._data

    @property
    def n(self) -> int:
        return self._n

    @property
    def variables(self) -> tuple[str, ...]:
        """The variables this oracle answers about, sorted."""
        return tuple(sorted(self._names))

    def bind(self, data: TabularData) -> Self:
        """A new instance of this score bound to ``data``; never mutates ``self``."""
        fresh = copy.copy(self)
        fresh._attach(data)
        return fresh

    def _attach(self, data: TabularData) -> None:
        complete = data.dropna()
        self._data = complete
        self._n = complete.n_rows
        self._names = frozenset(complete.variables)
        self._blocks = {}
        self._codes = {}
        self._cards = {}
        if self.needs == "numeric":
            self._blocks = _numeric_blocks(complete)
        elif self.needs == "codes":
            self._codes, self._cards = _codes_columns(complete)
        else:
            raise ValueError(f"{type(self).__name__}.needs must be 'numeric' or 'codes'")
        self._cache = {}
        self._n_evals = 0
        self._cache_hits = 0
        self._prepare()

    def _prepare(self) -> None:
        """Hook for one-time sufficient statistics after binding."""

    # -- the queries -------------------------------------------------------

    def local(self, node: str, parents: Iterable[str] = ()) -> float:
        """The cached local score of ``node`` given ``parents``."""
        self._require_bound()
        parent_set = frozenset(parents)
        self._validate(node, parent_set)
        key = (node, parent_set)
        hit = self._cache.get(key)
        if hit is not None:
            self._cache_hits += 1
            return hit
        value = self._local_value(node, parent_set)
        self._cache[key] = value
        self._n_evals += 1
        return value

    def local_diff(self, node: str, parents: Iterable[str], add: str) -> float:
        """The gain from adding ``add`` to the parents of ``node``.

        The default is two locals; an implementation may override it with a
        cheaper incremental form.
        """
        base = frozenset(parents)
        return self.local(node, base | {add}) - self.local(node, base)

    def enumerate_locals(
        self, node: str, candidates: Iterable[str], max_parents: int | None = None
    ) -> dict[frozenset[str], float]:
        """Every candidate parent set up to ``max_parents``, scored.

        The access pattern of the exact solvers (A*, dynamic programming,
        ILP): all the local scores up front under an in-degree cap.
        """
        pool = sorted(frozenset(candidates))
        cap = len(pool) if max_parents is None else min(max_parents, len(pool))
        out: dict[frozenset[str], float] = {}
        for size in range(cap + 1):
            for combo in combinations(pool, size):
                parent_set = frozenset(combo)
                out[parent_set] = self.local(node, parent_set)
        return out

    def graph_score(self, parents_of: dict[str, frozenset[str]]) -> float:
        """The decomposable total: the sum of locals over a parent map."""
        return float(sum(self.local(node, parents) for node, parents in parents_of.items()))

    def diagnostics(self) -> dict[str, float]:
        """Run counters: evaluations and cache hits."""
        return {"n_score_evals": float(self._n_evals), "cache_hits": float(self._cache_hits)}

    # -- subclass surface ---------------------------------------------------

    def local_from_arrays(self, y: npt.NDArray[np.float64], X: npt.NDArray[np.float64]) -> float:
        """The local score from arrays: ``y`` is ``(n, dy)``, ``X`` is ``(n, k)``."""
        raise NotImplementedError(f"{type(self).__name__} does not implement local_from_arrays")

    def local_from_codes(
        self,
        y: npt.NDArray[np.int64],
        y_card: int,
        zs: tuple[npt.NDArray[np.int64], ...],
        z_cards: tuple[int, ...],
    ) -> float:
        """The local score from codes: ``y`` is ``(n,)`` integer codes, ``zs`` the parents'."""
        raise NotImplementedError(f"{type(self).__name__} does not implement local_from_codes")

    # -- internals -----------------------------------------------------------

    def _local_value(self, node: str, parents: frozenset[str]) -> float:
        """Assemble the declared input form and score it; native scores may override."""
        ordered = sorted(parents)
        if self.needs == "numeric":
            y = self._blocks[node]
            X = (
                np.hstack([self._blocks[p] for p in ordered])
                if ordered
                else np.empty((self._n, 0), dtype=np.float64)
            )
            return float(self.local_from_arrays(y, X))
        zs = tuple(self._codes[p] for p in ordered)
        z_cards = tuple(self._cards[p] for p in ordered)
        return float(self.local_from_codes(self._codes[node], self._cards[node], zs, z_cards))

    def _require_bound(self) -> None:
        if not self.is_bound:
            raise RuntimeError(
                f"{type(self).__name__} is not bound to data; construct it with data "
                "or call bind(data)"
            )

    def _validate(self, node: str, parents: frozenset[str]) -> None:
        unknown = sorted(({node} | parents) - self._names)
        if unknown:
            raise ValueError(f"unknown variable(s) {unknown}; the data has {sorted(self._names)}")
        if node in parents:
            raise ValueError(f"{node!r} cannot be its own parent")

data property

The dataset this score is bound to, after dropping incomplete rows.

variables property

The variables this oracle answers about, sorted.

bind(data)

A new instance of this score bound to data; never mutates self.

Source code in src/ergodic/discovery/_score.py

def bind(self, data: TabularData) -> Self:
    """A new instance of this score bound to ``data``; never mutates ``self``."""
    fresh = copy.copy(self)
    fresh._attach(data)
    return fresh

diagnostics()

Run counters: evaluations and cache hits.

Source code in src/ergodic/discovery/_score.py

def diagnostics(self) -> dict[str, float]:
    """Run counters: evaluations and cache hits."""
    return {"n_score_evals": float(self._n_evals), "cache_hits": float(self._cache_hits)}

enumerate_locals(node, candidates, max_parents=None)

Every candidate parent set up to max_parents, scored.

The access pattern of the exact solvers (A*, dynamic programming, ILP): all the local scores up front under an in-degree cap.

Source code in src/ergodic/discovery/_score.py

def enumerate_locals(
    self, node: str, candidates: Iterable[str], max_parents: int | None = None
) -> dict[frozenset[str], float]:
    """Every candidate parent set up to ``max_parents``, scored.

    The access pattern of the exact solvers (A*, dynamic programming,
    ILP): all the local scores up front under an in-degree cap.
    """
    pool = sorted(frozenset(candidates))
    cap = len(pool) if max_parents is None else min(max_parents, len(pool))
    out: dict[frozenset[str], float] = {}
    for size in range(cap + 1):
        for combo in combinations(pool, size):
            parent_set = frozenset(combo)
            out[parent_set] = self.local(node, parent_set)
    return out

graph_score(parents_of)

The decomposable total: the sum of locals over a parent map.

Source code in src/ergodic/discovery/_score.py

def graph_score(self, parents_of: dict[str, frozenset[str]]) -> float:
    """The decomposable total: the sum of locals over a parent map."""
    return float(sum(self.local(node, parents) for node, parents in parents_of.items()))

local(node, parents=())

The cached local score of node given parents.

Source code in src/ergodic/discovery/_score.py

def local(self, node: str, parents: Iterable[str] = ()) -> float:
    """The cached local score of ``node`` given ``parents``."""
    self._require_bound()
    parent_set = frozenset(parents)
    self._validate(node, parent_set)
    key = (node, parent_set)
    hit = self._cache.get(key)
    if hit is not None:
        self._cache_hits += 1
        return hit
    value = self._local_value(node, parent_set)
    self._cache[key] = value
    self._n_evals += 1
    return value

local_diff(node, parents, add)

The gain from adding add to the parents of node.

The default is two locals; an implementation may override it with a cheaper incremental form.

Source code in src/ergodic/discovery/_score.py

def local_diff(self, node: str, parents: Iterable[str], add: str) -> float:
    """The gain from adding ``add`` to the parents of ``node``.

    The default is two locals; an implementation may override it with a
    cheaper incremental form.
    """
    base = frozenset(parents)
    return self.local(node, base | {add}) - self.local(node, base)

local_from_arrays(y, X)

The local score from arrays: y is (n, dy), X is (n, k).

Source code in src/ergodic/discovery/_score.py

def local_from_arrays(self, y: npt.NDArray[np.float64], X: npt.NDArray[np.float64]) -> float:
    """The local score from arrays: ``y`` is ``(n, dy)``, ``X`` is ``(n, k)``."""
    raise NotImplementedError(f"{type(self).__name__} does not implement local_from_arrays")

local_from_codes(y, y_card, zs, z_cards)

The local score from codes: y is (n,) integer codes, zs the parents'.

Source code in src/ergodic/discovery/_score.py

def local_from_codes(
    self,
    y: npt.NDArray[np.int64],
    y_card: int,
    zs: tuple[npt.NDArray[np.int64], ...],
    z_cards: tuple[int, ...],
) -> float:
    """The local score from codes: ``y`` is ``(n,)`` integer codes, ``zs`` the parents'."""
    raise NotImplementedError(f"{type(self).__name__} does not implement local_from_codes")

TimeSeriesDiscoveryResult dataclass

Bases: DiscoveryResult

A :class:DiscoveryResult over window nodes, with the per-link p-values.

links maps (source, target, lag) to the p-value of the test that kept the link. A lag of one or more is the lagged edge source[t-lag] -> target[t]. A lag of zero is a contemporaneous edge: an oriented edge is keyed (source, target), an undirected edge is keyed once with the pair in sorted order.

Source code in src/ergodic/discovery/_tsresult.py

@dataclass(frozen=True)
class TimeSeriesDiscoveryResult(DiscoveryResult):
    """A :class:`DiscoveryResult` over window nodes, with the per-link p-values.

    ``links`` maps ``(source, target, lag)`` to the p-value of the test that
    kept the link. A lag of one or more is the lagged edge
    ``source[t-lag] -> target[t]``. A lag of zero is a contemporaneous edge:
    an oriented edge is keyed ``(source, target)``, an undirected edge is
    keyed once with the pair in sorted order.
    """

    max_lag: int = 0
    links: Mapping[tuple[str, str, int], float] = field(default_factory=dict)

    def __repr__(self) -> str:
        n_lagged = sum(1 for _, _, lag in self.links if lag >= 1)
        n_contemporaneous = sum(1 for _, _, lag in self.links if lag == 0)
        ingredient = self.test if self.test is not None else self.score
        body = (
            f"{self.method}/{ingredient}: {self.graph.kind.value} on {len(self.graph.nodes)} "
            f"window nodes, max_lag={self.max_lag}, {n_lagged} lagged links, "
            f"{n_contemporaneous} contemporaneous"
        )
        if self.alpha is not None:
            body += f", alpha={self.alpha:g}"
        return f"TimeSeriesDiscoveryResult({body}, n={self.n})"

    def summary_edges(self) -> frozenset[tuple[str, str]]:
        """Variable-level pairs implied by the links.

        ``(x, y)`` is present when any lagged link ``x -> y`` exists or a
        contemporaneous edge ``x -> y`` is oriented in the window graph. An
        undirected contemporaneous pair contributes both orders, and a lagged
        self-link (autocorrelation) contributes ``(x, x)``.
        """
        out: set[tuple[str, str]] = set()
        for source, target, lag in self.links:
            if lag == 0:
                edge = self.graph.get_edge(lagged_node(source, 0), lagged_node(target, 0))
                if edge is not None and edge.is_undirected:
                    out.add((target, source))
            out.add((source, target))
        return frozenset(out)

summary_edges()

Variable-level pairs implied by the links.

(x, y) is present when any lagged link x -> y exists or a contemporaneous edge x -> y is oriented in the window graph. An undirected contemporaneous pair contributes both orders, and a lagged self-link (autocorrelation) contributes (x, x).

Source code in src/ergodic/discovery/_tsresult.py

def summary_edges(self) -> frozenset[tuple[str, str]]:
    """Variable-level pairs implied by the links.

    ``(x, y)`` is present when any lagged link ``x -> y`` exists or a
    contemporaneous edge ``x -> y`` is oriented in the window graph. An
    undirected contemporaneous pair contributes both orders, and a lagged
    self-link (autocorrelation) contributes ``(x, x)``.
    """
    out: set[tuple[str, str]] = set()
    for source, target, lag in self.links:
        if lag == 0:
            edge = self.graph.get_edge(lagged_node(source, 0), lagged_node(target, 0))
            if edge is not None and edge.is_undirected:
                out.add((target, source))
        out.add((source, target))
    return frozenset(out)

bootstrap(data, *, method='pc', n_bootstrap=100, seed=0, knowledge=None, block_length=None, **kwargs)

Rerun a discovery method on bootstrap resamples and count what survives.

Answers the question a single run cannot: how stable is this edge under resampling? Each replicate draws a resample, runs discover(sample, method=method, knowledge=knowledge, **kwargs), and the result aggregates adjacency and glyph frequencies over the runs that finished. Extra keyword arguments flow to the algorithm unchanged (alpha, max_cond_size, score=..., max_lag=...), so any registered method, iid or temporal, sits behind the same call.

An iid method on a :class:~ergodic.data.TabularData draws data.n_rows rows with replacement, the classic bootstrap; block_length does not apply there and raises. A temporal method ("pcmci") never resamples rows, which would shred the serial dependence it reads. On a :class:~ergodic.data.TimeSeriesData it runs a moving-block bootstrap: ceil(T / L) blocks of L consecutive rows (block_length, default max(ceil(sqrt(T)), 2 * (max_lag + 1))) are drawn with replacement and enter the replicate as separate entities of a panel, "block_0", "block_1", ... with fresh time stamps. Blocks-as-entities is the point: gluing blocks back into one series would build lagged windows that straddle two unrelated fragments and let fake adjacencies in, while per-entity pooling never reads across a block seam. On a :class:~ergodic.data.PanelData it runs the cluster bootstrap instead, drawing n_entities entities whole with replacement (relabeled in draw order so repeated draws stay distinct); block_length does not apply to a panel and raises. One honesty note: block resampling treats the series as stationary at the block scale, so a trend or a regime change inside the sampling window leaks into the frequencies. An iid method on temporal data still refuses, inside the replicate, exactly as :func:discover does.

The frequencies measure resampling variability, not Bayesian uncertainty; they are not posterior probabilities over graphs (graph posteriors belong to a later phase). The run is deterministic: the same data, seed, and settings give the same result. A temporal method keys the frequencies on window-graph glyphs like "X[t-1] --> Y[t]".

A replicate that raises ValueError or RuntimeError (degenerate resamples happen: a constant column, a knowledge conflict on noisy data) counts in n_failed and is skipped; frequencies are over the successful runs only. When every replicate fails, the wrapper raises RuntimeError chaining the last error.

Parameters:

Name	Type	Description	Default
data	TabularData | TimeSeriesData | PanelData	The dataset to resample: tabular for the iid methods, a time series or panel for the temporal ones.	required
method	str	A registered discovery method name, as in :func:discover.	'pc'
n_bootstrap	int	How many resamples to draw, at least 1.	100
seed	int	Seed for the resampling generator.	0
knowledge	DomainKnowledge | None	Hard constraints, handed to every replicate.	None
block_length	int | None	Moving-block length for a time series under a temporal method; None picks max(ceil(sqrt(T)), 2 * (max_lag + 1)).	None
**kwargs	Any	Passed through to discover and on to the algorithm.	{}

Source code in src/ergodic/discovery/_bootstrap.py

def bootstrap(
    data: TabularData | TimeSeriesData | PanelData,
    *,
    method: str = "pc",
    n_bootstrap: int = 100,
    seed: int = 0,
    knowledge: DomainKnowledge | None = None,
    block_length: int | None = None,
    **kwargs: Any,
) -> BootstrapResult:
    """Rerun a discovery method on bootstrap resamples and count what survives.

    Answers the question a single run cannot: how stable is this edge under
    resampling? Each replicate draws a resample, runs
    ``discover(sample, method=method, knowledge=knowledge, **kwargs)``, and
    the result aggregates adjacency and glyph frequencies over the runs that
    finished. Extra keyword arguments flow to the algorithm unchanged
    (``alpha``, ``max_cond_size``, ``score=...``, ``max_lag=...``), so any
    registered method, iid or temporal, sits behind the same call.

    An iid method on a :class:`~ergodic.data.TabularData` draws
    ``data.n_rows`` rows with replacement, the classic bootstrap;
    ``block_length`` does not apply there and raises. A temporal method
    (``"pcmci"``) never resamples rows, which would shred the serial
    dependence it reads. On a :class:`~ergodic.data.TimeSeriesData` it runs a
    moving-block bootstrap: ``ceil(T / L)`` blocks of ``L`` consecutive rows
    (``block_length``, default ``max(ceil(sqrt(T)), 2 * (max_lag + 1))``) are
    drawn with replacement and enter the replicate as separate entities of a
    panel, ``"block_0"``, ``"block_1"``, ... with fresh time stamps.
    Blocks-as-entities is the point: gluing blocks back into one series would
    build lagged windows that straddle two unrelated fragments and let fake
    adjacencies in, while per-entity pooling never reads across a block seam.
    On a :class:`~ergodic.data.PanelData` it runs the cluster bootstrap
    instead, drawing ``n_entities`` entities whole with replacement
    (relabeled in draw order so repeated draws stay distinct);
    ``block_length`` does not apply to a panel and raises. One honesty note:
    block resampling treats the series as stationary at the block scale, so a
    trend or a regime change inside the sampling window leaks into the
    frequencies. An iid method on temporal data still refuses, inside the
    replicate, exactly as :func:`discover` does.

    The frequencies measure resampling variability, not Bayesian uncertainty;
    they are not posterior probabilities over graphs (graph posteriors belong
    to a later phase). The run is deterministic: the same data, seed, and
    settings give the same result. A temporal method keys the frequencies on
    window-graph glyphs like ``"X[t-1] --> Y[t]"``.

    A replicate that raises ``ValueError`` or ``RuntimeError`` (degenerate
    resamples happen: a constant column, a knowledge conflict on noisy data)
    counts in ``n_failed`` and is skipped; frequencies are over the successful
    runs only. When every replicate fails, the wrapper raises ``RuntimeError``
    chaining the last error.

    Args:
        data: The dataset to resample: tabular for the iid methods, a time
            series or panel for the temporal ones.
        method: A registered discovery method name, as in :func:`discover`.
        n_bootstrap: How many resamples to draw, at least 1.
        seed: Seed for the resampling generator.
        knowledge: Hard constraints, handed to every replicate.
        block_length: Moving-block length for a time series under a temporal
            method; ``None`` picks ``max(ceil(sqrt(T)), 2 * (max_lag + 1))``.
        **kwargs: Passed through to ``discover`` and on to the algorithm.
    """
    if n_bootstrap < 1:
        raise ValueError(f"n_bootstrap must be >= 1, got {n_bootstrap}")
    temporal: TimeSeriesData | PanelData | None = None
    if method.lower() in TEMPORAL_METHODS:
        if not isinstance(data, TimeSeriesData | PanelData):
            raise ValueError(
                f"{method!r} reads temporal data, so its moving-block bootstrap needs "
                f"TimeSeriesData or PanelData; got {type(data).__name__}"
            )
        temporal = data
    elif block_length is not None:
        raise ValueError(
            f"block_length is for temporal data under a temporal method; {method!r} resamples rows"
        )
    max_lag = int(kwargs.get("max_lag", 1))  # PCMCI's own default
    rng = np.random.default_rng(seed)
    n = data.n_rows
    edge_counts: dict[frozenset[str], int] = {}
    mark_counts: dict[str, int] = {}
    n_runs = 0
    n_failed = 0
    last_error: Exception | None = None
    for _ in range(n_bootstrap):
        if temporal is not None:
            sample: TabularData | TimeSeriesData | PanelData = _temporal_sample(
                temporal, rng, block_length=block_length, max_lag=max_lag
            )
        else:
            idx = rng.integers(0, n, size=n)
            sample = data.take(idx)
        try:
            result = discover(sample, method=method, knowledge=knowledge, **kwargs)
        except (ValueError, RuntimeError) as error:
            n_failed += 1
            last_error = error
            continue
        n_runs += 1
        for edge in result.graph.edges:
            pair = frozenset((edge.u, edge.v))
            edge_counts[pair] = edge_counts.get(pair, 0) + 1
            glyph = edge.glyph()
            mark_counts[glyph] = mark_counts.get(glyph, 0) + 1
    if n_runs == 0:
        raise RuntimeError(
            f"every bootstrap replicate failed ({n_failed} of {n_bootstrap}); "
            "the chained error is the last failure"
        ) from last_error
    edges = {
        pair: count / n_runs
        for pair, count in sorted(edge_counts.items(), key=lambda item: tuple(sorted(item[0])))
    }
    marks = {glyph: count / n_runs for glyph, count in sorted(mark_counts.items())}
    return BootstrapResult(
        method=method,
        n_runs=n_runs,
        n_failed=n_failed,
        edges=edges,
        marks=marks,
        variables=data.variables,
    )

discover(data, method='pc', *, test=None, score=None, knowledge=None, alpha=0.05, **kwargs)

Learn causal structure from data with the named method.

test feeds a constraint method (a :class:CITest instance, bound or not, or an alias like "linear"); score feeds a score-based method (a :class:Score instance or an alias like "bic", "bdeu"); both default to "auto", picked from the column schema. Each algorithm declares which ingredient it consumes. Extra keyword arguments pass to the algorithm (max_cond_size, max_parents, node_limit, max_degree, tabu, max_lag).

The iid methods read a :class:TabularData and reject temporal shapes (rows are not exchangeable); method="pcmci" reads a :class:TimeSeriesData or :class:PanelData and returns a window graph over lagged nodes like "X[t-1]".

A few methods consume neither ingredient because the machinery is the algorithm: "lingam" (regression and residual independence) and "glasso" (a sparse precision estimate). They reject test= and score=, and alpha does not apply (their own knobs, like lam or prune, pass through the keyword arguments).

Source code in src/ergodic/discovery/_discover.py

def discover(
    data: TabularData | TimeSeriesData | PanelData,
    method: str = "pc",
    *,
    test: Any = None,
    score: Any = None,
    knowledge: DomainKnowledge | None = None,
    alpha: float = 0.05,
    **kwargs: Any,
) -> DiscoveryResult:
    """Learn causal structure from ``data`` with the named ``method``.

    ``test`` feeds a constraint method (a :class:`CITest` instance, bound or
    not, or an alias like ``"linear"``); ``score`` feeds a score-based method
    (a :class:`Score` instance or an alias like ``"bic"``, ``"bdeu"``); both
    default to ``"auto"``, picked from the column schema. Each algorithm
    declares which ingredient it consumes. Extra keyword arguments pass to
    the algorithm (``max_cond_size``, ``max_parents``, ``node_limit``,
    ``max_degree``, ``tabu``, ``max_lag``).

    The iid methods read a :class:`TabularData` and reject temporal shapes
    (rows are not exchangeable); ``method="pcmci"`` reads a
    :class:`TimeSeriesData` or :class:`PanelData` and returns a window graph
    over lagged nodes like ``"X[t-1]"``.

    A few methods consume neither ingredient because the machinery is the
    algorithm: ``"lingam"`` (regression and residual independence) and
    ``"glasso"`` (a sparse precision estimate). They reject ``test=`` and
    ``score=``, and ``alpha`` does not apply (their own knobs, like ``lam``
    or ``prune``, pass through the keyword arguments).
    """
    cls = algorithm_for(method)
    result: DiscoveryResult
    if cls.ingredient is None:
        if test is not None or score is not None:
            raise ValueError(f"method {method!r} consumes neither a test nor a score")
        result = cls(**kwargs).fit(data, knowledge=knowledge)
        return result
    if cls.ingredient == "test":
        if score is not None:
            raise ValueError(f"method {method!r} consumes a test, not a score")
        chosen_test = test if test is not None else "auto"
        result = cls(test=chosen_test, alpha=alpha, **kwargs).fit(data, knowledge=knowledge)
        return result
    if test is not None:
        raise ValueError(f"method {method!r} consumes a score, not a test")
    chosen_score = score if score is not None else "auto"
    result = cls(score=chosen_score, **kwargs).fit(data, knowledge=knowledge)
    return result

granger(data, *, max_lag=1, alpha=0.05)

Granger-test every ordered pair of variables in data.

For each effect Y, the full regression puts Y[t] on an intercept plus every variable's lags 1..max_lag; the restricted one drops the candidate cause's lags. All variables must be continuous. See the module docstring for what this does and does not mean.

Source code in src/ergodic/discovery/_granger.py

def granger(data: TimeSeriesData, *, max_lag: int = 1, alpha: float = 0.05) -> GrangerResult:
    """Granger-test every ordered pair of variables in ``data``.

    For each effect ``Y``, the full regression puts ``Y[t]`` on an intercept
    plus every variable's lags ``1..max_lag``; the restricted one drops the
    candidate cause's lags. All variables must be continuous. See the module
    docstring for what this does and does not mean.
    """
    if not isinstance(data, TimeSeriesData):
        raise TypeError(
            f"granger needs TimeSeriesData, got {type(data).__name__}; "
            "build one with ergodic.data.timeseries"
        )
    if not 0.0 < alpha < 1.0:
        raise ValueError(f"alpha must be in (0, 1), got {alpha}")
    if max_lag < 1:
        raise ValueError(f"max_lag must be >= 1, got {max_lag}")
    offenders = [name for name in data.variables if not data.schema[name].is_continuous]
    if offenders:
        raise ValueError(f"granger needs continuous variables; not continuous: {offenders}")
    frame = lagged_tabular(data, max_lag).dropna()
    variables = sorted(data.variables)
    lag_owners = [v for v in variables for _ in range(1, max_lag + 1)]
    lag_columns = [lagged_node(v, k) for v in variables for k in range(1, max_lag + 1)]
    n = frame.n_rows
    n_params = 1 + len(lag_columns)
    if n < n_params + 2:
        raise ValueError(
            f"need at least {n_params + 2} complete rows for {n_params} regressors, got {n}"
        )
    full = np.hstack([np.ones((n, 1), dtype=np.float64), frame.matrix(lag_columns)])
    pvalues: dict[tuple[str, str], float] = {}
    for effect in variables:
        response = frame.matrix([lagged_node(effect, 0)])[:, 0]
        rss_full = math.log(max(_rss(full, response), _RSS_FLOOR))
        for cause in variables:
            if cause == effect:
                continue
            keep = [0] + [i + 1 for i, owner in enumerate(lag_owners) if owner != cause]
            rss_restricted = math.log(max(_rss(full[:, keep], response), _RSS_FLOOR))
            statistic = max(0.0, n * (rss_restricted - rss_full))
            pvalues[(cause, effect)] = _chi2_sf(statistic, float(max_lag))
    significant = frozenset(pair for pair, p in pvalues.items() if p <= alpha)
    return GrangerResult(
        pvalues=pvalues, significant=significant, max_lag=max_lag, alpha=alpha, n=n
    )

lagged_node(name, lag)

The window-column name for name at lag steps back.

Lag zero reads "X[t]", lag k reads "X[t-k]". The name must carry no whitespace, since the graph layer parses edges on whitespace.

Source code in src/ergodic/discovery/_lagged.py

def lagged_node(name: str, lag: int) -> str:
    """The window-column name for ``name`` at ``lag`` steps back.

    Lag zero reads ``"X[t]"``, lag ``k`` reads ``"X[t-k]"``. The name must
    carry no whitespace, since the graph layer parses edges on whitespace.
    """
    if any(ch.isspace() for ch in name):
        raise ValueError(f"variable name {name!r} contains whitespace; lagged names cannot")
    if lag < 0:
        raise ValueError(f"lag must be >= 0, got {lag}")
    if lag == 0:
        return f"{name}[t]"
    return f"{name}[t-{lag}]"

lagged_tabular(data, max_lag)

Stack the windows of data into a tabular frame of lagged columns.

Row t holds every variable at lags 0..max_lag; the column for variable v at lag k is named lagged_node(v, k) and holds v at time t - k. Column types carry over: continuous stays continuous, categorical keeps its categories, ordinal keeps its level order.

A time series needs at least max_lag + 8 rows. Panel windows are built within each entity (sorted entity order) and stacked; entities with fewer than max_lag + 2 rows are skipped with a warning, and it is an error when no rows survive. Unequal time steps warn too: lags count rows, not calendar time.

Source code in src/ergodic/discovery/_lagged.py

def lagged_tabular(data: TimeSeriesData | PanelData, max_lag: int) -> TabularData:
    """Stack the windows of ``data`` into a tabular frame of lagged columns.

    Row ``t`` holds every variable at lags ``0..max_lag``; the column for
    variable ``v`` at lag ``k`` is named ``lagged_node(v, k)`` and holds ``v``
    at time ``t - k``. Column types carry over: continuous stays continuous,
    categorical keeps its categories, ordinal keeps its level order.

    A time series needs at least ``max_lag + 8`` rows. Panel windows are
    built within each entity (sorted entity order) and stacked; entities with
    fewer than ``max_lag + 2`` rows are skipped with a warning, and it is an
    error when no rows survive. Unequal time steps warn too: lags count rows,
    not calendar time.
    """
    if max_lag < 1:
        raise ValueError(f"max_lag must be >= 1, got {max_lag}")
    if isinstance(data, PanelData):
        columns, irregular, skipped = _panel_columns(data, max_lag)
        if skipped:
            warnings.warn(
                f"skipped {skipped} of {data.n_entities} entities with fewer than "
                f"max_lag + 2 = {max_lag + 2} rows",
                UserWarning,
                stacklevel=2,
            )
    elif isinstance(data, TimeSeriesData):
        columns, irregular = _series_columns(data, max_lag)
    else:
        raise TypeError(
            f"lagged_tabular needs TimeSeriesData or PanelData, got {type(data).__name__}"
        )
    if irregular:
        warnings.warn(
            "time steps are not equally spaced; lags are positional (rows back), not calendar",
            UserWarning,
            stacklevel=2,
        )
    return _typed_frame(data, columns, max_lag)

parse_lagged(node)

Invert :func:lagged_node: "X[t-2]" becomes ("X", 2).

Source code in src/ergodic/discovery/_lagged.py

def parse_lagged(node: str) -> tuple[str, int]:
    """Invert :func:`lagged_node`: ``"X[t-2]"`` becomes ``("X", 2)``."""
    cut = node.rfind("[t")
    if cut > 0 and node.endswith("]"):
        base, suffix = node[:cut], node[cut:]
        if suffix == "[t]":
            return base, 0
        digits = suffix[3:-1]
        canonical = digits.isdigit() and str(int(digits)) == digits and int(digits) >= 1
        if suffix.startswith("[t-") and canonical:
            return base, int(digits)
    raise ValueError(f"{node!r} is not a lagged node name; expected 'name[t]' or 'name[t-k]'")

Causal inference

Effect estimation: the estimators, the learners that fill their slots, refutation, and the quasi-experimental designs.

ergodic.inference

Causal inference: estimate causal effects from data and a causal model.

An estimator factors into independent pieces: the estimand family (from :mod:ergodic.identification), the strategy that turns it into a number (a :class:CausalEstimator), and the learner in each nuisance slot. estimate_effect runs identification then estimation and returns an :class:EffectEstimate.

>>> from ergodic import dag, tabular
>>> from ergodic.inference import estimate_effect
>>> est = estimate_effect(graph, data, "T", "Y")          # doctest: +SKIP
>>> est.value, est.interval()                              # doctest: +SKIP

The learner is structurally a sklearn estimator, so estimate_effect(..., learner="knn") or a slot mapping like learner={"outcome": GradientBoostingRegressor()} drops any sklearn or xgboost model into the slots with no adapter. The result is distribution-shaped: a point estimate with an analytic standard error, or bootstrap draws when there is no closed form.

The quasi-experimental designs read panels and time series instead of a graph: :func:did and :func:event_study over a :class:~ergodic.data.PanelData (parallel trends), :func:synthetic_control (a convex donor combination), and :func:interrupted_time_series over a :class:~ergodic.data.TimeSeriesData (a projected pre-trend). Each states its identification in a :class:Design record and returns the same :class:EffectEstimate surface.

AIPW

Bases: _AdjustmentBase

Combine an outcome model and a propensity model.

Consistent if either model is right. The ATE standard error comes from the efficient influence function (Neyman-orthogonal, so first-order insensitive to nuisance error); the ATT standard error falls back to bootstrap.

Source code in src/ergodic/inference/estimators/_doubly_robust.py

class AIPW(_AdjustmentBase):
    """Combine an outcome model and a propensity model.

    Consistent if either model is right. The ATE standard error comes from the
    efficient influence function (Neyman-orthogonal, so first-order insensitive to
    nuisance error); the ATT standard error falls back to bootstrap.
    """

    name = "aipw"
    slots = (Slot("outcome", Task.REGRESSION), Slot("propensity", Task.CLASSIFICATION))
    targets = frozenset({EffectTarget.ATE, EffectTarget.ATT})

    def __init__(
        self,
        learners: Any = None,
        *,
        clip: float = 1e-3,
        bayesian_bootstrap: bool = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(learners, **kwargs)
        self.clip = clip
        self.bayesian_bootstrap = bayesian_bootstrap

    def _nuisance(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        n = a.shape[0]
        outcome = make["outcome"]().fit(np.hstack([a.reshape(-1, 1), z]), y)
        ones = np.ones((n, 1), dtype=np.float64)
        zeros = np.zeros((n, 1), dtype=np.float64)
        mu1 = _mean_predict(outcome, np.hstack([ones, z]))
        mu0 = _mean_predict(outcome, np.hstack([zeros, z]))
        propensity = make["propensity"]().fit(z, a)
        e = np.clip(_mean_predict(propensity, z), self.clip, 1.0 - self.clip)
        return mu0, mu1, e

    def _scores(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> npt.NDArray[np.float64]:
        a = arrays["A"]
        y = arrays["Y"]
        mu0, mu1, e = self._nuisance(arrays, make)
        if target is EffectTarget.ATT:
            p1 = float(np.mean(a))
            if p1 <= 0.0:
                raise ValueError("no treated units for ATT")
            return _aipw_att_scores(a, y, mu0, e, p1)
        return _aipw_ate_scores(a, y, mu0, mu1, e)

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        return float(np.mean(self._scores(arrays, target, make)))

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        if target is not EffectTarget.ATE:
            return None
        psi = self._scores(arrays, target, make)
        return float(np.mean(psi)), psi

    def _posterior(self, state: _FitState, target: EffectTarget) -> npt.NDArray[np.float64] | None:
        """The doubly-robust posterior: the influence scores per nuisance draw.

        Fit the outcome and propensity once, read their ``(S, n)`` draws, and for
        each draw evaluate the same influence scores as the point path. With
        ``bayesian_bootstrap`` (the default) reweight the rows by a Dirichlet(1)
        draw, so the spread carries the nuisance posterior and the sampling
        uncertainty together. Needs both nuisances Bayesian; a mix returns ``None``.
        """
        if not all(state.bayesian.values()):
            return None
        arrays = state.arrays
        a, y, z = arrays["A"], arrays["Y"], arrays["Z"]
        n = a.shape[0]
        ones = np.ones((n, 1), dtype=np.float64)
        zeros = np.zeros((n, 1), dtype=np.float64)
        outcome = state.make["outcome"]().fit(np.hstack([a.reshape(-1, 1), z]), y)
        mu1d = _mean_predict_draws(outcome, np.hstack([ones, z]))
        mu0d = _mean_predict_draws(outcome, np.hstack([zeros, z]))
        propensity = state.make["propensity"]().fit(z, a)
        ed = np.clip(_mean_predict_draws(propensity, z), self.clip, 1.0 - self.clip)
        n_draws = min(mu1d.shape[0], mu0d.shape[0], ed.shape[0])
        p1 = float(np.mean(a))
        if target is EffectTarget.ATT and p1 <= 0.0:
            raise ValueError("no treated units for ATT")
        rng = np.random.default_rng(self._random_state)
        draws = np.empty(n_draws, dtype=np.float64)
        for s in range(n_draws):
            w = _dirichlet_weights(n, rng) if self.bayesian_bootstrap else None
            if target is EffectTarget.ATT:
                if w is None:
                    phi = _aipw_att_scores(a, y, mu0d[s], ed[s], p1)
                else:
                    # the treated share must follow the same reweighted measure
                    p1_w = float(np.sum(w * a))
                    if p1_w <= 0.0:
                        raise ValueError("a Bayesian-bootstrap draw left no treated weight for ATT")
                    phi = _aipw_att_scores(a, y, mu0d[s], ed[s], p1_w)
            else:
                phi = _aipw_ate_scores(a, y, mu0d[s], mu1d[s], ed[s])
            draws[s] = float(np.mean(phi)) if w is None else float(np.sum(w * phi))
        return draws

BayesianLearner

Bases: Protocol

A learner that also carries a posterior over its mean function.

Structurally a :class:Learner (fit then predict, where predict is the posterior predictive mean), plus predict_draws returning (S, n) draws of the mean function. A classifier variant exposes predict_proba and predict_proba_draws (draws of the class-1 probability) instead. Detection is duck-typed in :func:_is_bayesian, the same way a classifier is spotted by predict_proba, so any object that quacks this way counts.

Because predict stays the posterior mean, a Bayesian learner is also a valid point learner in every estimator; the posterior path is an addition, not a fork.

Source code in src/ergodic/inference/_learner.py

@runtime_checkable
class BayesianLearner(Protocol):
    """A learner that also carries a posterior over its mean function.

    Structurally a :class:`Learner` (``fit`` then ``predict``, where ``predict`` is
    the posterior predictive mean), plus ``predict_draws`` returning ``(S, n)``
    draws of the mean function. A classifier variant exposes ``predict_proba`` and
    ``predict_proba_draws`` (draws of the class-1 probability) instead. Detection is
    duck-typed in :func:`_is_bayesian`, the same way a classifier is spotted by
    ``predict_proba``, so any object that quacks this way counts.

    Because ``predict`` stays the posterior mean, a Bayesian learner is also a valid
    point learner in every estimator; the posterior path is an addition, not a fork.
    """

    def fit(
        self,
        X: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        sample_weight: npt.NDArray[np.float64] | None = None,
    ) -> Any: ...

    def predict(self, X: npt.NDArray[np.float64]) -> Any: ...

    def predict_draws(self, X: npt.NDArray[np.float64]) -> Any: ...

CATEResult dataclass

Bases: PerUnitEffect

Conditional average treatment effects, with an explainability hook.

A conditional average over the unobserved noise (the inference-pillar object). Carries the :class:~ergodic.identification.Identification it was fit under and the cached doubly-robust scores + covariate design, so :meth:explain projects honest summaries (BLP, GATES, CLAN, the calibration test, RATE) off the same orthogonal signal.

Source code in src/ergodic/inference/_perunit.py

@dataclass(frozen=True, eq=False)
class CATEResult(PerUnitEffect):
    """Conditional average treatment effects, with an explainability hook.

    A conditional average over the unobserved noise (the inference-pillar object).
    Carries the :class:`~ergodic.identification.Identification` it was fit under and
    the cached doubly-robust scores + covariate design, so :meth:`explain` projects
    honest summaries (BLP, GATES, CLAN, the calibration test, RATE) off the same
    orthogonal signal.
    """

    estimand: str = "cate"
    identification: Identification | None = None
    _scores: DRScores | None = field(default=None, repr=False)

    def explain(self) -> Heterogeneity:
        """Tier-B characterization of the heterogeneity, on the doubly-robust scores.

        Builds a :class:`~ergodic.inference.explain.Heterogeneity` record whose
        methods (``blp``, ``gates``, ``clan``, ``calibration_test``, ``rate``) all
        consume the cached cross-fit scores, the covariate design, and the fitted
        per-unit ``tau_hat``. Needs the scores; raises when they were not cached
        (an estimator that does not build them, or a result without them).
        """
        from .explain import Heterogeneity

        if self._scores is None:
            raise RuntimeError(
                "this CATEResult carries no doubly-robust scores to explain; "
                "explainability needs the cross-fit AIPW signal (the cate() facade "
                "and the doubly-robust meta-learners cache it)"
            )
        # the out-of-fold proxy ranking when available (the honesty recipe: the
        # ranking is fixed on data not used for the summary), else the in-sample fit
        tau_hat = self._scores.tau_hat if self._scores.tau_hat is not None else self.values
        return Heterogeneity(
            scores=self._scores,
            tau_hat=np.asarray(tau_hat, dtype=np.float64),
            treatment=self.treatment,
            outcome=self.outcome,
        )

explain()

Tier-B characterization of the heterogeneity, on the doubly-robust scores.

Builds a :class:~ergodic.inference.explain.Heterogeneity record whose methods (blp, gates, clan, calibration_test, rate) all consume the cached cross-fit scores, the covariate design, and the fitted per-unit tau_hat. Needs the scores; raises when they were not cached (an estimator that does not build them, or a result without them).

Source code in src/ergodic/inference/_perunit.py

def explain(self) -> Heterogeneity:
    """Tier-B characterization of the heterogeneity, on the doubly-robust scores.

    Builds a :class:`~ergodic.inference.explain.Heterogeneity` record whose
    methods (``blp``, ``gates``, ``clan``, ``calibration_test``, ``rate``) all
    consume the cached cross-fit scores, the covariate design, and the fitted
    per-unit ``tau_hat``. Needs the scores; raises when they were not cached
    (an estimator that does not build them, or a result without them).
    """
    from .explain import Heterogeneity

    if self._scores is None:
        raise RuntimeError(
            "this CATEResult carries no doubly-robust scores to explain; "
            "explainability needs the cross-fit AIPW signal (the cate() facade "
            "and the doubly-robust meta-learners cache it)"
        )
    # the out-of-fold proxy ranking when available (the honesty recipe: the
    # ranking is fixed on data not used for the summary), else the in-sample fit
    tau_hat = self._scores.tau_hat if self._scores.tau_hat is not None else self.values
    return Heterogeneity(
        scores=self._scores,
        tau_hat=np.asarray(tau_hat, dtype=np.float64),
        treatment=self.treatment,
        outcome=self.outcome,
    )

CausalEstimator

Bases: ABC

Base for the strategies. Stateful: fit then effect.

Source code in src/ergodic/inference/_estimator.py

class CausalEstimator(ABC):
    """Base for the strategies. Stateful: ``fit`` then ``effect``."""

    name: ClassVar[str]
    consumes: ClassVar[frozenset[str]]
    slots: ClassVar[tuple[Slot, ...]] = ()
    targets: ClassVar[frozenset[EffectTarget]]
    kind: ClassVar[str] = "analytic"

    def __init__(
        self,
        learners: Any = None,
        *,
        random_state: int | None = None,
        n_boot: int = 1000,
        treated: Any | None = None,
    ) -> None:
        self._spec = learners
        self._random_state = random_state
        self._n_boot = n_boot
        self._treated = treated
        self._fit_state: _FitState | None = None
        self._covariates: tuple[str, ...] = ()
        self._identification: Identification | None = None
        self._treatment_name: str = ""
        self._outcome_name: str = ""

    # -- learner resolution --------------------------------------------------

    def _slot_specs(self) -> dict[str, Any]:
        names = [s.name for s in self.slots]
        spec = self._spec
        if spec is None:
            return dict.fromkeys(names)
        if isinstance(spec, Mapping):
            unknown = set(spec) - set(names)
            if unknown:
                raise ValueError(f"{self.name} has no slot(s) {sorted(unknown)}; slots are {names}")
            return {name: spec.get(name) for name in names}
        return dict.fromkeys(names, spec)

    def _default_learner(self, slot: Slot, y_binary: bool) -> Learner:
        if slot.task is Task.CLASSIFICATION:
            return LogisticRegression()
        if slot.name == "outcome" and y_binary:
            return LogisticRegression()
        return OLS()

    def _resolve(
        self, y_binary: bool
    ) -> tuple[dict[str, Callable[[], Learner]], dict[str, str], dict[str, bool]]:
        specs = self._slot_specs()
        make: dict[str, Callable[[], Learner]] = {}
        names: dict[str, str] = {}
        bayesian: dict[str, bool] = {}
        for slot in self.slots:
            spec = specs[slot.name]
            want_classifier = slot.task is Task.CLASSIFICATION
            learner = (
                self._default_learner(slot, y_binary)
                if spec is None
                else _as_learner(spec, want_classifier=want_classifier)
            )
            if slot.task is Task.CLASSIFICATION and not _is_classifier(learner):
                raise TypeError(
                    f"slot {slot.name!r} needs a classifier with predict_proba; "
                    f"got {type(learner).__name__}"
                )
            make[slot.name] = partial(_fresh, learner)
            names[slot.name] = type(learner).__name__
            bayesian[slot.name] = _is_bayesian(learner)
        return make, names, bayesian

    # -- fit and effect ------------------------------------------------------

    def fit(self, identification: Identification, data: TabularData) -> Self:
        """Read the estimand and the data, resolve learners, store the fit state."""
        if not identification.identifiable:
            raise ValueError("effect is not identifiable; nothing to estimate")
        if identification.strategy not in self.consumes:
            raise ValueError(
                f"{self.name} consumes {sorted(self.consumes)}, "
                f"but the effect was identified by {identification.strategy!r}"
            )
        self._identification = identification
        arrays, y_binary, diagnostics = self._extract(identification, data)
        make, learner_names, bayesian = self._resolve(y_binary)
        self._fit_state = _FitState(
            arrays=arrays,
            make=make,
            learner_names=learner_names,
            n=data.n_rows,
            diagnostics=diagnostics,
            bayesian=bayesian,
        )
        return self

    @property
    def _state(self) -> _FitState:
        if self._fit_state is None:
            raise RuntimeError(f"{self.name} is not fitted; call fit first")
        return self._fit_state

    def _cate_covariates(self, data: TabularData | None) -> npt.NDArray[np.float64]:
        """The covariate design for CATE prediction: in-sample, or extracted from ``data``."""
        if data is None:
            return self._state.arrays["Z"]
        return _covariate_matrix(data, list(self._covariates))

    def _treatment_vector(
        self, data: TabularData, name: str, *, allow_continuous: bool = False
    ) -> npt.NDArray[np.float64]:
        """The treatment as floats, honoring ``treated`` and warning on an implicit coding."""
        if self._treated is not None:
            return _coded_treatment(data, name, self._treated)
        col = data.schema[name]
        if col.is_discrete and col.cardinality == 2:
            cats = tuple(col.categories or ())
            # {0, 1} also admits booleans, since False == 0 and True == 1
            if not set(cats) <= {0, 1}:
                warnings.warn(
                    f"treatment {name!r} was coded 1 for level {cats[1]!r} and 0 for "
                    f"{cats[0]!r} by level order; pass treated=... to choose",
                    UserWarning,
                    stacklevel=2,
                )
        if allow_continuous:
            return _numeric_vector(data, name, "treatment")
        return _binary_vector(data, name, "treatment")

    def effect(self, *, target: EffectTarget | str | None = None) -> EffectEstimate:
        """Report the effect, with an analytic standard error or a bootstrap one."""
        from ._result import EffectEstimate

        chosen = self._pick_target(target)
        state = self._state
        if self._has_bayesian_slot(state):
            return self._bayesian_effect(state, chosen)
        analytic = self._analytic(state.arrays, chosen, state.make)
        if analytic is not None:
            point, influence = analytic
            se = float(np.std(influence, ddof=1) / np.sqrt(influence.shape[0]))
            return EffectEstimate(
                target=chosen,
                estimator=self.name,
                learners=dict(state.learner_names),
                n=state.n,
                kind="analytic",
                diagnostics=dict(state.diagnostics),
                _point=point,
                _std_error=se,
            )
        rng = np.random.default_rng(self._random_state)
        point = self._point(state.arrays, chosen, state.make)
        draws = self._bootstrap(state, chosen, rng)
        diagnostics = {
            **state.diagnostics,
            "n_boot_requested": float(self._n_boot),
            "n_boot_returned": float(draws.shape[0]),
        }
        return EffectEstimate(
            target=chosen,
            estimator=self.name,
            learners=dict(state.learner_names),
            n=state.n,
            kind="bootstrap",
            diagnostics=diagnostics,
            _point=point,
            _draws=draws,
        )

    # -- the posterior path --------------------------------------------------

    def _has_bayesian_slot(self, state: _FitState) -> bool:
        """True when any resolved slot carries a posterior."""
        return any(state.bayesian.values())

    def _bayesian_effect(self, state: _FitState, target: EffectTarget) -> EffectEstimate:
        """Estimate with a Bayesian slot present, never bootstrapping.

        A bootstrap loop would refit the (expensive) Bayesian nuisance once per
        resample, so this path avoids it. If the estimator implements ``_posterior``
        for this target, return the posterior. Otherwise fall back to the analytic
        result when one exists, else a single point with no interval, with a note in
        the diagnostics. The point and the analytic both fit the nuisances once and
        read their posterior means.
        """
        from ._result import EffectEstimate

        draws = self._posterior(state, target)
        if draws is not None:
            return EffectEstimate(
                target=target,
                estimator=self.name,
                learners=dict(state.learner_names),
                n=state.n,
                kind="posterior",
                diagnostics=dict(state.diagnostics),
                _draws=draws,
            )
        diagnostics = {**state.diagnostics, "posterior_available": 0.0}
        analytic = self._analytic(state.arrays, target, state.make)
        if analytic is not None:
            point, influence = analytic
            se = float(np.std(influence, ddof=1) / np.sqrt(influence.shape[0]))
            return EffectEstimate(
                target=target,
                estimator=self.name,
                learners=dict(state.learner_names),
                n=state.n,
                kind="analytic",
                diagnostics=diagnostics,
                _point=point,
                _std_error=se,
            )
        point = self._point(state.arrays, target, state.make)
        return EffectEstimate(
            target=target,
            estimator=self.name,
            learners=dict(state.learner_names),
            n=state.n,
            kind="point",
            diagnostics=diagnostics,
            _point=point,
        )

    def _pick_target(self, target: EffectTarget | str | None) -> EffectTarget:
        if target is None:
            if EffectTarget.ATE in self.targets:
                return EffectTarget.ATE
            return sorted(self.targets, key=str)[0]
        chosen = as_target(target)
        if chosen not in self.targets:
            raise ValueError(
                f"{self.name} reports {sorted(str(t) for t in self.targets)}, not {chosen}"
            )
        return chosen

    def _bootstrap(
        self, state: _FitState, target: EffectTarget, rng: np.random.Generator
    ) -> npt.NDArray[np.float64]:
        n = state.n
        draws = np.empty(self._n_boot, dtype=np.float64)
        filled = 0
        attempts = 0
        budget = self._n_boot * 5
        while filled < self._n_boot and attempts < budget:
            attempts += 1
            idx = rng.integers(0, n, n)
            sub = {key: arr[idx] for key, arr in state.arrays.items()}
            try:
                value = self._point(sub, target, state.make)
            except (ValueError, ZeroDivisionError, np.linalg.LinAlgError):
                continue
            if not np.isfinite(value):
                continue
            draws[filled] = value
            filled += 1
        if filled == 0:
            raise RuntimeError(
                f"bootstrap for {self.name} produced no finite draws in {attempts} attempts; "
                "every resample failed to estimate"
            )
        return draws[:filled]

    # -- per-estimator -------------------------------------------------------

    @abstractmethod
    def _extract(
        self, identification: Identification, data: TabularData
    ) -> tuple[dict[str, npt.NDArray[np.float64]], bool, dict[str, float]]:
        """Pull the arrays the estimator needs, the outcome-is-binary flag, diagnostics."""

    @abstractmethod
    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        """Fit fresh learners on ``arrays`` and return the effect estimate."""

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        """Return ``(point, influence)`` when an influence function is known, else ``None``."""
        return None

    def _posterior(self, state: _FitState, target: EffectTarget) -> npt.NDArray[np.float64] | None:
        """Posterior draws of the effect when the estimator supports them.

        Called only when a slot is Bayesian. The default is ``None`` (no posterior
        path for this estimator or target), which sends ``_bayesian_effect`` to its
        analytic-or-point fallback. Implementers fit the nuisances once, read their
        ``predict_draws``, and evaluate the estimator's functional per draw.
        """
        return None

effect(*, target=None)

Report the effect, with an analytic standard error or a bootstrap one.

Source code in src/ergodic/inference/_estimator.py

def effect(self, *, target: EffectTarget | str | None = None) -> EffectEstimate:
    """Report the effect, with an analytic standard error or a bootstrap one."""
    from ._result import EffectEstimate

    chosen = self._pick_target(target)
    state = self._state
    if self._has_bayesian_slot(state):
        return self._bayesian_effect(state, chosen)
    analytic = self._analytic(state.arrays, chosen, state.make)
    if analytic is not None:
        point, influence = analytic
        se = float(np.std(influence, ddof=1) / np.sqrt(influence.shape[0]))
        return EffectEstimate(
            target=chosen,
            estimator=self.name,
            learners=dict(state.learner_names),
            n=state.n,
            kind="analytic",
            diagnostics=dict(state.diagnostics),
            _point=point,
            _std_error=se,
        )
    rng = np.random.default_rng(self._random_state)
    point = self._point(state.arrays, chosen, state.make)
    draws = self._bootstrap(state, chosen, rng)
    diagnostics = {
        **state.diagnostics,
        "n_boot_requested": float(self._n_boot),
        "n_boot_returned": float(draws.shape[0]),
    }
    return EffectEstimate(
        target=chosen,
        estimator=self.name,
        learners=dict(state.learner_names),
        n=state.n,
        kind="bootstrap",
        diagnostics=diagnostics,
        _point=point,
        _draws=draws,
    )

fit(identification, data)

Read the estimand and the data, resolve learners, store the fit state.

Source code in src/ergodic/inference/_estimator.py

def fit(self, identification: Identification, data: TabularData) -> Self:
    """Read the estimand and the data, resolve learners, store the fit state."""
    if not identification.identifiable:
        raise ValueError("effect is not identifiable; nothing to estimate")
    if identification.strategy not in self.consumes:
        raise ValueError(
            f"{self.name} consumes {sorted(self.consumes)}, "
            f"but the effect was identified by {identification.strategy!r}"
        )
    self._identification = identification
    arrays, y_binary, diagnostics = self._extract(identification, data)
    make, learner_names, bayesian = self._resolve(y_binary)
    self._fit_state = _FitState(
        arrays=arrays,
        make=make,
        learner_names=learner_names,
        n=data.n_rows,
        diagnostics=diagnostics,
        bayesian=bayesian,
    )
    return self

CausalForest

Bases: _MetaLearnerBase

An honest causal forest (GRF) for the conditional average treatment effect.

The forest is the CATE learner. The outcome and propensity nuisance slots take any learner (they residualize Y and W on the covariates before the forest grows); the forest machinery, honest splitting and the kernel readout, is native. The ATE is the average per-unit effect, and cate returns a :class:~ergodic.inference.CATEResult with per-unit infinitesimal-jackknife standard errors and intervals.

Hyperparameters: n_trees (forest size), min_leaf (the minimum estimate-half rows per leaf), max_depth, honest (split/estimate sample-splitting, on by default), subsample (the per-tree subsample fraction), and random_state. Binary treatment in this phase; a continuous treatment via the same local moment is deferred and raises.

Source code in src/ergodic/inference/estimators/_forest.py

class CausalForest(_MetaLearnerBase):
    """An honest causal forest (GRF) for the conditional average treatment effect.

    The forest is the CATE learner. The ``outcome`` and ``propensity`` nuisance slots
    take any learner (they residualize ``Y`` and ``W`` on the covariates before the
    forest grows); the forest machinery, honest splitting and the kernel readout, is
    native. The ATE is the average per-unit effect, and ``cate`` returns a
    :class:`~ergodic.inference.CATEResult` with per-unit infinitesimal-jackknife
    standard errors and intervals.

    Hyperparameters: ``n_trees`` (forest size), ``min_leaf`` (the minimum estimate-half
    rows per leaf), ``max_depth``, ``honest`` (split/estimate sample-splitting, on by
    default), ``subsample`` (the per-tree subsample fraction), and ``random_state``.
    Binary treatment in this phase; a continuous treatment via the same local moment
    is deferred and raises.
    """

    name = "causal_forest"
    slots = (
        Slot("outcome", Task.REGRESSION),
        Slot("propensity", Task.CLASSIFICATION),
    )
    targets = frozenset({EffectTarget.ATE})

    def __init__(
        self,
        learners: Any = None,
        *,
        n_trees: int = 200,
        min_leaf: int = 5,
        max_depth: int = 8,
        honest: bool = True,
        subsample: float = 0.5,
        max_features: float | int | None = None,
        alpha: float = 0.05,
        n_folds: int = 5,
        clip: float = 1e-3,
        **kwargs: Any,
    ) -> None:
        super().__init__(learners, **kwargs)
        if not 0.0 < subsample <= 1.0:
            raise ValueError(f"subsample must be in (0, 1], got {subsample}")
        if n_trees < 1:
            raise ValueError(f"n_trees must be at least 1, got {n_trees}")
        if min_leaf < 1:
            raise ValueError(f"min_leaf must be at least 1, got {min_leaf}")
        if not 0.0 <= alpha < 0.5:
            raise ValueError(f"alpha must be in [0, 0.5), got {alpha}")
        self.n_trees = n_trees
        self.min_leaf = min_leaf
        self.max_depth = max_depth
        self.honest = honest
        self.subsample = subsample
        self.max_features = max_features
        self.alpha = alpha
        self.n_folds = n_folds
        self.clip = clip

    # -- residualization (the forest R-learner front end) --------------------

    def _residualize(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        """Cross-fit ``m(x)`` and ``e(x)`` and return the residuals ``W - e``, ``Y - m``.

        Reuses the shared cross-fit nuisance routine: the outcome model carries the
        treatment as a feature (so ``m = e * mu1 + (1 - e) * mu0``), and the propensity
        is the treatment model. The residuals are what the forest splits on.
        """
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        mu0, mu1, e = _crossfit_nuisances(
            a,
            y,
            z,
            make_outcome=make["outcome"],
            make_propensity=make.get("propensity") or (lambda: LogisticRegression()),
            n_folds=self.n_folds,
            clip=self.clip,
            random_state=self._random_state,
        )
        m = e * mu1 + (1.0 - e) * mu0
        return a - e, y - m

    def _max_features_for(self, n_features: int) -> int:
        if self.max_features is None:
            # GRF default: ceil(sqrt(p) + 20), capped at p
            return min(n_features, int(np.ceil(np.sqrt(n_features) + 20)))
        if isinstance(self.max_features, int):
            return max(1, min(n_features, self.max_features))
        return max(1, min(n_features, int(np.ceil(self.max_features * n_features))))

    # -- the meta-learner hooks ----------------------------------------------

    def _fit_cate(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> _ForestModel:
        w_res, y_res = self._residualize(arrays, make)
        x = arrays["Z"]
        n, n_features = x.shape
        rng = np.random.default_rng(self._random_state)
        sub_n = max(2 * self.min_leaf, int(np.ceil(self.subsample * n)))
        sub_n = min(sub_n, n)
        n_try = self._max_features_for(n_features)
        trees: list[_Tree] = []
        importance = np.zeros(n_features, dtype=np.float64)
        for _ in range(self.n_trees):
            in_bag = np.asarray(rng.choice(n, size=sub_n, replace=False), dtype=np.int64)
            if self.honest and in_bag.shape[0] >= 2:
                shuffled = rng.permutation(in_bag)
                half = shuffled.shape[0] // 2
                split_rows = shuffled[:half]
                estimate_rows = shuffled[half:]
            else:
                split_rows = in_bag
                estimate_rows = in_bag
            builder = _TreeBuilder(
                x,
                w_res,
                y_res,
                min_leaf=self.min_leaf,
                max_depth=self.max_depth,
                n_features_try=n_try,
                alpha=self.alpha,
                rng=rng,
            )
            root = builder.build(split_rows, estimate_rows)
            trees.append(_Tree(root=root, in_bag=in_bag, split_importance=builder.importance))
            importance += builder.importance
        total = float(importance.sum())
        if total > 0:
            importance = importance / total
        return _ForestModel(
            trees=trees,
            x_train=x,
            w_res=w_res,
            y_res=y_res,
            importance=importance,
            covariate_names=self._covariates,
        )

    def _predict_cate(
        self, model: _ForestModel, z: npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64]:
        return model.predict(z)

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        """The doubly-robust AIPW ATE with its influence-function standard error.

        GRF averages per-unit AIPW scores for a doubly-robust, efficient ATE rather
        than bootstrapping the forest. The score is
        ``Gamma_i = tau_hat(X_i) + (W_i - e_i) / (e_i (1 - e_i)) (Y_i - mu_hat_i)``
        with ``mu_hat_i = W_i mu1_i + (1 - W_i) mu0_i``; its mean is the ATE and its
        scaled standard deviation the standard error. Reusing this for ``effect()``
        keeps the scalar ATE path off the (prohibitively slow) forest-refit bootstrap.
        Computed only on the fit arrays (``_analytic`` on a resample returns ``None``).
        """
        if target is not EffectTarget.ATE:
            return None
        model = getattr(self, "_fitted_cate", None)
        if model is None or arrays["A"].shape[0] != model.x_train.shape[0]:
            return None
        a, y, z = arrays["A"], arrays["Y"], arrays["Z"]
        try:
            mu0, mu1, e = _crossfit_nuisances(
                a,
                y,
                z,
                make_outcome=make["outcome"],
                make_propensity=make.get("propensity") or (lambda: LogisticRegression()),
                n_folds=self.n_folds,
                clip=self.clip,
                random_state=self._random_state,
            )
        except (ValueError, np.linalg.LinAlgError):
            return None
        tau_hat = model.predict(z)
        mu_hat = a * mu1 + (1.0 - a) * mu0
        gamma = tau_hat + (a - e) / (e * (1.0 - e)) * (y - mu_hat)
        return float(np.mean(gamma)), np.asarray(gamma, dtype=np.float64)

    # -- the per-unit result, with the infinitesimal-jackknife variance ------

    def cate(self, data: TabularData | None = None, *, uncertainty: bool = True) -> CATEResult:
        """Per-unit CATE with infinitesimal-jackknife standard errors and intervals.

        Overrides the base meta-learner ``cate`` (which would bootstrap-refit the
        whole forest). The point is the forest-kernel local effect; the per-unit
        variance is the infinitesimal jackknife over the trees (the
        bootstrap-of-little-bags grouping), valid because the forest is honest and
        subsampled. The result carries the cached doubly-robust scores for
        :meth:`~ergodic.inference.CATEResult.explain`.

        ``uncertainty=False`` returns the point effects with ``kind="point"`` and no
        standard error, skipping the infinitesimal-jackknife pass over the trees.
        """
        from .._perunit import CATEResult

        model = self._cate_model_forest()
        z = self._cate_covariates(data)
        point = np.asarray(model.predict(z), dtype=np.float64)
        n = point.shape[0]
        std_error = self._ij_standard_error(model, z, point) if uncertainty else None
        kind = "analytic" if uncertainty else "point"
        scores = self._cate_scores(data)
        return CATEResult(
            treatment=self._treatment_name,
            outcome=self._outcome_name,
            covariates=self._covariates,
            n=n,
            kind=kind,
            source=self.name,
            diagnostics=dict(self._state.diagnostics),
            identification=self._identification,
            _values=point,
            _index=self._cate_index(data, n),
            _std_error=std_error,
            _scores=scores,
        )

    def _cate_model_forest(self) -> _ForestModel:
        model = getattr(self, "_fitted_cate", None)
        if model is None:
            raise RuntimeError(f"{self.name} is not fitted; call fit first")
        return model  # type: ignore[no-any-return]

    def _ij_standard_error(
        self,
        model: _ForestModel,
        z: npt.NDArray[np.float64],
        point: npt.NDArray[np.float64],
    ) -> npt.NDArray[np.float64]:
        """Per-unit infinitesimal-jackknife standard errors over the trees.

        The IJ variance (Wager-Athey Eq. 8, the bootstrap-of-little-bags form)
        is ``V(x) = ((n - 1)/n)(n/(n - s))^2 sum_i Cov_b[t_b(x), N_ib]^2``, where
        ``t_b(x)`` is tree ``b``'s prediction at ``x``, ``N_ib`` is the inclusion count
        of training row ``i`` in tree ``b`` (here 0/1 under subsampling), and ``s`` is
        the subsample size. A Monte-Carlo correction subtracts the within-tree
        sampling noise so the estimate stays non-negative and is not inflated by a
        finite tree count. Honesty and subsampling are what make this valid.
        """
        n_trees = len(model.trees)
        n_query = z.shape[0]
        n_train = model.x_train.shape[0]
        if n_trees < 2 or n_train < 2:
            return np.zeros(n_query, dtype=np.float64)
        # per-tree predictions, (B, n_query)
        preds = np.empty((n_trees, n_query), dtype=np.float64)
        inbag = np.zeros((n_trees, n_train), dtype=np.float64)
        for b, tree in enumerate(model.trees):
            preds[b] = model.tree_predict(tree, z)
            inbag[b, tree.in_bag] = 1.0
        s = float(inbag.sum(axis=1).mean())
        pred_centered = preds - preds.mean(axis=0, keepdims=True)
        inbag_centered = inbag - inbag.mean(axis=0, keepdims=True)
        # Cov_b[t_b(x), N_ib] over trees, shape (n_query, n_train)
        cov = (pred_centered.T @ inbag_centered) / n_trees
        raw = np.sum(cov * cov, axis=1)
        # bootstrap-of-little-bags Monte-Carlo bias correction
        pred_var = pred_centered.var(axis=0)  # (n_query,)
        inbag_var = float(inbag_centered.var())
        correction = n_train * inbag_var * pred_var / n_trees
        var_unscaled = np.clip(raw - correction, 0.0, np.inf)
        if s >= n_train:
            scale = (n_train - 1.0) / n_train
        else:
            scale = ((n_train - 1.0) / n_train) * (n_train / (n_train - s)) ** 2
        var = scale * var_unscaled
        return np.asarray(np.sqrt(np.clip(var, 0.0, np.inf)), dtype=np.float64)

    # -- variable importance -------------------------------------------------

    def variable_importance(self) -> VariableImportance:
        """Depth-weighted, normalized split frequencies over the covariates.

        A driver signal: which covariates the forest splits on, weighted toward the
        shallow splits that move the effect most. Available after ``fit``.
        """
        model = self._cate_model_forest()
        return VariableImportance(names=model.covariate_names, values=model.importance)

cate(data=None, *, uncertainty=True)

Per-unit CATE with infinitesimal-jackknife standard errors and intervals.

Overrides the base meta-learner cate (which would bootstrap-refit the whole forest). The point is the forest-kernel local effect; the per-unit variance is the infinitesimal jackknife over the trees (the bootstrap-of-little-bags grouping), valid because the forest is honest and subsampled. The result carries the cached doubly-robust scores for :meth:~ergodic.inference.CATEResult.explain.

uncertainty=False returns the point effects with kind="point" and no standard error, skipping the infinitesimal-jackknife pass over the trees.

Source code in src/ergodic/inference/estimators/_forest.py

def cate(self, data: TabularData | None = None, *, uncertainty: bool = True) -> CATEResult:
    """Per-unit CATE with infinitesimal-jackknife standard errors and intervals.

    Overrides the base meta-learner ``cate`` (which would bootstrap-refit the
    whole forest). The point is the forest-kernel local effect; the per-unit
    variance is the infinitesimal jackknife over the trees (the
    bootstrap-of-little-bags grouping), valid because the forest is honest and
    subsampled. The result carries the cached doubly-robust scores for
    :meth:`~ergodic.inference.CATEResult.explain`.

    ``uncertainty=False`` returns the point effects with ``kind="point"`` and no
    standard error, skipping the infinitesimal-jackknife pass over the trees.
    """
    from .._perunit import CATEResult

    model = self._cate_model_forest()
    z = self._cate_covariates(data)
    point = np.asarray(model.predict(z), dtype=np.float64)
    n = point.shape[0]
    std_error = self._ij_standard_error(model, z, point) if uncertainty else None
    kind = "analytic" if uncertainty else "point"
    scores = self._cate_scores(data)
    return CATEResult(
        treatment=self._treatment_name,
        outcome=self._outcome_name,
        covariates=self._covariates,
        n=n,
        kind=kind,
        source=self.name,
        diagnostics=dict(self._state.diagnostics),
        identification=self._identification,
        _values=point,
        _index=self._cate_index(data, n),
        _std_error=std_error,
        _scores=scores,
    )

variable_importance()

Depth-weighted, normalized split frequencies over the covariates.

A driver signal: which covariates the forest splits on, weighted toward the shallow splits that move the effect most. Available after fit.

Source code in src/ergodic/inference/estimators/_forest.py

def variable_importance(self) -> VariableImportance:
    """Depth-weighted, normalized split frequencies over the covariates.

    A driver signal: which covariates the forest splits on, weighted toward the
    shallow splits that move the effect most. Available after ``fit``.
    """
    model = self._cate_model_forest()
    return VariableImportance(names=model.covariate_names, values=model.importance)

DML

Bases: _AdjustmentBase

Double machine learning, the partially linear model (partialling-out).

Residualize the outcome and the treatment against the covariates with cross-fitting, then regress one residual on the other. The treatment may be binary or continuous; a continuous treatment reports the average partial effect. The ATE standard error comes from the orthogonal moment.

Source code in src/ergodic/inference/estimators/_dml.py

class DML(_AdjustmentBase):
    """Double machine learning, the partially linear model (partialling-out).

    Residualize the outcome and the treatment against the covariates with
    cross-fitting, then regress one residual on the other. The treatment may be
    binary or continuous; a continuous treatment reports the average partial
    effect. The ATE standard error comes from the orthogonal moment.
    """

    name = "dml"
    slots = (Slot("outcome", Task.REGRESSION), Slot("treatment", Task.REGRESSION))
    targets = frozenset({EffectTarget.ATE})

    def __init__(self, learners: Any = None, *, n_folds: int = 5, **kwargs: Any) -> None:
        super().__init__(learners, **kwargs)
        self.n_folds = n_folds
        self._t_binary = False
        self._cache: (
            tuple[float, npt.NDArray[np.float64], npt.NDArray[np.float64], float] | None
        ) = None

    def _extract(
        self, identification: Identification, data: TabularData
    ) -> tuple[dict[str, npt.NDArray[np.float64]], bool, dict[str, float]]:
        treatment = _single(identification.treatment, "treatment")
        outcome = _single(identification.outcome, "outcome")
        covariates = identification.adjustment_set or frozenset()
        _require_complete(data, {treatment, outcome, *covariates})
        self._covariates = tuple(sorted(covariates))
        a = self._treatment_vector(data, treatment, allow_continuous=True)
        y, y_binary = _outcome_vector(data, outcome)
        z = _covariate_matrix(data, covariates)
        self._t_binary = set(np.unique(a).tolist()) <= {0.0, 1.0}
        arrays = {"A": a, "Y": y, "Z": z}
        diagnostics = {"n_covariates": float(z.shape[1]), "treatment_binary": float(self._t_binary)}
        return arrays, y_binary, diagnostics

    def _default_learner(self, slot: Slot, y_binary: bool) -> Learner:
        if slot.name == "treatment":
            return LogisticRegression() if self._t_binary else OLS()
        return super()._default_learner(slot, y_binary)

    def _residuals(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        n = a.shape[0]
        res_y = np.empty(n, dtype=np.float64)
        res_t = np.empty(n, dtype=np.float64)
        rng = np.random.default_rng(self._random_state)
        stratify = a if self._t_binary else None
        for test in _folds(n, self.n_folds, rng, stratify=stratify):
            train = _train_mask(n, test)
            res_y[test] = y[test] - _fit_predict(make["outcome"], z[train], y[train], z[test])
            res_t[test] = a[test] - _fit_predict(make["treatment"], z[train], a[train], z[test])
        return res_y, res_t

    def _solve(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[float, npt.NDArray[np.float64], npt.NDArray[np.float64], float]:
        res_y, res_t = self._residuals(arrays, make)
        denom = float(np.sum(res_t * res_t))
        if denom <= 0.0:
            raise ValueError("no treatment variation after residualizing")
        theta = float(np.sum(res_t * res_y) / denom)
        return theta, res_y, res_t, denom

    def fit(self, identification: Identification, data: TabularData) -> Self:
        super().fit(identification, data)
        state = self._state
        self._cache = self._solve(state.arrays, state.make)
        return self

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        return self._solve(arrays, make)[0]

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        if target is not EffectTarget.ATE or self._cache is None:
            return None
        theta, res_y, res_t, denom = self._cache
        n = res_t.shape[0]
        influence = res_t * (res_y - theta * res_t) / (denom / n)
        return theta, influence

    def cate(self, data: TabularData | None = None, *, uncertainty: bool = True) -> CATEResult:
        """The constant per-unit effect for the partially-linear model: ``theta``.

        The partially-linear DML assumes a homogeneous effect, so every unit's CATE
        is the single slope ``theta``, with the same analytic standard error from the
        orthogonal moment. Returned as a :class:`~ergodic.inference.CATEResult` for a
        uniform surface; a binary treatment additionally caches the doubly-robust
        scores so ``explain`` can test for the heterogeneity the model assumes away.
        ``uncertainty=False`` returns the point ``theta`` with ``kind="point"`` and no
        standard error.
        """
        from .._perunit import CATEResult

        if self._cache is None:
            raise RuntimeError("dml is not fitted; call fit first")
        theta = self._cache[0]
        z = self._cate_covariates(data)
        n = z.shape[0]
        values = np.full(n, theta, dtype=np.float64)
        std_error: npt.NDArray[np.float64] | None = None
        kind = "point"
        if uncertainty:
            analytic = self._analytic(self._state.arrays, EffectTarget.ATE, self._state.make)
            if analytic is not None:
                influence = analytic[1]
                se = float(np.std(influence, ddof=1) / np.sqrt(influence.shape[0]))
                std_error = np.full(n, se, dtype=np.float64)
                kind = "analytic"
        scores = self._dml_scores() if data is None else None
        return CATEResult(
            treatment=self._treatment_name,
            outcome=self._outcome_name,
            covariates=self._covariates,
            n=n,
            kind=kind,
            source=self.name,
            diagnostics=dict(self._state.diagnostics),
            identification=self._identification,
            _values=values,
            _index=np.arange(n, dtype=np.int64),
            _std_error=std_error,
            _scores=scores,
        )

    def _dml_scores(self) -> DRScores | None:
        """Doubly-robust scores for explainability, only for a binary treatment."""
        if not self._t_binary or self._identification is None:
            return None
        from .._scores import DRScores
        from .._scores import aipw_scores as _scores_fn

        state = self._state
        a, y, z = state.arrays["A"], state.arrays["Y"], state.arrays["Z"]
        try:
            mu0, mu1, e = _dml_crossfit_nuisances(
                a, y, z, state.make, self.n_folds, self._random_state
            )
        except (ValueError, np.linalg.LinAlgError):
            return None
        return DRScores(
            scores=_scores_fn(a, y, mu0, mu1, e),
            mu0=mu0,
            mu1=mu1,
            propensity=e,
            covariates=z,
            treatment=a,
            outcome=y,
            covariate_names=self._covariates,
        )

cate(data=None, *, uncertainty=True)

The constant per-unit effect for the partially-linear model: theta.

The partially-linear DML assumes a homogeneous effect, so every unit's CATE is the single slope theta, with the same analytic standard error from the orthogonal moment. Returned as a :class:~ergodic.inference.CATEResult for a uniform surface; a binary treatment additionally caches the doubly-robust scores so explain can test for the heterogeneity the model assumes away. uncertainty=False returns the point theta with kind="point" and no standard error.

Source code in src/ergodic/inference/estimators/_dml.py

def cate(self, data: TabularData | None = None, *, uncertainty: bool = True) -> CATEResult:
    """The constant per-unit effect for the partially-linear model: ``theta``.

    The partially-linear DML assumes a homogeneous effect, so every unit's CATE
    is the single slope ``theta``, with the same analytic standard error from the
    orthogonal moment. Returned as a :class:`~ergodic.inference.CATEResult` for a
    uniform surface; a binary treatment additionally caches the doubly-robust
    scores so ``explain`` can test for the heterogeneity the model assumes away.
    ``uncertainty=False`` returns the point ``theta`` with ``kind="point"`` and no
    standard error.
    """
    from .._perunit import CATEResult

    if self._cache is None:
        raise RuntimeError("dml is not fitted; call fit first")
    theta = self._cache[0]
    z = self._cate_covariates(data)
    n = z.shape[0]
    values = np.full(n, theta, dtype=np.float64)
    std_error: npt.NDArray[np.float64] | None = None
    kind = "point"
    if uncertainty:
        analytic = self._analytic(self._state.arrays, EffectTarget.ATE, self._state.make)
        if analytic is not None:
            influence = analytic[1]
            se = float(np.std(influence, ddof=1) / np.sqrt(influence.shape[0]))
            std_error = np.full(n, se, dtype=np.float64)
            kind = "analytic"
    scores = self._dml_scores() if data is None else None
    return CATEResult(
        treatment=self._treatment_name,
        outcome=self._outcome_name,
        covariates=self._covariates,
        n=n,
        kind=kind,
        source=self.name,
        diagnostics=dict(self._state.diagnostics),
        identification=self._identification,
        _values=values,
        _index=np.arange(n, dtype=np.int64),
        _std_error=std_error,
        _scores=scores,
    )

DRLearner

Bases: _MetaLearnerBase

Doubly-robust learning of the ATE and the CATE.

Cross-fit the AIPW score into a per-unit pseudo-outcome, average it for the ATE (with the influence-function standard error), and regress it on the covariates for the CATE through cate. The CATE result carries the cached doubly-robust scores, so its explain runs the Tier-B characterization.

Source code in src/ergodic/inference/estimators/_dml.py

class DRLearner(_MetaLearnerBase):
    """Doubly-robust learning of the ATE and the CATE.

    Cross-fit the AIPW score into a per-unit pseudo-outcome, average it for the
    ATE (with the influence-function standard error), and regress it on the
    covariates for the CATE through ``cate``. The CATE result carries the cached
    doubly-robust scores, so its ``explain`` runs the Tier-B characterization.
    """

    name = "dr_learner"
    slots = (
        Slot("outcome", Task.REGRESSION),
        Slot("propensity", Task.CLASSIFICATION),
        Slot("final", Task.REGRESSION),
    )
    targets = frozenset({EffectTarget.ATE})

    def __init__(
        self, learners: Any = None, *, n_folds: int = 5, clip: float = 1e-3, **kwargs: Any
    ) -> None:
        super().__init__(learners, **kwargs)
        self.n_folds = n_folds
        self.clip = clip
        self._pseudo: npt.NDArray[np.float64] | None = None

    def _crossfit(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        n = a.shape[0]
        mu0 = np.empty(n, dtype=np.float64)
        mu1 = np.empty(n, dtype=np.float64)
        e = np.empty(n, dtype=np.float64)
        rng = np.random.default_rng(self._random_state)
        for test in _folds(n, self.n_folds, rng, stratify=a):
            train = _train_mask(n, test)
            design = np.hstack([a[train].reshape(-1, 1), z[train]])
            outcome = make["outcome"]().fit(design, y[train])
            ones = np.ones((test.shape[0], 1), dtype=np.float64)
            zeros = np.zeros((test.shape[0], 1), dtype=np.float64)
            mu1[test] = _mean_predict(outcome, np.hstack([ones, z[test]]))
            mu0[test] = _mean_predict(outcome, np.hstack([zeros, z[test]]))
            propensity = make["propensity"]().fit(z[train], a[train])
            e[test] = np.clip(_mean_predict(propensity, z[test]), self.clip, 1.0 - self.clip)
        return mu0, mu1, e

    def _pseudo_outcome(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> npt.NDArray[np.float64]:
        a = arrays["A"]
        y = arrays["Y"]
        mu0, mu1, e = self._crossfit(arrays, make)
        return (mu1 - mu0) + a * (y - mu1) / e - (1.0 - a) * (y - mu0) / (1.0 - e)

    def _dr_nuisances(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        # reuse the estimator's own cross-fit for the explainability scores
        return self._crossfit(arrays, make)

    def _fit_cate(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> _DRModel:
        pseudo = self._pseudo_outcome(arrays, make)
        final = make["final"]().fit(arrays["Z"], pseudo)
        return _DRModel(pseudo, final)

    def _predict_cate(self, model: _DRModel, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return _mean_predict(model.final, z)

    def fit(self, identification: Identification, data: TabularData) -> Self:
        super().fit(identification, data)
        self._pseudo = self._cate_model().pseudo
        return self

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        if target is not EffectTarget.ATE or self._pseudo is None:
            return None
        return float(np.mean(self._pseudo)), self._pseudo

    def _cate_analytic_se(
        self, model: _DRModel, z: npt.NDArray[np.float64], data: TabularData | None
    ) -> npt.NDArray[np.float64] | None:
        """Per-unit standard errors from the linear final's sandwich covariance.

        When the final stage is the native linear :class:`OLS`, the pseudo-outcome
        regression ``tau_hat(z) = [1, z] beta`` has a heteroskedasticity-robust
        coefficient covariance, and each unit's variance is the quadratic form
        ``x' Cov(beta) x``. Returns ``None`` for any other (nonlinear) final, where the
        per-unit point predictions have no closed-form interval and the bootstrap
        path takes over.
        """
        final = model.final
        if not isinstance(final, OLS) or final._coef is None:
            return None
        z_train = self._state.arrays["Z"]
        x_train = _ols_design(final, z_train)
        x_eval = _ols_design(final, z)
        resid = model.pseudo - x_train @ final._coef
        xtx_inv = np.linalg.pinv(x_train.T @ x_train)
        meat = (x_train * resid[:, None]).T @ (x_train * resid[:, None])
        cov = xtx_inv @ meat @ xtx_inv
        var = np.einsum("ij,jk,ik->i", x_eval, cov, x_eval)
        return np.asarray(np.sqrt(np.clip(var, 0.0, np.inf)), dtype=np.float64)

DRScores dataclass

The cross-fit doubly-robust object: the AIPW signal plus its nuisances.

scores is the per-unit Gamma_i (its mean is the AIPW ATE). mu0, mu1, and propensity are the cross-fit out-of-fold nuisances, kept so the calibration and projection tools reuse them without refitting. treatment, outcome, and covariates carry the design back to the result.

Source code in src/ergodic/inference/_scores.py

@dataclass(frozen=True)
class DRScores:
    """The cross-fit doubly-robust object: the AIPW signal plus its nuisances.

    ``scores`` is the per-unit ``Gamma_i`` (its mean is the AIPW ATE). ``mu0``,
    ``mu1``, and ``propensity`` are the cross-fit out-of-fold nuisances, kept so the
    calibration and projection tools reuse them without refitting. ``treatment``,
    ``outcome``, and ``covariates`` carry the design back to the result.
    """

    scores: npt.NDArray[np.float64]
    mu0: npt.NDArray[np.float64]
    mu1: npt.NDArray[np.float64]
    propensity: npt.NDArray[np.float64]
    covariates: npt.NDArray[np.float64]
    treatment: npt.NDArray[np.float64]
    outcome: npt.NDArray[np.float64]
    covariate_names: tuple[str, ...]
    tau_hat: npt.NDArray[np.float64] | None = None

Design dataclass

How a quasi-experimental estimator claims identification.

strategy is one of "did", "synthetic_control", or "interrupted_time_series". assumptions are the untestable claims the design rests on, stated in words (parallel trends, convex donor span, no co-timed shock); they are what a reader should challenge, so they ride with every result.

Source code in src/ergodic/inference/_design.py

@dataclass(frozen=True)
class Design:
    """How a quasi-experimental estimator claims identification.

    ``strategy`` is one of ``"did"``, ``"synthetic_control"``, or
    ``"interrupted_time_series"``. ``assumptions`` are the untestable claims the
    design rests on, stated in words (parallel trends, convex donor span, no
    co-timed shock); they are what a reader should challenge, so they ride with
    every result.
    """

    strategy: str
    outcome: str
    treatment: str | None
    assumptions: tuple[str, ...]
    estimand: str = ""

DiDEffect dataclass

Bases: EffectEstimate

A difference-in-differences result: the estimate plus the design it rests on.

The :class:EffectEstimate surface (value, std_error, interval, draws) reads the headline ATT. flavor records which estimator ran, comparison which pool stood in for the counterfactual, by_group_time the per-cell detail of the group-time flavor, and by_event_time the event-study aggregation when :func:event_study built the result.

Source code in src/ergodic/inference/estimators/_did.py

@dataclass(frozen=True)
class DiDEffect(EffectEstimate):
    """A difference-in-differences result: the estimate plus the design it rests on.

    The :class:`EffectEstimate` surface (``value``, ``std_error``, ``interval``,
    ``draws``) reads the headline ATT. ``flavor`` records which estimator ran,
    ``comparison`` which pool stood in for the counterfactual, ``by_group_time``
    the per-cell detail of the group-time flavor, and ``by_event_time`` the
    event-study aggregation when :func:`event_study` built the result.
    """

    design: Design | None = None
    flavor: str = ""
    comparison: str = ""
    by_group_time: tuple[GroupTimeEffect, ...] = field(default=(), repr=False)
    by_event_time: tuple[EventTimeEffect, ...] = field(default=(), repr=False)

EffectEstimate dataclass

A causal effect with its uncertainty.

Built by a :class:CausalEstimator. Either _point/_std_error are set (an analytic result) or _draws is set (bootstrap or posterior), and the accessors read whichever is present.

Source code in src/ergodic/inference/_result.py

@dataclass(frozen=True)
class EffectEstimate:
    """A causal effect with its uncertainty.

    Built by a :class:`CausalEstimator`. Either ``_point``/``_std_error`` are set
    (an analytic result) or ``_draws`` is set (bootstrap or posterior), and the
    accessors read whichever is present.
    """

    target: EffectTarget
    estimator: str
    learners: Mapping[str, str]
    n: int
    kind: str  # "analytic", "bootstrap", or "posterior"
    diagnostics: Mapping[str, float] = field(default_factory=dict)
    _point: float | None = None
    _std_error: float | None = None
    _draws: npt.NDArray[np.float64] | None = field(default=None, repr=False)

    @property
    def value(self) -> float:
        """The point estimate, or the mean of the draws."""
        if self._point is not None:
            return self._point
        if self._draws is not None:
            return float(np.mean(self._draws))
        raise ValueError("estimate has neither a point nor draws")

    @property
    def std_error(self) -> float | None:
        """The analytic standard error, or the standard deviation of the draws."""
        if self._std_error is not None:
            return self._std_error
        if self._draws is not None:
            return float(np.std(self._draws, ddof=1))
        return None

    def interval(self, level: float = 0.95) -> tuple[float, float]:
        """A confidence or credible interval at ``level``.

        From the draw quantiles when draws are present, else the normal interval
        around the point estimate.
        """
        if not 0.0 < level < 1.0:
            raise ValueError(f"level must be in (0, 1), got {level}")
        if self._draws is not None:
            lo = float(np.quantile(self._draws, (1.0 - level) / 2.0))
            hi = float(np.quantile(self._draws, (1.0 + level) / 2.0))
            return (lo, hi)
        if self._point is not None and self._std_error is not None:
            z = _norm_ppf((1.0 + level) / 2.0)
            return (self._point - z * self._std_error, self._point + z * self._std_error)
        raise ValueError("estimate has no uncertainty to form an interval from")

    @property
    def draws(self) -> npt.NDArray[np.float64] | None:
        """The bootstrap or posterior draws, or ``None`` for a pure analytic result."""
        return self._draws

    def __repr__(self) -> str:
        se = self.std_error
        body = f"{self.target}={self.value:.4g}"
        if se is not None:
            lo, hi = self.interval()
            body += f", se={se:.3g}, 95% CI [{lo:.3g}, {hi:.3g}]"
        return f"EffectEstimate({self.estimator}, {self.kind}: {body}, n={self.n})"

draws property

The bootstrap or posterior draws, or None for a pure analytic result.

std_error property

The analytic standard error, or the standard deviation of the draws.

value property

The point estimate, or the mean of the draws.

interval(level=0.95)

A confidence or credible interval at level.

From the draw quantiles when draws are present, else the normal interval around the point estimate.

Source code in src/ergodic/inference/_result.py

def interval(self, level: float = 0.95) -> tuple[float, float]:
    """A confidence or credible interval at ``level``.

    From the draw quantiles when draws are present, else the normal interval
    around the point estimate.
    """
    if not 0.0 < level < 1.0:
        raise ValueError(f"level must be in (0, 1), got {level}")
    if self._draws is not None:
        lo = float(np.quantile(self._draws, (1.0 - level) / 2.0))
        hi = float(np.quantile(self._draws, (1.0 + level) / 2.0))
        return (lo, hi)
    if self._point is not None and self._std_error is not None:
        z = _norm_ppf((1.0 + level) / 2.0)
        return (self._point - z * self._std_error, self._point + z * self._std_error)
    raise ValueError("estimate has no uncertainty to form an interval from")

EffectTarget

Bases: Enum

Which effect a :class:CausalEstimator reports.

E1 returns ATE and ATT for the adjustment and front-door families and LATE for instruments. ATC, CATE, and ITE are reserved for the heterogeneous-effect phase.

Source code in src/ergodic/inference/_target.py

class EffectTarget(Enum):
    """Which effect a :class:`CausalEstimator` reports.

    E1 returns ``ATE`` and ``ATT`` for the adjustment and front-door families and
    ``LATE`` for instruments. ``ATC``, ``CATE``, and ``ITE`` are reserved for the
    heterogeneous-effect phase.
    """

    ATE = "ate"
    ATT = "att"
    ATC = "atc"
    LATE = "late"
    CATE = "cate"
    ITE = "ite"

    def __str__(self) -> str:
        return self.value

EventTimeEffect dataclass

The effect at one event time, aggregated over cohorts with cohort-size weights.

std_error and interval come from the unit bootstrap: the standard deviation of the draws and their 95% percentile interval. n_treated counts the treated units contributing a cell at this event time.

Source code in src/ergodic/inference/estimators/_did.py

@dataclass(frozen=True)
class EventTimeEffect:
    """The effect at one event time, aggregated over cohorts with cohort-size weights.

    ``std_error`` and ``interval`` come from the unit bootstrap: the standard
    deviation of the draws and their 95% percentile interval. ``n_treated``
    counts the treated units contributing a cell at this event time.
    """

    event_time: int
    estimate: float
    std_error: float | None
    interval: tuple[float, float]
    n_treated: int

GroupTimeEffect dataclass

One ATT(g, t) cell: the cohort adopting at group, read at period.

event_time is the position of period minus the position of group in the sorted period sequence, so -1 is the base period and never appears. n_treated and n_comparison count the units observed at both the base and the target period on each side of the comparison.

Source code in src/ergodic/inference/estimators/_did.py

@dataclass(frozen=True)
class GroupTimeEffect:
    """One ATT(g, t) cell: the cohort adopting at ``group``, read at ``period``.

    ``event_time`` is the position of ``period`` minus the position of ``group``
    in the sorted period sequence, so -1 is the base period and never appears.
    ``n_treated`` and ``n_comparison`` count the units observed at both the base
    and the target period on each side of the comparison.
    """

    group: Any
    period: Any
    event_time: int
    estimate: float
    n_treated: int
    n_comparison: int

Heterogeneity dataclass

Tier-B heterogeneity characterization on shared doubly-robust scores.

Carries the cross-fit :class:~ergodic.inference._scores.DRScores and the fitted per-unit tau_hat. The honesty register: the rankings and proxies are fixed (computed cross-fit), so the orthogonal summaries below inherit valid inference while the CATE point predictions themselves never carry a confidence interval.

Source code in src/ergodic/inference/explain/_record.py

@dataclass(frozen=True)
class Heterogeneity:
    """Tier-B heterogeneity characterization on shared doubly-robust scores.

    Carries the cross-fit :class:`~ergodic.inference._scores.DRScores` and the fitted
    per-unit ``tau_hat``. The honesty register: the rankings and proxies are fixed
    (computed cross-fit), so the orthogonal summaries below inherit valid inference
    while the CATE point predictions themselves never carry a confidence interval.
    """

    scores: DRScores
    tau_hat: npt.NDArray[np.float64]
    treatment: str
    outcome: str

    def blp(self, *, level: float = 0.9) -> BLPResult:
        """Best Linear Projection of the CATE on the covariates (HC3 inference)."""
        return best_linear_projection(
            self.scores.scores,
            self.scores.covariates,
            self.scores.covariate_names,
            level=level,
        )

    def gates(self, *, n_groups: int = 5, level: float = 0.9) -> GATESResult:
        """Sorted Group ATEs across ``n_groups`` predicted-effect quantiles."""
        return gates(self.scores.scores, self.tau_hat, n_groups=n_groups, level=level)

    def clan(self, *, n_groups: int = 5, level: float = 0.9) -> CLANResult:
        """Covariate-mean differences between the most- and least-affected groups."""
        return clan(
            self.scores.covariates,
            self.scores.covariate_names,
            self.tau_hat,
            n_groups=n_groups,
            level=level,
        )

    def calibration_test(self) -> CalibrationResult:
        """The omnibus calibration / heterogeneity test (one-sided HTE p-value)."""
        return calibration_test(self.scores.scores, self.tau_hat)

    def rate(
        self, metric: str = "autoc", *, n_boot: int = 200, random_state: int | None = None
    ) -> RATEResult:
        """The Rank-Weighted ATE of the ``tau_hat`` ranking: ``"autoc"`` or ``"qini"``."""
        return rate(
            self.scores.scores,
            self.tau_hat,
            metric=metric,
            n_boot=n_boot,
            random_state=random_state,
        )

    def policy_tree(
        self,
        covariates: str | tuple[str, ...] | list[str] | None = None,
        *,
        depth: int = 2,
    ) -> PolicyTree:
        """The depth-limited, DR-welfare-optimal policy tree off these scores.

        A convenience that forwards the shared doubly-robust scores to
        :func:`~ergodic.inference.policy_tree`. ``covariates`` selects the split
        columns by name (``None`` uses all of them). The "act" companion to
        :meth:`rate`: where ``rate`` scores a ranking, this returns an interpretable
        treat-or-not rule and its estimated policy value.
        """
        from ._policy import policy_tree as _policy_tree

        return _policy_tree(self.scores, covariates, depth=depth)

blp(*, level=0.9)

Best Linear Projection of the CATE on the covariates (HC3 inference).

Source code in src/ergodic/inference/explain/_record.py

def blp(self, *, level: float = 0.9) -> BLPResult:
    """Best Linear Projection of the CATE on the covariates (HC3 inference)."""
    return best_linear_projection(
        self.scores.scores,
        self.scores.covariates,
        self.scores.covariate_names,
        level=level,
    )

calibration_test()

The omnibus calibration / heterogeneity test (one-sided HTE p-value).

Source code in src/ergodic/inference/explain/_record.py

def calibration_test(self) -> CalibrationResult:
    """The omnibus calibration / heterogeneity test (one-sided HTE p-value)."""
    return calibration_test(self.scores.scores, self.tau_hat)

clan(*, n_groups=5, level=0.9)

Covariate-mean differences between the most- and least-affected groups.

Source code in src/ergodic/inference/explain/_record.py

def clan(self, *, n_groups: int = 5, level: float = 0.9) -> CLANResult:
    """Covariate-mean differences between the most- and least-affected groups."""
    return clan(
        self.scores.covariates,
        self.scores.covariate_names,
        self.tau_hat,
        n_groups=n_groups,
        level=level,
    )

gates(*, n_groups=5, level=0.9)

Sorted Group ATEs across n_groups predicted-effect quantiles.

Source code in src/ergodic/inference/explain/_record.py

def gates(self, *, n_groups: int = 5, level: float = 0.9) -> GATESResult:
    """Sorted Group ATEs across ``n_groups`` predicted-effect quantiles."""
    return gates(self.scores.scores, self.tau_hat, n_groups=n_groups, level=level)

policy_tree(covariates=None, *, depth=2)

The depth-limited, DR-welfare-optimal policy tree off these scores.

A convenience that forwards the shared doubly-robust scores to :func:~ergodic.inference.policy_tree. covariates selects the split columns by name (None uses all of them). The "act" companion to :meth:rate: where rate scores a ranking, this returns an interpretable treat-or-not rule and its estimated policy value.

Source code in src/ergodic/inference/explain/_record.py

def policy_tree(
    self,
    covariates: str | tuple[str, ...] | list[str] | None = None,
    *,
    depth: int = 2,
) -> PolicyTree:
    """The depth-limited, DR-welfare-optimal policy tree off these scores.

    A convenience that forwards the shared doubly-robust scores to
    :func:`~ergodic.inference.policy_tree`. ``covariates`` selects the split
    columns by name (``None`` uses all of them). The "act" companion to
    :meth:`rate`: where ``rate`` scores a ranking, this returns an interpretable
    treat-or-not rule and its estimated policy value.
    """
    from ._policy import policy_tree as _policy_tree

    return _policy_tree(self.scores, covariates, depth=depth)

rate(metric='autoc', *, n_boot=200, random_state=None)

The Rank-Weighted ATE of the tau_hat ranking: "autoc" or "qini".

Source code in src/ergodic/inference/explain/_record.py

def rate(
    self, metric: str = "autoc", *, n_boot: int = 200, random_state: int | None = None
) -> RATEResult:
    """The Rank-Weighted ATE of the ``tau_hat`` ranking: ``"autoc"`` or ``"qini"``."""
    return rate(
        self.scores.scores,
        self.tau_hat,
        metric=metric,
        n_boot=n_boot,
        random_state=random_state,
    )

IPW

Bases: _AdjustmentBase

Weight by the inverse propensity of treatment.

scheme="hajek" (the default) self-normalizes the weights; scheme="ht" is the raw Horvitz-Thompson estimator. Propensities are clipped to [clip, 1 - clip] to bound the weights.

Source code in src/ergodic/inference/estimators/_weighting.py

class IPW(_AdjustmentBase):
    """Weight by the inverse propensity of treatment.

    ``scheme="hajek"`` (the default) self-normalizes the weights; ``scheme="ht"``
    is the raw Horvitz-Thompson estimator. Propensities are clipped to
    ``[clip, 1 - clip]`` to bound the weights.
    """

    name = "ipw"
    slots = (Slot("propensity", Task.CLASSIFICATION),)
    targets = frozenset({EffectTarget.ATE, EffectTarget.ATT})

    def __init__(
        self, learners: Any = None, *, scheme: str = "hajek", clip: float = 1e-3, **kwargs: Any
    ) -> None:
        super().__init__(learners, **kwargs)
        if scheme not in ("hajek", "ht"):
            raise ValueError(f"scheme must be 'hajek' or 'ht', got {scheme!r}")
        self.scheme = scheme
        self.clip = clip

    def _propensity(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> npt.NDArray[np.float64]:
        a = arrays["A"]
        z = arrays["Z"]
        model = make["propensity"]().fit(z, a)
        e = _mean_predict(model, z)
        return np.clip(e, self.clip, 1.0 - self.clip)

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        a = arrays["A"]
        y = arrays["Y"]
        e = self._propensity(arrays, make)
        if target is EffectTarget.ATT:
            odds = e / (1.0 - e)
            w0 = (1.0 - a) * odds
            denom = float(np.sum(w0))
            if denom <= 0.0:
                raise ValueError("no control weight for ATT")
            mu0 = float(np.sum(w0 * y) / denom)
            return float(np.mean(y[a == 1.0])) - mu0
        if self.scheme == "hajek":
            w1 = a / e
            w0 = (1.0 - a) / (1.0 - e)
            s1 = float(np.sum(w1))
            s0 = float(np.sum(w0))
            if s1 <= 0.0 or s0 <= 0.0:
                raise ValueError("a treatment arm has no weight")
            return float(np.sum(w1 * y) / s1 - np.sum(w0 * y) / s0)
        return float(np.mean(a * y / e - (1.0 - a) * y / (1.0 - e)))

ITEResult dataclass

Bases: PerUnitEffect

Individual treatment effects from a structural model (the twin network).

The individual's own counterfactual contrast (the SCM-pillar object), not a conditional average. Carries the factual outcome column, the contrast framing (levels= or shift=), and a structural provenance note.

:meth:explain is deferred: the Tier-B machinery is defined over the CATE's doubly-robust scores, and the inference-valid descriptive regression of abducted ITEs is a later phase, so it raises with a pointer rather than faking inference.

Source code in src/ergodic/inference/_perunit.py

@dataclass(frozen=True, eq=False)
class ITEResult(PerUnitEffect):
    """Individual treatment effects from a structural model (the twin network).

    The individual's own counterfactual contrast (the SCM-pillar object), not a
    conditional average. Carries the factual outcome column, the contrast framing
    (``levels=`` or ``shift=``), and a structural provenance note.

    :meth:`explain` is deferred: the Tier-B machinery is defined over the CATE's
    doubly-robust scores, and the inference-valid descriptive regression of abducted
    ITEs is a later phase, so it raises with a pointer rather than faking inference.
    """

    estimand: str = "ite"
    framing: str = ""
    note: str = ""
    _factual: npt.NDArray[np.float64] | None = field(default=None, repr=False)

    @property
    def factual(self) -> npt.NDArray[np.float64] | None:
        """The per-unit factual outcome the contrast was taken against, when carried."""
        if self._factual is None:
            return None
        return np.asarray(self._factual, dtype=np.float64)

    def explain(self) -> Heterogeneity:
        """Deferred in phase 1: the Tier-B tools need the CATE's doubly-robust scores."""
        raise NotImplementedError(
            "ITEResult.explain is not available yet: the Tier-B characterization "
            "(BLP, GATES, RATE) is defined over the doubly-robust CATE scores, and the "
            "descriptive regression of structural ITEs is a later phase. Use a CATE "
            "meta-learner (cate(...)) for inference-valid heterogeneity characterization"
        )

factual property

The per-unit factual outcome the contrast was taken against, when carried.

explain()

Deferred in phase 1: the Tier-B tools need the CATE's doubly-robust scores.

Source code in src/ergodic/inference/_perunit.py

def explain(self) -> Heterogeneity:
    """Deferred in phase 1: the Tier-B tools need the CATE's doubly-robust scores."""
    raise NotImplementedError(
        "ITEResult.explain is not available yet: the Tier-B characterization "
        "(BLP, GATES, RATE) is defined over the doubly-robust CATE scores, and the "
        "descriptive regression of structural ITEs is a later phase. Use a CATE "
        "meta-learner (cate(...)) for inference-valid heterogeneity characterization"
    )

ITSEffect dataclass

Bases: EffectEstimate

A segmented-regression fit around one break.

level_change is the jump at the break and slope_change the change in trend per period, each with a Newey-West standard error. value is the ATT: the average post-period gap between fitted and counterfactual, where counterfactual projects the pre-period trend over every period with the break terms zeroed and the seasonal terms kept.

Source code in src/ergodic/inference/estimators/_its.py

@dataclass(frozen=True)
class ITSEffect(EffectEstimate):
    """A segmented-regression fit around one break.

    ``level_change`` is the jump at the break and ``slope_change`` the change
    in trend per period, each with a Newey-West standard error. ``value`` is
    the ATT: the average post-period gap between ``fitted`` and
    ``counterfactual``, where ``counterfactual`` projects the pre-period trend
    over every period with the break terms zeroed and the seasonal terms kept.
    """

    design: Design | None = None
    at: Any = None
    level_change: float = float("nan")
    level_change_se: float = float("nan")
    slope_change: float = float("nan")
    slope_change_se: float = float("nan")
    times: tuple[Any, ...] = ()
    observed: tuple[float, ...] = ()
    counterfactual: tuple[float, ...] = ()
    fitted: tuple[float, ...] = ()

KNNClassifier

Bases: _KNNBase

K-nearest-neighbors classification for a binary 0/1 label.

Source code in src/ergodic/inference/learners/_neighbors.py

class KNNClassifier(_KNNBase):
    """K-nearest-neighbors classification for a binary 0/1 label."""

    def fit(
        self,
        X: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        sample_weight: npt.NDArray[np.float64] | None = None,
    ) -> KNNClassifier:
        self._fit_xy(X, y)
        return self

    def predict_proba(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        prob = np.mean(self._neighbor_targets(X), axis=1)
        return np.asarray(np.column_stack([1.0 - prob, prob]), dtype=np.float64)

    def predict(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return np.asarray(self.predict_proba(X)[:, 1] >= 0.5, dtype=np.float64)

KNNRegressor

Bases: _KNNBase

K-nearest-neighbors regression: the prediction is the neighbor mean.

Source code in src/ergodic/inference/learners/_neighbors.py

class KNNRegressor(_KNNBase):
    """K-nearest-neighbors regression: the prediction is the neighbor mean."""

    def fit(
        self,
        X: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        sample_weight: npt.NDArray[np.float64] | None = None,
    ) -> KNNRegressor:
        self._fit_xy(X, y)
        return self

    def predict(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return np.asarray(np.mean(self._neighbor_targets(X), axis=1), dtype=np.float64)

Learner

Bases: Protocol

The sklearn estimator surface the estimators rely on.

Source code in src/ergodic/inference/_learner.py

@runtime_checkable
class Learner(Protocol):
    """The sklearn estimator surface the estimators rely on."""

    def fit(
        self,
        X: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        sample_weight: npt.NDArray[np.float64] | None = None,
    ) -> Any: ...

    def predict(self, X: npt.NDArray[np.float64]) -> Any: ...

LogisticRegression

Binary logistic regression by Newton-Raphson (IRLS), ridge-stabilized.

Source code in src/ergodic/inference/learners/_linear.py

class LogisticRegression:
    """Binary logistic regression by Newton-Raphson (IRLS), ridge-stabilized."""

    def __init__(
        self,
        *,
        fit_intercept: bool = True,
        l2: float = 1e-6,
        max_iter: int = 100,
        tol: float = 1e-8,
    ) -> None:
        self.fit_intercept = fit_intercept
        self.l2 = l2
        self.max_iter = max_iter
        self.tol = tol
        self._coef: npt.NDArray[np.float64] | None = None

    def fit(
        self,
        X: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        sample_weight: npt.NDArray[np.float64] | None = None,
    ) -> LogisticRegression:
        design = _design(X, self.fit_intercept)
        label = np.asarray(y, dtype=np.float64).reshape(-1)
        n, k = design.shape
        weight = (
            np.ones(n, dtype=np.float64)
            if sample_weight is None
            else np.asarray(sample_weight, dtype=np.float64).reshape(-1)
        )
        penalty = self.l2 * np.eye(k, dtype=np.float64)
        if self.fit_intercept:
            penalty[0, 0] = 0.0
        beta = np.zeros(k, dtype=np.float64)
        for _ in range(self.max_iter):
            prob = _sigmoid(design @ beta)
            grad = design.T @ (weight * (label - prob)) - penalty @ beta
            reweight = weight * prob * (1.0 - prob)
            hessian = design.T @ (design * reweight[:, None]) + penalty
            try:
                step = np.linalg.solve(hessian, grad)
            except np.linalg.LinAlgError:
                step = np.asarray(np.linalg.lstsq(hessian, grad, rcond=None)[0], dtype=np.float64)
            beta = beta + step
            if float(np.max(np.abs(step))) < self.tol:
                break
        self._coef = np.asarray(beta, dtype=np.float64)
        return self

    def predict_proba(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        if self._coef is None:
            raise RuntimeError("LogisticRegression.predict_proba called before fit")
        prob = _sigmoid(_design(X, self.fit_intercept) @ self._coef)
        return np.asarray(np.column_stack([1.0 - prob, prob]), dtype=np.float64)

    def predict(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return np.asarray(self.predict_proba(X)[:, 1] >= 0.5, dtype=np.float64)

    def get_params(self, deep: bool = True) -> dict[str, Any]:
        return {
            "fit_intercept": self.fit_intercept,
            "l2": self.l2,
            "max_iter": self.max_iter,
            "tol": self.tol,
        }

    def set_params(self, **params: Any) -> LogisticRegression:
        for key, value in params.items():
            setattr(self, key, value)
        return self

MediationEstimate dataclass

A mediation decomposition with its uncertainty.

nde, nie, and total are the natural direct effect, the natural indirect effect, and the total effect (nde + nie), each a distribution- shaped :class:~ergodic.inference.EffectEstimate with a bootstrap standard error and interval. cde maps each mediator level to the controlled direct effect at that level (also an :class:EffectEstimate). design names the (untestable) assumptions the decomposition rests on, the same role :class:~ergodic.inference.Design plays for the quasi-experimental estimators.

proportion_mediated is nie / total, the share of the total effect that flows through the mediator. It is unstable when the total effect is near zero (a small denominator), so it is offered with that caveat and a flagged diagnostic.

Source code in src/ergodic/inference/estimators/_mediation.py

@dataclass(frozen=True)
class MediationEstimate:
    """A mediation decomposition with its uncertainty.

    ``nde``, ``nie``, and ``total`` are the natural direct effect, the natural
    indirect effect, and the total effect (``nde + nie``), each a distribution-
    shaped :class:`~ergodic.inference.EffectEstimate` with a bootstrap standard
    error and interval. ``cde`` maps each mediator level to the controlled direct
    effect at that level (also an :class:`EffectEstimate`). ``design`` names the
    (untestable) assumptions the decomposition rests on, the same role
    :class:`~ergodic.inference.Design` plays for the quasi-experimental estimators.

    ``proportion_mediated`` is ``nie / total``, the share of the total effect that
    flows through the mediator. It is unstable when the total effect is near zero
    (a small denominator), so it is offered with that caveat and a ``flagged``
    diagnostic.
    """

    nde: EffectEstimate
    nie: EffectEstimate
    total: EffectEstimate
    cde: dict[Any, EffectEstimate]
    design: Design
    mediators: frozenset[str]
    covariates: frozenset[str]

    @property
    def proportion_mediated(self) -> float:
        """``NIE / TE``: the share of the total effect carried by the mediator.

        Unstable when the total effect is near zero; see
        :attr:`proportion_mediated_flagged`.
        """
        te = self.total.value
        return self.nie.value / te if te != 0.0 else float("nan")

    @property
    def proportion_mediated_flagged(self) -> bool:
        """``True`` when the total effect is too small to trust the proportion.

        The proportion mediated has the total effect in its denominator, so a
        total effect close to zero makes it numerically unstable: a tiny shift in
        either piece swings the ratio. Flagged when ``|TE|`` is below one bootstrap
        standard error (or essentially zero).
        """
        te = self.total.value
        se = self.total.std_error or 0.0
        return abs(te) <= max(se, 1e-8)

    def __repr__(self) -> str:
        pm = self.proportion_mediated
        flag = " (unstable)" if self.proportion_mediated_flagged else ""
        return (
            f"MediationEstimate(NDE={self.nde.value:.4g}, NIE={self.nie.value:.4g}, "
            f"TE={self.total.value:.4g}, proportion_mediated={pm:.3g}{flag})"
        )

proportion_mediated property

NIE / TE: the share of the total effect carried by the mediator.

Unstable when the total effect is near zero; see :attr:proportion_mediated_flagged.

proportion_mediated_flagged property

True when the total effect is too small to trust the proportion.

The proportion mediated has the total effect in its denominator, so a total effect close to zero makes it numerically unstable: a tiny shift in either piece swings the ratio. Flagged when |TE| is below one bootstrap standard error (or essentially zero).

NullEffect

Bases: CausalEstimator

Report the exact zero effect for a null estimand.

Identification returns strategy="null" when the graph has no directed path from the treatment to the outcome, so the intervention cannot move the outcome and the effect is zero by construction. No nuisance slots, no uncertainty: the point is 0.0 with a zero standard error.

Source code in src/ergodic/inference/estimators/_null.py

class NullEffect(CausalEstimator):
    """Report the exact zero effect for a ``null`` estimand.

    Identification returns ``strategy="null"`` when the graph has no directed
    path from the treatment to the outcome, so the intervention cannot move the
    outcome and the effect is zero by construction. No nuisance slots, no
    uncertainty: the point is 0.0 with a zero standard error.
    """

    name = "null_effect"
    consumes = frozenset({NULL})
    targets = frozenset({EffectTarget.ATE, EffectTarget.ATT})

    def _extract(
        self, identification: Identification, data: TabularData
    ) -> tuple[dict[str, npt.NDArray[np.float64]], bool, dict[str, float]]:
        treatment = _single(identification.treatment, "treatment")
        outcome = _single(identification.outcome, "outcome")
        missing = [name for name in (treatment, outcome) if name not in data.schema]
        if missing:
            raise ValueError(f"columns {missing} are not in the data")
        return {}, False, {}

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        return 0.0

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        return 0.0, np.zeros(self._state.n, dtype=np.float64)

OLS

Ordinary least squares, with optional ridge penalty and sample weights.

Source code in src/ergodic/inference/learners/_linear.py

class OLS:
    """Ordinary least squares, with optional ridge penalty and sample weights."""

    def __init__(self, *, fit_intercept: bool = True, ridge: float = 0.0) -> None:
        self.fit_intercept = fit_intercept
        self.ridge = ridge
        self._coef: npt.NDArray[np.float64] | None = None

    def fit(
        self,
        X: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        sample_weight: npt.NDArray[np.float64] | None = None,
    ) -> OLS:
        design = _design(X, self.fit_intercept)
        target = np.asarray(y, dtype=np.float64).reshape(-1)
        n, k = design.shape
        weight = (
            np.ones(n, dtype=np.float64)
            if sample_weight is None
            else np.asarray(sample_weight, dtype=np.float64).reshape(-1)
        )
        root = np.sqrt(weight)
        xw = design * root[:, None]
        yw = target * root
        if self.ridge > 0.0:
            penalty = np.sqrt(self.ridge) * np.eye(k, dtype=np.float64)
            if self.fit_intercept:
                penalty[0, 0] = 0.0
            xw = np.vstack([xw, penalty])
            yw = np.concatenate([yw, np.zeros(k, dtype=np.float64)])
        solution = np.linalg.lstsq(xw, yw, rcond=None)
        self._coef = np.asarray(solution[0], dtype=np.float64)
        return self

    def predict(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        if self._coef is None:
            raise RuntimeError("OLS.predict called before fit")
        return np.asarray(_design(X, self.fit_intercept) @ self._coef, dtype=np.float64)

    def get_params(self, deep: bool = True) -> dict[str, Any]:
        return {"fit_intercept": self.fit_intercept, "ridge": self.ridge}

    def set_params(self, **params: Any) -> OLS:
        for key, value in params.items():
            setattr(self, key, value)
        return self

OutcomeRegression

Bases: _MetaLearnerBase

Fit a model for the outcome and average the predicted contrast.

With per_arm=False (the default) one model carries the treatment as a feature (the S-learner / g-formula). With per_arm=True a separate model is fit on each arm (the T-learner). A Bayesian outcome learner flows its posterior into the per-unit CATE result and into cate_draws/cate_interval.

Source code in src/ergodic/inference/estimators/_regression.py

class OutcomeRegression(_MetaLearnerBase):
    """Fit a model for the outcome and average the predicted contrast.

    With ``per_arm=False`` (the default) one model carries the treatment as a
    feature (the S-learner / g-formula). With ``per_arm=True`` a separate model is
    fit on each arm (the T-learner). A Bayesian outcome learner flows its posterior
    into the per-unit CATE result and into ``cate_draws``/``cate_interval``.
    """

    name = "outcome_regression"
    slots = (Slot("outcome", Task.REGRESSION),)
    targets = frozenset({EffectTarget.ATE, EffectTarget.ATT})

    def __init__(self, learners: Any = None, *, per_arm: bool = False, **kwargs: Any) -> None:
        super().__init__(learners, **kwargs)
        self.per_arm = per_arm
        self._model: Learner | None = None
        self._arm_models: tuple[Learner, Learner] | None = None

    def _fit_cate(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> Learner | tuple[Learner, Learner]:
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        if self.per_arm:
            model1 = make["outcome"]()
            model1.fit(z[a == 1.0], y[a == 1.0])
            model0 = make["outcome"]()
            model0.fit(z[a == 0.0], y[a == 0.0])
            return (model0, model1)
        model = make["outcome"]()
        model.fit(np.hstack([a.reshape(-1, 1), z]), y)
        return model

    def _predict_cate(
        self, model: Learner | tuple[Learner, Learner], z: npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64]:
        if isinstance(model, tuple):
            model0, model1 = model
            return _mean_predict(model1, z) - _mean_predict(model0, z)
        n = z.shape[0]
        ones = np.ones((n, 1), dtype=np.float64)
        zeros = np.zeros((n, 1), dtype=np.float64)
        return _mean_predict(model, np.hstack([ones, z])) - _mean_predict(
            model, np.hstack([zeros, z])
        )

    def fit(self, identification: Identification, data: TabularData) -> Self:
        super().fit(identification, data)
        model = self._cate_model()
        if isinstance(model, tuple):
            self._arm_models = model
        else:
            self._model = model
        return self

    def _cate_draws(
        self, model: Learner | tuple[Learner, Learner], z: npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64] | None:
        mu1, mu0 = self._outcome_draws(z)
        return np.asarray(mu1 - mu0, dtype=np.float64)

    # -- the posterior path (Bayesian outcome learner) -----------------------

    def _outcome_draws(
        self, z: npt.NDArray[np.float64]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        """Posterior draws of mu1 and mu0 at ``z``, each ``(S, n)``.

        For the S-learner the two predictions read one coefficient draw, so the
        contrast is coherent within a draw; for the T-learner the arms are
        independent posteriors paired by index.
        """
        if self.per_arm:
            if self._arm_models is None:
                raise RuntimeError("outcome_regression is not fitted; call fit first")
            model0, model1 = self._arm_models
            mu1 = _mean_predict_draws(model1, z)
            mu0 = _mean_predict_draws(model0, z)
        else:
            if self._model is None:
                raise RuntimeError("outcome_regression is not fitted; call fit first")
            n = z.shape[0]
            ones = np.ones((n, 1), dtype=np.float64)
            zeros = np.zeros((n, 1), dtype=np.float64)
            mu1 = _mean_predict_draws(self._model, np.hstack([ones, z]))
            mu0 = _mean_predict_draws(self._model, np.hstack([zeros, z]))
        s = min(mu1.shape[0], mu0.shape[0])
        return mu1[:s], mu0[:s]

    def cate_draws(self, data: TabularData | None = None) -> npt.NDArray[np.float64]:
        """Per-unit CATE posterior (mu1 - mu0), shape ``(S, n)``.

        Needs a Bayesian outcome learner; in-sample when ``data`` is ``None``.
        """
        z = self._cate_covariates(data)
        mu1, mu0 = self._outcome_draws(z)
        return np.asarray(mu1 - mu0, dtype=np.float64)

    def cate_interval(
        self, data: TabularData | None = None, *, level: float = 0.95
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        """Per-unit credible interval for the CATE: the column quantiles of the draws."""
        draws = self.cate_draws(data)
        lo = np.quantile(draws, (1.0 - level) / 2.0, axis=0)
        hi = np.quantile(draws, (1.0 + level) / 2.0, axis=0)
        return np.asarray(lo, dtype=np.float64), np.asarray(hi, dtype=np.float64)

    def _posterior(self, state: _FitState, target: EffectTarget) -> npt.NDArray[np.float64] | None:
        a = state.arrays["A"]
        mu1, mu0 = self._outcome_draws(state.arrays["Z"])
        diff = mu1 - mu0
        if target is EffectTarget.ATT:
            return np.asarray(diff[:, a == 1.0].mean(axis=1), dtype=np.float64)
        return np.asarray(diff.mean(axis=1), dtype=np.float64)

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        n = a.shape[0]
        if self.per_arm:
            treated = a == 1.0
            control = a == 0.0
            model1 = make["outcome"]().fit(z[treated], y[treated])
            model0 = make["outcome"]().fit(z[control], y[control])
            mu1 = _mean_predict(model1, z)
            mu0 = _mean_predict(model0, z)
        else:
            model = make["outcome"]().fit(np.hstack([a.reshape(-1, 1), z]), y)
            ones = np.ones((n, 1), dtype=np.float64)
            zeros = np.zeros((n, 1), dtype=np.float64)
            mu1 = _mean_predict(model, np.hstack([ones, z]))
            mu0 = _mean_predict(model, np.hstack([zeros, z]))
        diff = mu1 - mu0
        if target is EffectTarget.ATT:
            return float(np.mean(diff[a == 1.0]))
        return float(np.mean(diff))

cate_draws(data=None)

Per-unit CATE posterior (mu1 - mu0), shape (S, n).

Needs a Bayesian outcome learner; in-sample when data is None.

Source code in src/ergodic/inference/estimators/_regression.py

def cate_draws(self, data: TabularData | None = None) -> npt.NDArray[np.float64]:
    """Per-unit CATE posterior (mu1 - mu0), shape ``(S, n)``.

    Needs a Bayesian outcome learner; in-sample when ``data`` is ``None``.
    """
    z = self._cate_covariates(data)
    mu1, mu0 = self._outcome_draws(z)
    return np.asarray(mu1 - mu0, dtype=np.float64)

cate_interval(data=None, *, level=0.95)

Per-unit credible interval for the CATE: the column quantiles of the draws.

Source code in src/ergodic/inference/estimators/_regression.py

def cate_interval(
    self, data: TabularData | None = None, *, level: float = 0.95
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Per-unit credible interval for the CATE: the column quantiles of the draws."""
    draws = self.cate_draws(data)
    lo = np.quantile(draws, (1.0 - level) / 2.0, axis=0)
    hi = np.quantile(draws, (1.0 + level) / 2.0, axis=0)
    return np.asarray(lo, dtype=np.float64), np.asarray(hi, dtype=np.float64)

PLIV

Bases: _IVBase

Orthogonal instrumental variables (the partially linear IV model).

Residualize the outcome, treatment, and instrument against the covariates with cross-fitting, then the IV moment on the residuals. The flexible-learner instrument estimator that two-stage least squares is not. Reports LATE with an analytic standard error.

Source code in src/ergodic/inference/estimators/_iv.py

class PLIV(_IVBase):
    """Orthogonal instrumental variables (the partially linear IV model).

    Residualize the outcome, treatment, and instrument against the covariates with
    cross-fitting, then the IV moment on the residuals. The flexible-learner
    instrument estimator that two-stage least squares is not. Reports LATE with an
    analytic standard error.
    """

    name = "pliv"
    slots = (
        Slot("outcome", Task.REGRESSION),
        Slot("treatment", Task.REGRESSION),
        Slot("instrument", Task.REGRESSION),
    )
    targets = frozenset({EffectTarget.LATE})

    def __init__(self, learners: Any = None, *, n_folds: int = 5, **kwargs: Any) -> None:
        super().__init__(learners, **kwargs)
        self.n_folds = n_folds
        self._cache: tuple[float, npt.NDArray[np.float64]] | None = None

    def _residualize(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        a = arrays["A"].reshape(-1)
        y = arrays["Y"]
        zi = arrays["Zi"]
        w = arrays["W"]
        n = y.shape[0]
        res_y = np.empty(n, dtype=np.float64)
        res_t = np.empty(n, dtype=np.float64)
        res_z = np.empty((n, zi.shape[1]), dtype=np.float64)
        rng = np.random.default_rng(self._random_state)
        has_w = w.shape[1] > 0
        for test in _folds(n, self.n_folds, rng):
            train = _train_mask(n, test)
            if has_w:
                res_y[test] = y[test] - _fit_predict(make["outcome"], w[train], y[train], w[test])
                res_t[test] = a[test] - _fit_predict(make["treatment"], w[train], a[train], w[test])
                for j in range(zi.shape[1]):
                    fitted = make["instrument"]().fit(w[train], zi[train, j])
                    res_z[test, j] = zi[test, j] - _mean_predict(fitted, w[test])
            else:
                res_y[test] = y[test] - float(np.mean(y[train]))
                res_t[test] = a[test] - float(np.mean(a[train]))
                for j in range(zi.shape[1]):
                    res_z[test, j] = zi[test, j] - float(np.mean(zi[train, j]))
        return res_y, res_t, res_z

    def _solve(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[float, npt.NDArray[np.float64]]:
        res_y, res_t, res_z = self._residualize(arrays, make)
        beta, influence = _two_sls(res_y, res_t.reshape(-1, 1), res_z)
        return float(beta[0]), influence[:, 0]

    def fit(self, identification: Identification, data: TabularData) -> Self:
        super().fit(identification, data)
        self._cache = self._solve(self._state.arrays, self._state.make)
        return self

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        return self._solve(arrays, make)[0]

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        return self._cache

PerUnitEffect dataclass

Per-unit treatment effects with their uncertainty.

_values is the per-unit point vector (n,), ordered to the input rows. Either _std_error (per-unit analytic SEs) or _draws (S, n) (bootstrap or posterior) carries the uncertainty, and the accessors read whichever is present. index are explicit row ids; estimand is set by the subclass.

Source code in src/ergodic/inference/_perunit.py

@dataclass(frozen=True, eq=False)
class PerUnitEffect:
    """Per-unit treatment effects with their uncertainty.

    ``_values`` is the per-unit point vector ``(n,)``, ordered to the input rows.
    Either ``_std_error`` (per-unit analytic SEs) or ``_draws`` ``(S, n)`` (bootstrap
    or posterior) carries the uncertainty, and the accessors read whichever is
    present. ``index`` are explicit row ids; ``estimand`` is set by the subclass.
    """

    treatment: str
    outcome: str
    covariates: tuple[str, ...]
    n: int
    kind: str  # "point", "analytic", "bootstrap", or "posterior"
    estimand: str = ""
    source: str = ""
    diagnostics: Mapping[str, float] = field(default_factory=dict)
    _values: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty(0))
    _index: npt.NDArray[Any] | None = field(default=None, repr=False)
    _std_error: npt.NDArray[np.float64] | None = field(default=None, repr=False)
    _draws: npt.NDArray[np.float64] | None = field(default=None, repr=False)

    # -- the per-unit surface ------------------------------------------------

    @property
    def values(self) -> npt.NDArray[np.float64]:
        """The per-unit point estimates, ``(n,)``, ordered to the input rows.

        When only draws are present, the posterior or bootstrap mean per unit.
        """
        if self._values.size or self._draws is None:
            return np.asarray(self._values, dtype=np.float64)
        return np.asarray(self._draws.mean(axis=0), dtype=np.float64)

    def __array__(self, dtype: Any = None) -> npt.NDArray[np.float64]:
        """The per-unit point vector, so ``np.asarray(result)`` and arithmetic work."""
        out = self.values
        return out.astype(dtype) if dtype is not None else out

    @property
    def shape(self) -> tuple[int, ...]:
        """The shape of the per-unit point vector, ``(n,)``."""
        return self.values.shape

    def __getitem__(self, item: Any) -> Any:
        """Index into the per-unit point vector (so ``result[i]`` is the unit's effect)."""
        return self.values[item]

    def __iter__(self) -> Any:
        return iter(self.values)

    # Element-wise comparisons against the point vector, so a wrapped result is a
    # drop-in for the bare array the SCM and the meta-learners returned before.
    def __eq__(self, other: Any) -> Any:
        return self.values == np.asarray(other)

    def __ne__(self, other: Any) -> Any:
        return self.values != np.asarray(other)

    def __lt__(self, other: Any) -> Any:
        return self.values < np.asarray(other)

    def __le__(self, other: Any) -> Any:
        return self.values <= np.asarray(other)

    def __gt__(self, other: Any) -> Any:
        return self.values > np.asarray(other)

    def __ge__(self, other: Any) -> Any:
        return self.values >= np.asarray(other)

    __hash__ = None  # type: ignore[assignment]

    @property
    def index(self) -> npt.NDArray[Any]:
        """The row ids for the units, positional ``arange(n)`` when none were given."""
        if self._index is not None:
            return np.asarray(self._index)
        return np.arange(self.n)

    @property
    def std_error(self) -> npt.NDArray[np.float64] | None:
        """Per-unit standard errors, or the per-unit draw standard deviations."""
        if self._std_error is not None:
            return np.asarray(self._std_error, dtype=np.float64)
        if self._draws is not None:
            return np.asarray(self._draws.std(axis=0, ddof=1), dtype=np.float64)
        return None

    def interval(
        self, level: float = 0.9
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        """Per-unit confidence or credible intervals, ``(lo[n], hi[n])`` at ``level``.

        Draw quantiles when draws are present, else the normal interval from the
        per-unit point and standard error. Raises when there is no uncertainty.
        """
        if not 0.0 < level < 1.0:
            raise ValueError(f"level must be in (0, 1), got {level}")
        if self._draws is not None:
            lo = np.quantile(self._draws, (1.0 - level) / 2.0, axis=0)
            hi = np.quantile(self._draws, (1.0 + level) / 2.0, axis=0)
            return np.asarray(lo, dtype=np.float64), np.asarray(hi, dtype=np.float64)
        if self._std_error is not None:
            z = _norm_ppf((1.0 + level) / 2.0)
            point = self.values
            se = np.asarray(self._std_error, dtype=np.float64)
            return point - z * se, point + z * se
        raise ValueError("per-unit effect has no uncertainty to form an interval from")

    def draws(self) -> npt.NDArray[np.float64] | None:
        """The ``(S, n)`` bootstrap or posterior draws, or ``None`` for a point result."""
        if self._draws is None:
            return None
        return np.array(self._draws, copy=True)

    def ate(self) -> EffectEstimate:
        """The population summary: the average per-unit effect, as an ``EffectEstimate``.

        The bridge back to the scalar family. With draws, the per-unit draws average
        to a draw vector of the mean effect (so the population interval reflects the
        joint uncertainty). With per-unit analytic SEs, the point is the mean and the
        standard error combines them assuming independence. With neither, a point.
        """
        target = EffectTarget.CATE if self.estimand == "cate" else EffectTarget.ITE
        learners: dict[str, str] = {}
        if self._draws is not None:
            mean_draws = np.asarray(self._draws.mean(axis=1), dtype=np.float64)
            return EffectEstimate(
                target=target,
                estimator=self.source,
                learners=learners,
                n=self.n,
                kind=self.kind if self.kind in ("bootstrap", "posterior") else "bootstrap",
                diagnostics=dict(self.diagnostics),
                _draws=mean_draws,
            )
        point = float(np.mean(self.values))
        if self._std_error is not None:
            se = np.asarray(self._std_error, dtype=np.float64)
            combined = float(np.sqrt(np.sum(se * se)) / se.shape[0])
            return EffectEstimate(
                target=target,
                estimator=self.source,
                learners=learners,
                n=self.n,
                kind="analytic",
                diagnostics=dict(self.diagnostics),
                _point=point,
                _std_error=combined,
            )
        return EffectEstimate(
            target=target,
            estimator=self.source,
            learners=learners,
            n=self.n,
            kind="point",
            diagnostics=dict(self.diagnostics),
            _point=point,
        )

    def to_frame(self) -> pd.DataFrame:
        """A per-unit table (index, effect, and lo/hi/se when available). Pandas only here."""
        import pandas as pd

        columns: dict[str, npt.NDArray[Any]] = {"effect": self.values}
        se = self.std_error
        if se is not None:
            columns["std_error"] = se
        try:
            lo, hi = self.interval()
            columns["ci_low"] = lo
            columns["ci_high"] = hi
        except ValueError:
            pass
        return pd.DataFrame(columns, index=pd.Index(self.index, name="unit"))

    def __len__(self) -> int:
        return self.n

    def __repr__(self) -> str:
        v = self.values
        head = f"{self.estimand or 'effect'}"
        body = (
            f"ate={float(np.mean(v)):.4g}, "
            f"min={float(np.min(v)):.3g}, median={float(np.median(v)):.3g}, "
            f"max={float(np.max(v)):.3g}"
        )
        name = type(self).__name__
        return f"{name}({self.source or head}, {self.kind}: {body}, n={self.n})"

index property

The row ids for the units, positional arange(n) when none were given.

shape property

The shape of the per-unit point vector, (n,).

std_error property

Per-unit standard errors, or the per-unit draw standard deviations.

values property

The per-unit point estimates, (n,), ordered to the input rows.

When only draws are present, the posterior or bootstrap mean per unit.

__array__(dtype=None)

The per-unit point vector, so np.asarray(result) and arithmetic work.

Source code in src/ergodic/inference/_perunit.py

def __array__(self, dtype: Any = None) -> npt.NDArray[np.float64]:
    """The per-unit point vector, so ``np.asarray(result)`` and arithmetic work."""
    out = self.values
    return out.astype(dtype) if dtype is not None else out

__getitem__(item)

Index into the per-unit point vector (so result[i] is the unit's effect).

Source code in src/ergodic/inference/_perunit.py

def __getitem__(self, item: Any) -> Any:
    """Index into the per-unit point vector (so ``result[i]`` is the unit's effect)."""
    return self.values[item]

ate()

The population summary: the average per-unit effect, as an EffectEstimate.

The bridge back to the scalar family. With draws, the per-unit draws average to a draw vector of the mean effect (so the population interval reflects the joint uncertainty). With per-unit analytic SEs, the point is the mean and the standard error combines them assuming independence. With neither, a point.

Source code in src/ergodic/inference/_perunit.py

def ate(self) -> EffectEstimate:
    """The population summary: the average per-unit effect, as an ``EffectEstimate``.

    The bridge back to the scalar family. With draws, the per-unit draws average
    to a draw vector of the mean effect (so the population interval reflects the
    joint uncertainty). With per-unit analytic SEs, the point is the mean and the
    standard error combines them assuming independence. With neither, a point.
    """
    target = EffectTarget.CATE if self.estimand == "cate" else EffectTarget.ITE
    learners: dict[str, str] = {}
    if self._draws is not None:
        mean_draws = np.asarray(self._draws.mean(axis=1), dtype=np.float64)
        return EffectEstimate(
            target=target,
            estimator=self.source,
            learners=learners,
            n=self.n,
            kind=self.kind if self.kind in ("bootstrap", "posterior") else "bootstrap",
            diagnostics=dict(self.diagnostics),
            _draws=mean_draws,
        )
    point = float(np.mean(self.values))
    if self._std_error is not None:
        se = np.asarray(self._std_error, dtype=np.float64)
        combined = float(np.sqrt(np.sum(se * se)) / se.shape[0])
        return EffectEstimate(
            target=target,
            estimator=self.source,
            learners=learners,
            n=self.n,
            kind="analytic",
            diagnostics=dict(self.diagnostics),
            _point=point,
            _std_error=combined,
        )
    return EffectEstimate(
        target=target,
        estimator=self.source,
        learners=learners,
        n=self.n,
        kind="point",
        diagnostics=dict(self.diagnostics),
        _point=point,
    )

draws()

The (S, n) bootstrap or posterior draws, or None for a point result.

Source code in src/ergodic/inference/_perunit.py

def draws(self) -> npt.NDArray[np.float64] | None:
    """The ``(S, n)`` bootstrap or posterior draws, or ``None`` for a point result."""
    if self._draws is None:
        return None
    return np.array(self._draws, copy=True)

interval(level=0.9)

Per-unit confidence or credible intervals, (lo[n], hi[n]) at level.

Draw quantiles when draws are present, else the normal interval from the per-unit point and standard error. Raises when there is no uncertainty.

Source code in src/ergodic/inference/_perunit.py

def interval(
    self, level: float = 0.9
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Per-unit confidence or credible intervals, ``(lo[n], hi[n])`` at ``level``.

    Draw quantiles when draws are present, else the normal interval from the
    per-unit point and standard error. Raises when there is no uncertainty.
    """
    if not 0.0 < level < 1.0:
        raise ValueError(f"level must be in (0, 1), got {level}")
    if self._draws is not None:
        lo = np.quantile(self._draws, (1.0 - level) / 2.0, axis=0)
        hi = np.quantile(self._draws, (1.0 + level) / 2.0, axis=0)
        return np.asarray(lo, dtype=np.float64), np.asarray(hi, dtype=np.float64)
    if self._std_error is not None:
        z = _norm_ppf((1.0 + level) / 2.0)
        point = self.values
        se = np.asarray(self._std_error, dtype=np.float64)
        return point - z * se, point + z * se
    raise ValueError("per-unit effect has no uncertainty to form an interval from")

to_frame()

A per-unit table (index, effect, and lo/hi/se when available). Pandas only here.

Source code in src/ergodic/inference/_perunit.py

def to_frame(self) -> pd.DataFrame:
    """A per-unit table (index, effect, and lo/hi/se when available). Pandas only here."""
    import pandas as pd

    columns: dict[str, npt.NDArray[Any]] = {"effect": self.values}
    se = self.std_error
    if se is not None:
        columns["std_error"] = se
    try:
        lo, hi = self.interval()
        columns["ci_low"] = lo
        columns["ci_high"] = hi
    except ValueError:
        pass
    return pd.DataFrame(columns, index=pd.Index(self.index, name="unit"))

PolicyTree dataclass

A shallow doubly-robust optimal-policy tree with its value and regret.

root is the learned tree over the named covariates; :meth:predict and :meth:assign evaluate it. value/std_error are the estimated policy value W(pi_hat) and its standard error from the per-unit influence terms; treat_all_value and treat_none_value are the two trivial baselines, and :meth:regret is the value gap against either. depth is the requested depth and exhaustive records whether the globally optimal search was used.

Source code in src/ergodic/inference/explain/_policy.py

@dataclass(frozen=True)
class PolicyTree:
    """A shallow doubly-robust optimal-policy tree with its value and regret.

    ``root`` is the learned tree over the named ``covariates``; :meth:`predict` and
    :meth:`assign` evaluate it. ``value``/``std_error`` are the estimated policy
    value ``W(pi_hat)`` and its standard error from the per-unit influence terms;
    ``treat_all_value`` and ``treat_none_value`` are the two trivial baselines, and
    :meth:`regret` is the value gap against either. ``depth`` is the requested depth
    and ``exhaustive`` records whether the globally optimal search was used.
    """

    root: PolicyNode
    covariates: tuple[str, ...]
    n: int
    depth: int
    value: float
    std_error: float
    treat_all_value: float
    treat_none_value: float
    exhaustive: bool
    treatment: str = ""
    outcome: str = ""
    _scores: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty(0), repr=False)
    _design: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty((0, 0)), repr=False)

    # -- evaluation ----------------------------------------------------------

    def predict(self, covariates: npt.NDArray[np.float64]) -> npt.NDArray[np.int64]:
        """The per-row action (0/1) for a covariate matrix aligned to ``covariates``."""
        x = np.asarray(covariates, dtype=np.float64)
        if x.ndim == 1:
            x = x.reshape(1, -1)
        if x.shape[1] != len(self.covariates):
            raise ValueError(
                f"covariate matrix has {x.shape[1]} columns, expected {len(self.covariates)}"
            )
        return np.asarray([self.root.predict(row) for row in x], dtype=np.int64)

    def assign(self) -> npt.NDArray[np.int64]:
        """The fitted policy's action for each unit it was learned on."""
        return self.predict(self._design)

    def interval(self, level: float = 0.9) -> tuple[float, float]:
        """A normal confidence interval for the policy value at ``level``."""
        from .._stats import _norm_ppf

        if not 0.0 < level < 1.0:
            raise ValueError(f"level must be in (0, 1), got {level}")
        z = _norm_ppf((1.0 + level) / 2.0)
        return self.value - z * self.std_error, self.value + z * self.std_error

    def regret(self, baseline: str = "treat_all") -> float:
        """The learned policy's welfare GAIN over a trivial baseline (the baseline's regret).

        ``baseline`` is ``"treat_all"``, ``"treat_none"``, or ``"best_trivial"`` (the
        better of the two). Returns ``value - baseline_value``: how much welfare the
        learned rule adds over treating everyone / no one. The trivial constant
        policies are depth-0 leaves inside the tree's class, so the optimal tree is at
        least as good and this gain is non-negative up to sampling noise (it is the
        baseline's regret against the learned optimum, not the optimum's own regret,
        which would need the unknown oracle value).
        """
        if baseline == "treat_all":
            ref = self.treat_all_value
        elif baseline == "treat_none":
            ref = self.treat_none_value
        elif baseline == "best_trivial":
            ref = max(self.treat_all_value, self.treat_none_value)
        else:
            raise ValueError(
                f"baseline must be 'treat_all', 'treat_none', or 'best_trivial', got {baseline!r}"
            )
        return self.value - ref

    def leaves(self) -> list[PolicyNode]:
        """The leaf nodes, left to right (each carries its action, count, and value)."""
        out: list[PolicyNode] = []

        def walk(node: PolicyNode) -> None:
            if node.is_leaf:
                out.append(node)
                return
            assert node.left is not None and node.right is not None
            walk(node.left)
            walk(node.right)

        walk(self.root)
        return out

    def describe(self) -> str:
        """A readable ``if-else`` rendering of the tree (the interpretable rule)."""
        lines: list[str] = []

        def walk(node: PolicyNode, indent: str) -> None:
            if node.is_leaf:
                lines.append(f"{indent}-> treat={node.action} (n={node.n})")
                return
            feat = node.feature
            thr = node.threshold
            assert thr is not None
            lines.append(f"{indent}if {feat} <= {thr:.4g}:")
            assert node.left is not None and node.right is not None
            walk(node.left, indent + "  ")
            lines.append(f"{indent}else:  # {feat} > {thr:.4g}")
            walk(node.right, indent + "  ")

        walk(self.root, "")
        return "\n".join(lines)

assign()

The fitted policy's action for each unit it was learned on.

Source code in src/ergodic/inference/explain/_policy.py

def assign(self) -> npt.NDArray[np.int64]:
    """The fitted policy's action for each unit it was learned on."""
    return self.predict(self._design)

describe()

A readable if-else rendering of the tree (the interpretable rule).

Source code in src/ergodic/inference/explain/_policy.py

def describe(self) -> str:
    """A readable ``if-else`` rendering of the tree (the interpretable rule)."""
    lines: list[str] = []

    def walk(node: PolicyNode, indent: str) -> None:
        if node.is_leaf:
            lines.append(f"{indent}-> treat={node.action} (n={node.n})")
            return
        feat = node.feature
        thr = node.threshold
        assert thr is not None
        lines.append(f"{indent}if {feat} <= {thr:.4g}:")
        assert node.left is not None and node.right is not None
        walk(node.left, indent + "  ")
        lines.append(f"{indent}else:  # {feat} > {thr:.4g}")
        walk(node.right, indent + "  ")

    walk(self.root, "")
    return "\n".join(lines)

interval(level=0.9)

A normal confidence interval for the policy value at level.

Source code in src/ergodic/inference/explain/_policy.py

def interval(self, level: float = 0.9) -> tuple[float, float]:
    """A normal confidence interval for the policy value at ``level``."""
    from .._stats import _norm_ppf

    if not 0.0 < level < 1.0:
        raise ValueError(f"level must be in (0, 1), got {level}")
    z = _norm_ppf((1.0 + level) / 2.0)
    return self.value - z * self.std_error, self.value + z * self.std_error

leaves()

The leaf nodes, left to right (each carries its action, count, and value).

Source code in src/ergodic/inference/explain/_policy.py

def leaves(self) -> list[PolicyNode]:
    """The leaf nodes, left to right (each carries its action, count, and value)."""
    out: list[PolicyNode] = []

    def walk(node: PolicyNode) -> None:
        if node.is_leaf:
            out.append(node)
            return
        assert node.left is not None and node.right is not None
        walk(node.left)
        walk(node.right)

    walk(self.root)
    return out

predict(covariates)

The per-row action (0/1) for a covariate matrix aligned to covariates.

Source code in src/ergodic/inference/explain/_policy.py

def predict(self, covariates: npt.NDArray[np.float64]) -> npt.NDArray[np.int64]:
    """The per-row action (0/1) for a covariate matrix aligned to ``covariates``."""
    x = np.asarray(covariates, dtype=np.float64)
    if x.ndim == 1:
        x = x.reshape(1, -1)
    if x.shape[1] != len(self.covariates):
        raise ValueError(
            f"covariate matrix has {x.shape[1]} columns, expected {len(self.covariates)}"
        )
    return np.asarray([self.root.predict(row) for row in x], dtype=np.int64)

regret(baseline='treat_all')

The learned policy's welfare GAIN over a trivial baseline (the baseline's regret).

baseline is "treat_all", "treat_none", or "best_trivial" (the better of the two). Returns value - baseline_value: how much welfare the learned rule adds over treating everyone / no one. The trivial constant policies are depth-0 leaves inside the tree's class, so the optimal tree is at least as good and this gain is non-negative up to sampling noise (it is the baseline's regret against the learned optimum, not the optimum's own regret, which would need the unknown oracle value).

Source code in src/ergodic/inference/explain/_policy.py

def regret(self, baseline: str = "treat_all") -> float:
    """The learned policy's welfare GAIN over a trivial baseline (the baseline's regret).

    ``baseline`` is ``"treat_all"``, ``"treat_none"``, or ``"best_trivial"`` (the
    better of the two). Returns ``value - baseline_value``: how much welfare the
    learned rule adds over treating everyone / no one. The trivial constant
    policies are depth-0 leaves inside the tree's class, so the optimal tree is at
    least as good and this gain is non-negative up to sampling noise (it is the
    baseline's regret against the learned optimum, not the optimum's own regret,
    which would need the unknown oracle value).
    """
    if baseline == "treat_all":
        ref = self.treat_all_value
    elif baseline == "treat_none":
        ref = self.treat_none_value
    elif baseline == "best_trivial":
        ref = max(self.treat_all_value, self.treat_none_value)
    else:
        raise ValueError(
            f"baseline must be 'treat_all', 'treat_none', or 'best_trivial', got {baseline!r}"
        )
    return self.value - ref

PyMCClassifier

Bases: PyMCLearner

A Bayesian classification learner over a PyMC model.

Name the probability node "p" (the per-row P(class=1)). Exposes predict_proba and predict_proba_draws, so it fills a propensity or a mediator slot and drives the posterior path there.

Source code in src/ergodic/inference/learners/_pymc.py

class PyMCClassifier(PyMCLearner):
    """A Bayesian classification learner over a PyMC model.

    Name the probability node ``"p"`` (the per-row P(class=1)). Exposes
    ``predict_proba`` and ``predict_proba_draws``, so it fills a propensity or a
    mediator slot and drives the posterior path there.
    """

    _mean_node = "p"

    def predict_proba_draws(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return self._mean_draws(X)

    def predict_proba(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        p = self._mean_draws(X).mean(axis=0)
        return np.asarray(np.column_stack([1.0 - p, p]), dtype=np.float64)

    def predict(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return np.asarray(self.predict_proba(X)[:, 1] >= 0.5, dtype=np.float64)

PyMCLearner

A Bayesian regression learner over a PyMC model.

Implement :meth:build to return a pymc.Model. predict is the posterior predictive mean and predict_draws the (S, n) posterior of the mean node, so the learner drops into a regression slot and also drives the posterior path.

Source code in src/ergodic/inference/learners/_pymc.py

class PyMCLearner:
    """A Bayesian regression learner over a PyMC model.

    Implement :meth:`build` to return a ``pymc.Model``. ``predict`` is the
    posterior predictive mean and ``predict_draws`` the ``(S, n)`` posterior of the
    mean node, so the learner drops into a regression slot and also drives the
    posterior path.
    """

    _mean_node = "mu"
    _data_name = "X"

    def __init__(
        self,
        *,
        draws: int = 1000,
        tune: int = 1000,
        chains: int = 2,
        target_accept: float = 0.9,
        random_state: int | None = None,
        progressbar: bool = False,
        **build_kwargs: Any,
    ) -> None:
        self.draws = draws
        self.tune = tune
        self.chains = chains
        self.target_accept = target_accept
        self.random_state = random_state
        self.progressbar = progressbar
        self.build_kwargs = build_kwargs
        self._model: Any = None
        self._idata: Any = None

    def build(self, X: npt.NDArray[np.float64], y: npt.NDArray[np.float64]) -> Any:
        """Return a ``pymc.Model``. Subclasses implement this."""
        raise NotImplementedError("subclass PyMCLearner and implement build")

    def fit(
        self,
        X: npt.NDArray[np.float64],
        y: npt.NDArray[np.float64],
        sample_weight: npt.NDArray[np.float64] | None = None,
    ) -> PyMCLearner:
        import pymc as pm

        x = np.asarray(X, dtype=np.float64)
        target = np.asarray(y, dtype=np.float64).reshape(-1)
        self._model = self.build(x, target)
        with self._model:
            self._idata = pm.sample(
                draws=self.draws,
                tune=self.tune,
                chains=self.chains,
                target_accept=self.target_accept,
                random_seed=self.random_state,
                progressbar=self.progressbar,
            )
        return self

    def _mean_draws(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        import pymc as pm

        if self._idata is None:
            raise RuntimeError(f"{type(self).__name__}.predict called before fit")
        x = np.asarray(X, dtype=np.float64)
        with self._model:
            pm.set_data({self._data_name: x})
            posterior = pm.sample_posterior_predictive(
                self._idata,
                var_names=[self._mean_node],
                progressbar=self.progressbar,
                random_seed=self.random_state,
            )
        arr = np.asarray(
            posterior.posterior_predictive[self._mean_node].to_numpy(), dtype=np.float64
        )
        return arr.reshape(-1, arr.shape[-1])

    def predict_draws(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return self._mean_draws(X)

    def predict(self, X: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return np.asarray(self._mean_draws(X).mean(axis=0), dtype=np.float64)

    def get_params(self, deep: bool = True) -> dict[str, Any]:
        params = {key: getattr(self, key) for key in _SAMPLER_KEYS}
        params.update(self.build_kwargs)
        return params

    def set_params(self, **params: Any) -> PyMCLearner:
        for key, value in params.items():
            if key in _SAMPLER_KEYS:
                setattr(self, key, value)
            else:
                self.build_kwargs[key] = value
        return self

build(X, y)

Return a pymc.Model. Subclasses implement this.

Source code in src/ergodic/inference/learners/_pymc.py

def build(self, X: npt.NDArray[np.float64], y: npt.NDArray[np.float64]) -> Any:
    """Return a ``pymc.Model``. Subclasses implement this."""
    raise NotImplementedError("subclass PyMCLearner and implement build")

RLearner

Bases: _MetaLearnerBase

The R-learner: residualize the outcome and treatment, then fit the CATE.

Robinson's residual-on-residual moment with a heterogeneous final model, weighted by the squared treatment residual. The ATE is the average of the fitted CATE, with a bootstrap standard error.

Source code in src/ergodic/inference/estimators/_meta.py

class RLearner(_MetaLearnerBase):
    """The R-learner: residualize the outcome and treatment, then fit the CATE.

    Robinson's residual-on-residual moment with a heterogeneous final model,
    weighted by the squared treatment residual. The ATE is the average of the
    fitted CATE, with a bootstrap standard error.
    """

    name = "r_learner"
    slots = (
        Slot("outcome", Task.REGRESSION),
        Slot("propensity", Task.CLASSIFICATION),
        Slot("final", Task.REGRESSION),
    )
    targets = frozenset({EffectTarget.ATE})

    def __init__(
        self, learners: Any = None, *, n_folds: int = 5, clip: float = 1e-3, **kwargs: Any
    ) -> None:
        super().__init__(learners, **kwargs)
        self.n_folds = n_folds
        self.clip = clip

    def _residualize(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        n = a.shape[0]
        res_y = np.empty(n, dtype=np.float64)
        res_t = np.empty(n, dtype=np.float64)
        rng = np.random.default_rng(self._random_state)
        for test in _folds(n, self.n_folds, rng, stratify=a):
            train = _train_mask(n, test)
            res_y[test] = y[test] - _fit_predict(make["outcome"], z[train], y[train], z[test])
            e = np.clip(
                _mean_predict(make["propensity"]().fit(z[train], a[train]), z[test]),
                self.clip,
                1.0 - self.clip,
            )
            res_t[test] = a[test] - e
        return res_y, res_t

    def _fit_cate(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> Learner:
        res_y, res_t = self._residualize(arrays, make)
        pseudo = res_y / res_t
        weight = res_t * res_t
        final = make["final"]()
        try:
            final.fit(arrays["Z"], pseudo, sample_weight=weight)
        except TypeError as exc:
            raise TypeError(
                "slot 'final' needs a learner whose fit accepts sample_weight; "
                f"got {type(final).__name__}"
            ) from exc
        return final

    def _predict_cate(self, model: Learner, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        return _mean_predict(model, z)

Refutation dataclass

The result of one refutation check.

estimand_value is the effect on the real data. refuted_value is the mean effect over the perturbed runs, with refuted_std their spread. passed is the convenience verdict (the rule is in the method's section of :func:refute_effect); read the numbers, not just the flag. p_value is the permutation tail for the placebo check and None for the others.

Source code in src/ergodic/inference/_refute.py

@dataclass(frozen=True)
class Refutation:
    """The result of one refutation check.

    ``estimand_value`` is the effect on the real data. ``refuted_value`` is the
    mean effect over the perturbed runs, with ``refuted_std`` their spread.
    ``passed`` is the convenience verdict (the rule is in the method's section of
    :func:`refute_effect`); read the numbers, not just the flag. ``p_value`` is
    the permutation tail for the placebo check and ``None`` for the others.
    """

    method: str
    estimand_value: float
    refuted_value: float
    refuted_std: float
    p_value: float | None
    passed: bool
    n_simulations: int
    diagnostics: Mapping[str, float] = field(default_factory=dict)

    def __repr__(self) -> str:
        verdict = "passed" if self.passed else "FAILED"
        body = (
            f"{self.method}: estimate={self.estimand_value:.4g}, "
            f"under perturbation={self.refuted_value:.4g}"
        )
        if self.p_value is not None:
            body += f", p={self.p_value:.3g}"
        return f"Refutation({body}, {verdict}, n={self.n_simulations})"

SelectionReport dataclass

The ranked CATE model-selection result with the chosen and aggregated models.

ranking is the candidates sorted by ascending DR tau-risk (best first). selected is the fitted :class:~ergodic.inference.CATEResult of the lowest risk; selected_method names it. aggregate is the Q-aggregation ensemble result and weights its convex weights (aligned to aggregate_methods), with aggregate_tau_risk the ensemble's DR tau-risk.

Source code in src/ergodic/inference/_select.py

@dataclass(frozen=True)
class SelectionReport:
    """The ranked CATE model-selection result with the chosen and aggregated models.

    ``ranking`` is the candidates sorted by ascending DR tau-risk (best first).
    ``selected`` is the fitted :class:`~ergodic.inference.CATEResult` of the lowest
    risk; ``selected_method`` names it. ``aggregate`` is the Q-aggregation ensemble
    result and ``weights`` its convex weights (aligned to ``aggregate_methods``),
    with ``aggregate_tau_risk`` the ensemble's DR tau-risk.
    """

    ranking: tuple[CandidateScore, ...]
    selected_method: str
    selected: CATEResult
    aggregate: CATEResult
    aggregate_methods: tuple[str, ...]
    weights: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty(0))
    aggregate_tau_risk: float = float("nan")

    def risk(self, method: str) -> float:
        """The DR tau-risk of a candidate by name (raises when it was not scored)."""
        for score in self.ranking:
            if score.method == method:
                return score.tau_risk
        raise KeyError(f"{method!r} is not among the scored candidates")

    @property
    def best_single_tau_risk(self) -> float:
        """The DR tau-risk of the best single candidate (the first in the ranking)."""
        return self.ranking[0].tau_risk

    def frame(self) -> dict[str, float]:
        """A method-to-risk mapping (the aggregate included), handy for tables/tests."""
        out = {score.method: score.tau_risk for score in self.ranking}
        out["q_aggregate"] = self.aggregate_tau_risk
        return out

best_single_tau_risk property

The DR tau-risk of the best single candidate (the first in the ranking).

frame()

A method-to-risk mapping (the aggregate included), handy for tables/tests.

Source code in src/ergodic/inference/_select.py

def frame(self) -> dict[str, float]:
    """A method-to-risk mapping (the aggregate included), handy for tables/tests."""
    out = {score.method: score.tau_risk for score in self.ranking}
    out["q_aggregate"] = self.aggregate_tau_risk
    return out

risk(method)

The DR tau-risk of a candidate by name (raises when it was not scored).

Source code in src/ergodic/inference/_select.py

def risk(self, method: str) -> float:
    """The DR tau-risk of a candidate by name (raises when it was not scored)."""
    for score in self.ranking:
        if score.method == method:
            return score.tau_risk
    raise KeyError(f"{method!r} is not among the scored candidates")

Slot dataclass

A named nuisance position with the learning task it requires.

Source code in src/ergodic/inference/_estimator.py

@dataclass(frozen=True)
class Slot:
    """A named nuisance position with the learning task it requires."""

    name: str
    task: Task

Stratification

Bases: _AdjustmentBase

Propensity stratification: bin on the propensity and difference within strata.

The estimate is the stratum-size-weighted (ATE) or treated-size-weighted (ATT) average of the within-stratum mean difference. Standard error by bootstrap.

Source code in src/ergodic/inference/estimators/_weighting.py

class Stratification(_AdjustmentBase):
    """Propensity stratification: bin on the propensity and difference within strata.

    The estimate is the stratum-size-weighted (ATE) or treated-size-weighted (ATT)
    average of the within-stratum mean difference. Standard error by bootstrap.
    """

    name = "stratification"
    slots = (Slot("propensity", Task.CLASSIFICATION),)
    targets = frozenset({EffectTarget.ATE, EffectTarget.ATT})

    def __init__(
        self, learners: Any = None, *, n_strata: int = 5, clip: float = 1e-3, **kwargs: Any
    ) -> None:
        super().__init__(learners, **kwargs)
        if n_strata < 2:
            raise ValueError(f"n_strata must be >= 2, got {n_strata}")
        self.n_strata = n_strata
        self.clip = clip

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        e = np.clip(_mean_predict(make["propensity"]().fit(z, a), z), self.clip, 1.0 - self.clip)
        edges = np.quantile(e, np.linspace(0.0, 1.0, self.n_strata + 1))
        bins = np.searchsorted(edges[1:-1], e, side="right")
        total = 0.0
        accum = 0.0
        for stratum in range(self.n_strata):
            here = bins == stratum
            treated = here & (a == 1.0)
            control = here & (a == 0.0)
            n_treated = int(np.sum(treated))
            n_control = int(np.sum(control))
            if n_treated == 0 or n_control == 0:
                continue
            diff = float(np.mean(y[treated]) - np.mean(y[control]))
            weight = (
                float(n_treated) if target is EffectTarget.ATT else float(n_treated + n_control)
            )
            accum += weight * diff
            total += weight
        if total <= 0.0:
            raise ValueError("no stratum has both a treated and a control unit")
        return accum / total

SyntheticControlEffect dataclass

Bases: EffectEstimate

A synthetic-control fit: the weights, the two paths, and placebo inference.

observed and synthetic run over periods, and gaps is their difference. pre_rmspe is the tracking error before at, the fit diagnostic to judge before trusting the effect. p_value ranks the treated unit's post-to-pre RMSPE ratio among the donor placebos, and placebo_gaps holds each donor's post-period placebo gaps. There is no sampling story for one treated unit, so std_error stays None on purpose and interval is a placebo band instead.

Source code in src/ergodic/inference/estimators/_synthetic_control.py

@dataclass(frozen=True)
class SyntheticControlEffect(EffectEstimate):
    """A synthetic-control fit: the weights, the two paths, and placebo inference.

    ``observed`` and ``synthetic`` run over ``periods``, and ``gaps`` is their
    difference. ``pre_rmspe`` is the tracking error before ``at``, the fit
    diagnostic to judge before trusting the effect. ``p_value`` ranks the
    treated unit's post-to-pre RMSPE ratio among the donor placebos, and
    ``placebo_gaps`` holds each donor's post-period placebo gaps. There is no
    sampling story for one treated unit, so ``std_error`` stays ``None`` on
    purpose and ``interval`` is a placebo band instead.
    """

    design: Design | None = None
    treated_unit: Any = None
    at: Any = None
    weights: Mapping[Any, float] = field(default_factory=dict)
    periods: tuple[Any, ...] = ()
    observed: tuple[float, ...] = ()
    synthetic: tuple[float, ...] = ()
    pre_rmspe: float = float("nan")
    post_rmspe: float = float("nan")
    p_value: float | None = None
    placebo_gaps: Mapping[Any, tuple[float, ...]] = field(default_factory=dict)

    @property
    def gaps(self) -> tuple[float, ...]:
        """Observed minus synthetic, one value per period."""
        return tuple(o - s for o, s in zip(self.observed, self.synthetic, strict=True))

    def interval(self, level: float = 0.95) -> tuple[float, float]:
        """A placebo-permutation band around the estimate, not a sampling CI.

        The half width is the ``level`` quantile of the absolute placebo
        post-gap means: the gap a unit with no treatment can show. An estimate
        whose band excludes zero clears the placebo bar at that level.
        """
        if not 0.0 < level < 1.0:
            raise ValueError(f"level must be in (0, 1), got {level}")
        if not self.placebo_gaps:
            raise ValueError("estimate has no placebo gaps to form a band from")
        means = np.asarray(
            [abs(float(np.mean(gaps))) for gaps in self.placebo_gaps.values()],
            dtype=np.float64,
        )
        half = float(np.quantile(means, level))
        return (self.value - half, self.value + half)

gaps property

Observed minus synthetic, one value per period.

interval(level=0.95)

A placebo-permutation band around the estimate, not a sampling CI.

The half width is the level quantile of the absolute placebo post-gap means: the gap a unit with no treatment can show. An estimate whose band excludes zero clears the placebo bar at that level.

Source code in src/ergodic/inference/estimators/_synthetic_control.py

def interval(self, level: float = 0.95) -> tuple[float, float]:
    """A placebo-permutation band around the estimate, not a sampling CI.

    The half width is the ``level`` quantile of the absolute placebo
    post-gap means: the gap a unit with no treatment can show. An estimate
    whose band excludes zero clears the placebo bar at that level.
    """
    if not 0.0 < level < 1.0:
        raise ValueError(f"level must be in (0, 1), got {level}")
    if not self.placebo_gaps:
        raise ValueError("estimate has no placebo gaps to form a band from")
    means = np.asarray(
        [abs(float(np.mean(gaps))) for gaps in self.placebo_gaps.values()],
        dtype=np.float64,
    )
    half = float(np.quantile(means, level))
    return (self.value - half, self.value + half)

Task

Bases: Enum

What a nuisance slot asks of its learner.

Source code in src/ergodic/inference/_estimator.py

class Task(Enum):
    """What a nuisance slot asks of its learner."""

    REGRESSION = "regression"
    CLASSIFICATION = "classification"  # propensity or a discrete outcome

    def __str__(self) -> str:
        return self.value

TwoStageLeastSquares

Bases: _IVBase

Two-stage least squares. OLS in both stages, robust standard errors.

The first stage regresses the treatment on the instruments and any conditioning covariates; the second regresses the outcome on the fitted treatment. Reports the local average treatment effect (LATE).

Source code in src/ergodic/inference/estimators/_iv.py

class TwoStageLeastSquares(_IVBase):
    """Two-stage least squares. OLS in both stages, robust standard errors.

    The first stage regresses the treatment on the instruments and any
    conditioning covariates; the second regresses the outcome on the fitted
    treatment. Reports the local average treatment effect (LATE).
    """

    name = "two_stage_least_squares"
    targets = frozenset({EffectTarget.LATE})

    def __init__(self, learners: Any = None, **kwargs: Any) -> None:
        super().__init__(learners, **kwargs)

    def _design(
        self, arrays: dict[str, npt.NDArray[np.float64]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], int]:
        a = arrays["A"]
        y = arrays["Y"]
        zi = arrays["Zi"]
        w = arrays["W"]
        n = y.shape[0]
        ones = np.ones((n, 1), dtype=np.float64)
        endog = np.hstack([ones, w, a])
        instr = np.hstack([ones, w, zi])
        t_index = 1 + w.shape[1]
        return y, endog, instr, t_index

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        y, endog, instr, t_index = self._design(arrays)
        beta, _ = _two_sls(y, endog, instr)
        return float(beta[t_index])

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        y, endog, instr, t_index = self._design(arrays)
        beta, influence = _two_sls(y, endog, instr)
        return float(beta[t_index]), influence[:, t_index]

TwoStageRegression

Bases: _FrontdoorBase

Estimate the front-door functional from a mediator model and an outcome model.

For binary treatment A and binary mediator M: a classifier gives P(M | A), a regression gives E[Y | M, A], and the effect is the front-door sum E[Y | do(A=1)] - E[Y | do(A=0)]. Continuous and multi-level mediators come later.

Source code in src/ergodic/inference/estimators/_frontdoor.py

class TwoStageRegression(_FrontdoorBase):
    """Estimate the front-door functional from a mediator model and an outcome model.

    For binary treatment ``A`` and binary mediator ``M``: a classifier gives
    ``P(M | A)``, a regression gives ``E[Y | M, A]``, and the effect is the
    front-door sum ``E[Y | do(A=1)] - E[Y | do(A=0)]``. Continuous and multi-level
    mediators come later.
    """

    name = "two_stage_regression"
    slots = (Slot("mediator", Task.CLASSIFICATION), Slot("outcome", Task.REGRESSION))
    targets = frozenset({EffectTarget.ATE})

    def __init__(self, learners: Any = None, **kwargs: Any) -> None:
        super().__init__(learners, **kwargs)

    def _point(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> float:
        a = arrays["A"]
        m = arrays["M"]
        y = arrays["Y"]
        mediator = make["mediator"]().fit(a.reshape(-1, 1), m)
        p_m1_a1 = float(_mean_predict(mediator, np.array([[1.0]]))[0])
        p_m1_a0 = float(_mean_predict(mediator, np.array([[0.0]]))[0])
        outcome = make["outcome"]().fit(np.column_stack([a, m, a * m]), y)

        def e_outcome(m_value: float, a_value: float) -> float:
            row = np.array([[a_value, m_value, a_value * m_value]], dtype=np.float64)
            return float(_mean_predict(outcome, row)[0])

        p_a1 = float(np.mean(a))

        def do(a_value: float) -> float:
            p_m1 = p_m1_a1 if a_value == 1.0 else p_m1_a0
            total = 0.0
            for m_value, p_m in ((1.0, p_m1), (0.0, 1.0 - p_m1)):
                inner = e_outcome(m_value, 1.0) * p_a1 + e_outcome(m_value, 0.0) * (1.0 - p_a1)
                total += p_m * inner
            return total

        return do(1.0) - do(0.0)

VariableImportance dataclass

Depth-weighted, normalized split frequencies: a driver signal over covariates.

names are the covariate columns and values the matching importances, which sum to one (or are all zero for a forest that never split). A split at depth d contributes 1 / (d + 1), so shallow splits, the ones that move the effect most, weigh more, following the GRF variable-importance convention.

Source code in src/ergodic/inference/estimators/_forest.py

@dataclass(frozen=True)
class VariableImportance:
    """Depth-weighted, normalized split frequencies: a driver signal over covariates.

    ``names`` are the covariate columns and ``values`` the matching importances, which
    sum to one (or are all zero for a forest that never split). A split at depth ``d``
    contributes ``1 / (d + 1)``, so shallow splits, the ones that move the effect most,
    weigh more, following the GRF variable-importance convention.
    """

    names: tuple[str, ...]
    values: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty(0))

    def top(self, k: int | None = None) -> list[tuple[str, float]]:
        """The covariates by descending importance (all of them when ``k`` is ``None``)."""
        order = np.argsort(self.values)[::-1]
        pairs = [(self.names[i], float(self.values[i])) for i in order]
        return pairs if k is None else pairs[:k]

    def to_dict(self) -> dict[str, float]:
        return {name: float(value) for name, value in zip(self.names, self.values, strict=True)}

top(k=None)

The covariates by descending importance (all of them when k is None).

Source code in src/ergodic/inference/estimators/_forest.py

def top(self, k: int | None = None) -> list[tuple[str, float]]:
    """The covariates by descending importance (all of them when ``k`` is ``None``)."""
    order = np.argsort(self.values)[::-1]
    pairs = [(self.names[i], float(self.values[i])) for i in order]
    return pairs if k is None else pairs[:k]

Wald

Bases: TwoStageLeastSquares

The Wald ratio: two-stage least squares with a single instrument.

Source code in src/ergodic/inference/estimators/_iv.py

class Wald(TwoStageLeastSquares):
    """The Wald ratio: two-stage least squares with a single instrument."""

    name = "wald"

    def _iv_columns(
        self, instruments: list[str], conditioning: list[str]
    ) -> tuple[list[str], list[str]]:
        if len(instruments) != 1:
            raise NotImplementedError(
                "Wald needs exactly one instrument; use two_stage_least_squares for more"
            )
        return instruments, conditioning

XLearner

Bases: _MetaLearnerBase

The X-learner: per-arm outcome models, imputed effects, propensity-weighted.

Fit one outcome model per arm, impute each unit's effect against the other arm's model, regress those imputed effects per arm, and combine the two with the propensity. The ATE is the average of the fitted CATE, with a bootstrap standard error.

Source code in src/ergodic/inference/estimators/_meta.py

class XLearner(_MetaLearnerBase):
    """The X-learner: per-arm outcome models, imputed effects, propensity-weighted.

    Fit one outcome model per arm, impute each unit's effect against the other
    arm's model, regress those imputed effects per arm, and combine the two with
    the propensity. The ATE is the average of the fitted CATE, with a bootstrap
    standard error.
    """

    name = "x_learner"
    slots = (Slot("outcome", Task.REGRESSION), Slot("propensity", Task.CLASSIFICATION))
    targets = frozenset({EffectTarget.ATE})

    def __init__(self, learners: Any = None, *, clip: float = 1e-3, **kwargs: Any) -> None:
        super().__init__(learners, **kwargs)
        self.clip = clip

    def _fit_cate(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[Learner, Learner, Learner]:
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        treated = a == 1.0
        control = a == 0.0
        mu1 = make["outcome"]().fit(z[treated], y[treated])
        mu0 = make["outcome"]().fit(z[control], y[control])
        imputed_treated = y[treated] - _mean_predict(mu0, z[treated])
        imputed_control = _mean_predict(mu1, z[control]) - y[control]
        tau1 = make["outcome"]().fit(z[treated], imputed_treated)
        tau0 = make["outcome"]().fit(z[control], imputed_control)
        propensity = make["propensity"]().fit(z, a)
        return tau0, tau1, propensity

    def _predict_cate(
        self, model: tuple[Learner, Learner, Learner], z: npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64]:
        tau0, tau1, propensity = model
        e = np.clip(_mean_predict(propensity, z), self.clip, 1.0 - self.clip)
        return e * _mean_predict(tau0, z) + (1.0 - e) * _mean_predict(tau1, z)

candidates(identification)

Every estimator that consumes the identified family, default first.

Source code in src/ergodic/inference/_registry.py

def candidates(identification: Identification) -> list[type[CausalEstimator]]:
    """Every estimator that consumes the identified family, default first."""
    return list(_BY_FAMILY.get(identification.strategy, []))

cate(graph, data, treatment, outcome, *, method='auto', learner=None, treated=None, random_state=None, n_boot=1000, uncertainty=True, **kwargs)

Identify the effect of treatment on outcome, then estimate the CATE.

Identification-gated exactly like :func:estimate_effect: identify the effect, select a CATE-capable estimator for the family, fit it, and return a :class:~ergodic.inference.CATEResult. method names the meta-learner ("dr_learner" by default, the doubly-robust choice that shares the explainability score helper; also "r_learner", "x_learner", "t_learner", "s_learner", "dml"). learner fills the nuisance slots (an alias, a learner, or a slot mapping). The result carries per-unit uncertainty (a posterior with a Bayesian slot, the bootstrap otherwise, the analytic standard error for the linear DR final) and the cached doubly-robust scores, so :meth:~ergodic.inference.CATEResult.explain runs the Tier-B characterization. Pass uncertainty=False to skip that step for a fast point-only result (explain still works off the cached scores).

Source code in src/ergodic/inference/_estimate.py

def cate(
    graph: MixedGraph,
    data: TabularData,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    method: str = "auto",
    learner: Any = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    uncertainty: bool = True,
    **kwargs: Any,
) -> CATEResult:
    """Identify the effect of ``treatment`` on ``outcome``, then estimate the CATE.

    Identification-gated exactly like :func:`estimate_effect`: identify the effect,
    select a CATE-capable estimator for the family, fit it, and return a
    :class:`~ergodic.inference.CATEResult`. ``method`` names the meta-learner
    (``"dr_learner"`` by default, the doubly-robust choice that shares the
    explainability score helper; also ``"r_learner"``, ``"x_learner"``,
    ``"t_learner"``, ``"s_learner"``, ``"dml"``). ``learner`` fills the nuisance
    slots (an alias, a learner, or a slot mapping). The result carries per-unit
    uncertainty (a posterior with a Bayesian slot, the bootstrap otherwise, the
    analytic standard error for the linear DR final) and the cached doubly-robust
    scores, so :meth:`~ergodic.inference.CATEResult.explain` runs the Tier-B
    characterization. Pass ``uncertainty=False`` to skip that step for a fast
    point-only result (``explain`` still works off the cached scores).
    """
    identification = identify_effect(graph, treatment, outcome)
    if not identification.identifiable:
        raise ValueError(
            f"effect of {treatment} on {outcome} is not identifiable: {identification.estimand}"
        )
    estimator_cls = cate_estimator_for(method)
    if identification.strategy not in estimator_cls.consumes:
        raise ValueError(
            f"the CATE estimator {estimator_cls.name!r} consumes "
            f"{sorted(estimator_cls.consumes)}, but the effect was identified by "
            f"{identification.strategy!r}; CATE estimation needs an adjustment-identified "
            "effect (back-door). Front-door, instruments, and the general-ID functional "
            "have no per-unit meta-learner yet"
        )
    from typing import cast

    from ._registry import HasCate

    method_kwargs = {**cate_kwargs_for(method), **kwargs}
    estimator = estimator_cls(
        learner, random_state=random_state, n_boot=n_boot, treated=treated, **method_kwargs
    )
    estimator.fit(identification, data)
    return cast("HasCate", estimator).cate(uncertainty=uncertainty)

did(data, outcome, treatment, *, controls=(), comparison='never_treated', strategy='auto', n_boot=1000, random_state=None)

Difference-in-differences: the average effect of adoption on the treated (ATT).

The panel is read through adoption timing: a unit's adoption period is the first period its treatment equals 1, units that never show a 1 are never treated, and the treatment must be absorbing (never switching from 1 back to 0 within a unit). Periods compare by their position in the sorted period sequence, so integer and datetime periods both work.

Flavors. strategy="canonical" needs a single adoption period g: each unit's outcome change from its pre-period mean (periods before g) to its post-period mean (periods at or after g), treated minus comparison. Units lacking a pre or a post observation are dropped with a warning and counted in diagnostics["n_dropped_units"]. Controls enter through a unit-level regression of that change on the treated flag and the unit's pre-period control means. strategy="group_time" handles staggered adoption in the style of Callaway and Sant'Anna: one ATT(g, t) per adoption cohort and period, each differenced against the period just before adoption, then aggregated over post periods with cohort-size weights; the cells land in by_group_time. strategy="twfe" is the static two-way fixed-effects regression (iterated within transform, so unbalanced panels work). Under staggered adoption with effects that change over time it is a weighted mix of cell effects with possibly negative weights (Goodman-Bacon), so prefer group_time there. strategy="auto" picks canonical for one adoption period and group_time for several, never twfe.

Comparison pools. comparison="never_treated" uses units that never adopt and raises when none exist. comparison="not_yet_treated" also admits units that adopt only after the periods a cell reads, the pool a panel without never-treated units needs. In the canonical flavor the two pools coincide because no later adopters exist, and twfe ignores the pool.

Inference is a unit bootstrap: resample units with replacement, recompute, and keep n_boot draws (a failed resample is skipped, with a budget of five attempts per draw). The result is kind="bootstrap"; the canonical no-controls fit also stores the two-sample analytic standard error in diagnostics["analytic_se"]. The parallel-trends and no-anticipation assumptions ride along in design.

Raises ValueError for an unknown column, a missing value in the outcome, treatment, or controls, a treatment not coded 0/1 or boolean, a non-absorbing treatment, no treated units, or an empty comparison pool.

Source code in src/ergodic/inference/estimators/_did.py

def did(
    data: PanelData,
    outcome: str,
    treatment: str,
    *,
    controls: Iterable[str] = (),
    comparison: str = "never_treated",
    strategy: str = "auto",
    n_boot: int = 1000,
    random_state: int | None = None,
) -> DiDEffect:
    """Difference-in-differences: the average effect of adoption on the treated (ATT).

    The panel is read through adoption timing: a unit's adoption period is the
    first period its treatment equals 1, units that never show a 1 are never
    treated, and the treatment must be absorbing (never switching from 1 back
    to 0 within a unit). Periods compare by their position in the sorted period
    sequence, so integer and datetime periods both work.

    Flavors. ``strategy="canonical"`` needs a single adoption period g: each
    unit's outcome change from its pre-period mean (periods before g) to its
    post-period mean (periods at or after g), treated minus comparison. Units
    lacking a pre or a post observation are dropped with a warning and counted
    in ``diagnostics["n_dropped_units"]``. Controls enter through a unit-level
    regression of that change on the treated flag and the unit's pre-period
    control means. ``strategy="group_time"`` handles staggered adoption in the
    style of Callaway and Sant'Anna: one ATT(g, t) per adoption cohort and
    period, each differenced against the period just before adoption, then
    aggregated over post periods with cohort-size weights; the cells land in
    ``by_group_time``. ``strategy="twfe"`` is the static two-way fixed-effects
    regression (iterated within transform, so unbalanced panels work). Under
    staggered adoption with effects that change over time it is a weighted mix
    of cell effects with possibly negative weights (Goodman-Bacon), so prefer
    ``group_time`` there. ``strategy="auto"`` picks canonical for one adoption
    period and group_time for several, never twfe.

    Comparison pools. ``comparison="never_treated"`` uses units that never
    adopt and raises when none exist. ``comparison="not_yet_treated"`` also
    admits units that adopt only after the periods a cell reads, the pool a
    panel without never-treated units needs. In the canonical flavor the two
    pools coincide because no later adopters exist, and twfe ignores the pool.

    Inference is a unit bootstrap: resample units with replacement, recompute,
    and keep ``n_boot`` draws (a failed resample is skipped, with a budget of
    five attempts per draw). The result is ``kind="bootstrap"``; the canonical
    no-controls fit also stores the two-sample analytic standard error in
    ``diagnostics["analytic_se"]``. The parallel-trends and no-anticipation
    assumptions ride along in ``design``.

    Raises ValueError for an unknown column, a missing value in the outcome,
    treatment, or controls, a treatment not coded 0/1 or boolean, a
    non-absorbing treatment, no treated units, or an empty comparison pool.
    """
    control_list = [str(name) for name in controls]
    _check_common(comparison, n_boot)
    if strategy not in ("auto", "canonical", "twfe", "group_time"):
        raise ValueError(
            f"strategy must be 'auto', 'canonical', 'twfe', or 'group_time', got {strategy!r}"
        )
    prep = _prepare(data, outcome, treatment, control_list)
    adoption_periods = sorted(set(prep.adoption[prep.adoption >= 0].tolist()))
    if not adoption_periods:
        raise ValueError(f"no treated units: treatment {treatment!r} is never 1")
    flavor = strategy
    if strategy == "auto":
        flavor = "canonical" if len(adoption_periods) == 1 else "group_time"
    if flavor == "canonical" and len(adoption_periods) > 1:
        raise ValueError(
            f"canonical did needs one adoption period, found {len(adoption_periods)} "
            f"({[prep.periods[g] for g in adoption_periods]}); use strategy='group_time'"
        )
    rng = np.random.default_rng(random_state)
    design = _design_record(outcome, treatment, comparison)
    if flavor == "canonical":
        return _fit_canonical(prep, comparison, n_boot, rng, design)
    if flavor == "twfe":
        return _fit_twfe(prep, comparison, n_boot, rng, design)
    return _fit_group_time(prep, comparison, n_boot, rng, design, event_window=None)

dr_scores(identification, data, *, learners=None, folds=5, clip=0.001, treated=None, random_state=None)

Cross-fit doubly-robust scores for an adjustment-identified effect.

learners is an alias, a learner, or a slot mapping over outcome and propensity (the same surface every estimator takes); the defaults are the native :class:OLS (or logistic for a binary outcome) and logistic propensity. folds sets the cross-fit count. Returns a :class:DRScores carrying the signal and the out-of-fold nuisances.

Source code in src/ergodic/inference/_scores.py

def dr_scores(
    identification: Identification,
    data: TabularData,
    *,
    learners: object = None,
    folds: int = 5,
    clip: float = 1e-3,
    treated: object = None,
    random_state: int | None = None,
) -> DRScores:
    """Cross-fit doubly-robust scores for an adjustment-identified effect.

    ``learners`` is an alias, a learner, or a slot mapping over ``outcome`` and
    ``propensity`` (the same surface every estimator takes); the defaults are the
    native :class:`OLS` (or logistic for a binary outcome) and logistic propensity.
    ``folds`` sets the cross-fit count. Returns a :class:`DRScores` carrying the
    signal and the out-of-fold nuisances.
    """
    treatment_name = _single(identification.treatment, "treatment")
    outcome_name = _single(identification.outcome, "outcome")
    covariates = identification.adjustment_set or frozenset()
    _require_complete(data, {treatment_name, outcome_name, *covariates})
    names = tuple(sorted(covariates))

    helper = _ScoreHelper(treated=treated, random_state=random_state)
    a = helper._treatment_vector(data, treatment_name)
    y, y_binary = _outcome_vector(data, outcome_name)
    z = _covariate_matrix(data, covariates)

    make_outcome, make_propensity = _resolve_score_learners(learners, y_binary)
    mu0, mu1, e = _crossfit_nuisances(
        a,
        y,
        z,
        make_outcome=make_outcome,
        make_propensity=make_propensity,
        n_folds=folds,
        clip=clip,
        random_state=random_state,
    )
    scores = aipw_scores(a, y, mu0, mu1, e)
    return DRScores(
        scores=scores,
        mu0=mu0,
        mu1=mu1,
        propensity=e,
        covariates=z,
        treatment=a,
        outcome=y,
        covariate_names=names,
    )

estimate(identification, data, *, strategy=None, learner=None, target=None, treated=None, random_state=None, n_boot=1000, **kwargs)

Estimate an identified effect over tabular data.

strategy names the estimator (else the family default). learner is an alias, a learner, or a slot-to-learner mapping (else the slot defaults). treated names the treatment level that codes to 1 for a two-level treatment (else level order decides, with a warning when the levels are not 0/1). Extra keyword arguments pass to the estimator (per_arm, scheme, clip).

Source code in src/ergodic/inference/_estimate.py

def estimate(
    identification: Identification,
    data: TabularData,
    *,
    strategy: str | None = None,
    learner: Any = None,
    target: EffectTarget | str | None = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    **kwargs: Any,
) -> EffectEstimate:
    """Estimate an identified effect over tabular data.

    ``strategy`` names the estimator (else the family default). ``learner`` is an
    alias, a learner, or a slot-to-learner mapping (else the slot defaults).
    ``treated`` names the treatment level that codes to 1 for a two-level
    treatment (else level order decides, with a warning when the levels are not
    0/1). Extra keyword arguments pass to the estimator (``per_arm``, ``scheme``,
    ``clip``).
    """
    if not identification.identifiable:
        raise ValueError(f"effect is not identifiable: {identification.estimand}")
    if identification.strategy == "id":
        if strategy is not None:
            raise ValueError(
                f"a general ID functional has no named estimator to select; got "
                f"strategy={strategy!r}. The discrete plug-in evaluates it directly"
            )
        return plugin_estimate(
            identification,
            data,
            target=target,
            treated=treated,
            random_state=random_state,
            n_boot=n_boot,
        )
    estimator_cls = estimator_for(identification, strategy)
    estimator = estimator_cls(
        learner, random_state=random_state, n_boot=n_boot, treated=treated, **kwargs
    )
    estimator.fit(identification, data)
    chosen = target if target is not None else default_target(identification)
    return estimator.effect(target=chosen)

estimate_effect(graph, data, treatment, outcome, *, strategy=None, learner=None, target=None, treated=None, random_state=None, n_boot=1000, **kwargs)

Identify the effect of treatment on outcome in graph, then estimate it.

treated names the treatment level that codes to 1 (see :func:estimate).

Source code in src/ergodic/inference/_estimate.py

def estimate_effect(
    graph: MixedGraph,
    data: TabularData,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    strategy: str | None = None,
    learner: Any = None,
    target: EffectTarget | str | None = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    **kwargs: Any,
) -> EffectEstimate:
    """Identify the effect of ``treatment`` on ``outcome`` in ``graph``, then estimate it.

    ``treated`` names the treatment level that codes to 1 (see :func:`estimate`).
    """
    identification = identify_effect(graph, treatment, outcome)
    if not identification.identifiable:
        raise ValueError(
            f"effect of {treatment} on {outcome} is not identifiable: {identification.estimand}"
        )
    return estimate(
        identification,
        data,
        strategy=strategy,
        learner=learner,
        target=target,
        treated=treated,
        random_state=random_state,
        n_boot=n_boot,
        **kwargs,
    )

estimate_mediation(graph, data, treatment, outcome, mediators, *, learner=None, treated=None, random_state=None, n_boot=1000)

Identify and estimate a mediation decomposition by g-computation.

Identifies the decomposition with :func:~ergodic.identification.identify_mediation (raising if the natural direct and indirect effects are not identifiable), then estimates NDE, NIE, the total effect, and the controlled direct effect at each mediator level by g-computation: fit mu(x, m, w) = E[Y | x, m, w] and p(m | x, w), plug them into the mediation formula, and average over the empirical covariates.

For a binary treatment and a discrete (binary or low-cardinality) mediator, with a continuous or binary outcome. learner is an alias, a learner, or a {"outcome": ..., "mediator": ...} mapping (the mediator slot needs a classifier). treated names the treatment level coded to 1. The standard errors are a nonparametric bootstrap over the rows at random_state.

A continuous mediator raises (the density / multiply-robust path is not built); a non-binary treatment raises. Returns a :class:MediationEstimate with a :class:~ergodic.inference.Design record naming the (untestable) cross-world assumption.

Source code in src/ergodic/inference/estimators/_mediation.py

def estimate_mediation(
    graph: MixedGraph,
    data: TabularData,
    treatment: str,
    outcome: str,
    mediators: str | frozenset[str],
    *,
    learner: Any = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
) -> MediationEstimate:
    """Identify and estimate a mediation decomposition by g-computation.

    Identifies the decomposition with
    :func:`~ergodic.identification.identify_mediation` (raising if the natural
    direct and indirect effects are not identifiable), then estimates NDE, NIE, the
    total effect, and the controlled direct effect at each mediator level by
    g-computation: fit ``mu(x, m, w) = E[Y | x, m, w]`` and ``p(m | x, w)``, plug
    them into the mediation formula, and average over the empirical covariates.

    For a binary treatment and a discrete (binary or low-cardinality) mediator,
    with a continuous or binary outcome. ``learner`` is an alias, a learner, or a
    ``{"outcome": ..., "mediator": ...}`` mapping (the mediator slot needs a
    classifier). ``treated`` names the treatment level coded to 1. The standard
    errors are a nonparametric bootstrap over the rows at ``random_state``.

    A continuous mediator raises (the density / multiply-robust path is not built);
    a non-binary treatment raises. Returns a :class:`MediationEstimate` with a
    :class:`~ergodic.inference.Design` record naming the (untestable) cross-world
    assumption.
    """
    mset = frozenset({mediators}) if isinstance(mediators, str) else frozenset(mediators)
    ident: MediationID = identify_mediation(graph, treatment, outcome, mset)
    if not ident.identifiable:
        raise ValueError(
            "the mediation decomposition is not identifiable: "
            f"{ident.assumptions[0]} The controlled direct effect "
            f"{'is' if ident.cde_identifiable else 'is not'} identifiable; for the "
            "natural effects, consider a path-specific analysis."
        )
    covariates = ident.covariates or frozenset()
    design = _build_design(
        data,
        treatment,
        outcome,
        mset,
        covariates,
        learner=learner,
        treated=treated,
    )
    if float(np.sum(design.a == 1.0)) == 0.0 or float(np.sum(design.a == 0.0)) == 0.0:
        raise ValueError(
            f"treatment {treatment!r} has an empty arm; both levels must appear to "
            "estimate a direct/indirect contrast"
        )

    n = data.n_rows
    full = np.arange(n, dtype=np.int64)
    point = _compute(design, full)
    boot = _bootstrap(design, n, random_state, n_boot)

    cde_levels = _mediator_level_values(data, next(iter(mset)))
    cde: dict[Any, EffectEstimate] = {}
    for code in design.levels:
        cde[cde_levels[code]] = _to_estimate(point.cde[code], boot["cde"][:, code], n, "cde")

    design_record = Design(
        strategy="mediation",
        outcome=outcome,
        treatment=treatment,
        assumptions=ident.assumptions,
        estimand=f"NDE = {ident.nde_estimand}; NIE = {ident.nie_estimand}",
    )
    return MediationEstimate(
        nde=_to_estimate(point.nde, boot["nde"], n, "nde"),
        nie=_to_estimate(point.nie, boot["nie"], n, "nie"),
        total=_to_estimate(point.total, boot["total"], n, "total"),
        cde=cde,
        design=design_record,
        mediators=mset,
        covariates=covariates,
    )

event_study(data, outcome, treatment, *, controls=(), comparison='never_treated', leads=None, lags=None, n_boot=1000, random_state=None)

The effect by event time: leads as the placebo check, lags as the dynamic path.

Runs the group-time machinery of :func:did and aggregates the ATT(g, t) cells that share an event time, the position of the period minus the position of the adoption period, weighting cohorts by size. Event time -1 is the reference: every cell differences against the period just before adoption, so the effect there is zero by construction and it is omitted. leads and lags bound the reported window (leads=2 keeps event times from -2, lags=3 up to 3); None keeps everything available.

The headline value is the overall post-adoption ATT, the same aggregate :func:did reports under strategy="group_time", and by_event_time carries one :class:EventTimeEffect per event time with a bootstrap standard error and 95% percentile interval from the same unit bootstrap that drives the headline draws. A lead whose interval sits away from zero is evidence against parallel trends.

Validation and the comparison pools match :func:did.

Source code in src/ergodic/inference/estimators/_did.py

def event_study(
    data: PanelData,
    outcome: str,
    treatment: str,
    *,
    controls: Iterable[str] = (),
    comparison: str = "never_treated",
    leads: int | None = None,
    lags: int | None = None,
    n_boot: int = 1000,
    random_state: int | None = None,
) -> DiDEffect:
    """The effect by event time: leads as the placebo check, lags as the dynamic path.

    Runs the group-time machinery of :func:`did` and aggregates the ATT(g, t)
    cells that share an event time, the position of the period minus the
    position of the adoption period, weighting cohorts by size. Event time -1
    is the reference: every cell differences against the period just before
    adoption, so the effect there is zero by construction and it is omitted.
    ``leads`` and ``lags`` bound the reported window (``leads=2`` keeps event
    times from -2, ``lags=3`` up to 3); ``None`` keeps everything available.

    The headline ``value`` is the overall post-adoption ATT, the same aggregate
    :func:`did` reports under ``strategy="group_time"``, and ``by_event_time``
    carries one :class:`EventTimeEffect` per event time with a bootstrap
    standard error and 95% percentile interval from the same unit bootstrap
    that drives the headline draws. A lead whose interval sits away from zero
    is evidence against parallel trends.

    Validation and the comparison pools match :func:`did`.
    """
    control_list = [str(name) for name in controls]
    _check_common(comparison, n_boot)
    if leads is not None and leads < 0:
        raise ValueError(f"leads must be non-negative, got {leads}")
    if lags is not None and lags < 0:
        raise ValueError(f"lags must be non-negative, got {lags}")
    prep = _prepare(data, outcome, treatment, control_list)
    if not bool((prep.adoption >= 0).any()):
        raise ValueError(f"no treated units: treatment {treatment!r} is never 1")
    rng = np.random.default_rng(random_state)
    design = _design_record(outcome, treatment, comparison)
    return _fit_group_time(prep, comparison, n_boot, rng, design, event_window=(leads, lags))

interrupted_time_series(data, outcome, *, at, seasonality=None, hac_lags=None)

Fit an interrupted time series for outcome with a break at at.

The break index is the first row whose time value is at least at; both sides need at least 3 points. The design matrix over row positions t = 0..T-1 is [1, t, post, post * (t - t0)], plus, with seasonality=m (between 2 and T // 2), Fourier pairs sin/cos(2 pi k t / m) for k = 1..min(3, m // 2). OLS fits it, and the coefficient covariance is Newey-West with Bartlett weights over hac_lags lags, default floor(0.75 * T ** (1/3)).

level_change is the post coefficient and slope_change the interaction. The headline value is the ATT, the mean post-period gap between the fit and the pre-trend projection, which equals level_change + slope_change * mean(t - t0 | post); its standard error comes from the matching linear combination of the HAC covariance. The counterfactual path zeroes the break terms but keeps the trend and seasonal terms, all at the fitted coefficients.

A NaN outcome raises, as does a break too close to either end.

Source code in src/ergodic/inference/estimators/_its.py

def interrupted_time_series(
    data: TimeSeriesData,
    outcome: str,
    *,
    at: Any,
    seasonality: int | None = None,
    hac_lags: int | None = None,
) -> ITSEffect:
    """Fit an interrupted time series for ``outcome`` with a break at ``at``.

    The break index is the first row whose time value is at least ``at``; both
    sides need at least 3 points. The design matrix over row positions
    ``t = 0..T-1`` is ``[1, t, post, post * (t - t0)]``, plus, with
    ``seasonality=m`` (between 2 and ``T // 2``), Fourier pairs
    ``sin/cos(2 pi k t / m)`` for ``k = 1..min(3, m // 2)``. OLS fits it, and
    the coefficient covariance is Newey-West with Bartlett weights over
    ``hac_lags`` lags, default ``floor(0.75 * T ** (1/3))``.

    ``level_change`` is the post coefficient and ``slope_change`` the
    interaction. The headline value is the ATT, the mean post-period gap
    between the fit and the pre-trend projection, which equals
    ``level_change + slope_change * mean(t - t0 | post)``; its standard error
    comes from the matching linear combination of the HAC covariance. The
    ``counterfactual`` path zeroes the break terms but keeps the trend and
    seasonal terms, all at the fitted coefficients.

    A NaN outcome raises, as does a break too close to either end.
    """
    if outcome not in data.variables:
        raise ValueError(
            f"outcome {outcome!r} is not a variable of the series; have {sorted(data.variables)}"
        )
    n = data.n_rows
    y = data.matrix([outcome])[:, 0]
    if np.isnan(y).any():
        raise ValueError(
            f"outcome {outcome!r} has missing values; interrupted time series "
            "needs a complete series"
        )
    if seasonality is not None and not 2 <= seasonality <= n // 2:
        raise ValueError(
            f"seasonality must be between 2 and half the series length ({n // 2}), "
            f"got {seasonality}"
        )
    if hac_lags is not None and hac_lags < 0:
        raise ValueError(f"hac_lags must be nonnegative, got {hac_lags}")

    times_list = list(data.times.tolist())
    try:
        post_flags = [bool(value >= at) for value in times_list]
    except TypeError as exc:
        raise ValueError(f"at={at!r} is not comparable to the time column {data.time!r}") from exc
    t0 = post_flags.index(True) if any(post_flags) else n
    n_pre = t0
    n_post = n - t0
    if n_pre < 3 or n_post < 3:
        raise ValueError(
            "the break needs at least 3 points on each side, "
            f"got {n_pre} pre and {n_post} post for at={at!r}"
        )

    t = np.arange(n, dtype=np.float64)
    post = (t >= t0).astype(np.float64)
    cols = [np.ones(n, dtype=np.float64), t, post, post * (t - float(t0))]
    if seasonality is not None:
        for k in range(1, min(3, seasonality // 2) + 1):
            angle = 2.0 * np.pi * k * t / float(seasonality)
            cols.append(np.sin(angle))
            cols.append(np.cos(angle))
    x = np.asarray(np.column_stack(cols), dtype=np.float64)

    beta = np.asarray(np.linalg.lstsq(x, y, rcond=None)[0], dtype=np.float64)
    resid = np.asarray(y - x @ beta, dtype=np.float64)
    lags = hac_lags if hac_lags is not None else int(np.floor(0.75 * float(n) ** (1.0 / 3.0)))
    cov = _hac_cov(x, resid, lags)

    level_change = float(beta[2])
    slope_change = float(beta[3])
    level_change_se = float(np.sqrt(max(float(cov[2, 2]), 0.0)))
    slope_change_se = float(np.sqrt(max(float(cov[3, 3]), 0.0)))
    contrast = np.zeros(x.shape[1], dtype=np.float64)
    contrast[2] = 1.0
    contrast[3] = float(np.mean(t[t0:] - float(t0)))
    value = float(contrast @ beta)
    std_error = float(np.sqrt(max(float(contrast @ cov @ contrast), 0.0)))

    fitted = np.asarray(x @ beta, dtype=np.float64)
    x_counter = x.copy()
    x_counter[:, 2] = 0.0
    x_counter[:, 3] = 0.0
    counterfactual = np.asarray(x_counter @ beta, dtype=np.float64)
    ss_res = float(np.sum(resid**2))
    ss_tot = float(np.sum((y - float(np.mean(y))) ** 2))
    r2 = 1.0 - ss_res / ss_tot if ss_tot > 0.0 else float("nan")

    design = Design(
        strategy="interrupted_time_series",
        outcome=outcome,
        treatment=None,
        assumptions=(
            "the pre-period trend would have continued without the break",
            "no co-timed shock at the break",
        ),
        estimand="average post-period gap to the projected pre-trend",
    )
    return ITSEffect(
        target=EffectTarget.ATT,
        estimator="interrupted_time_series",
        learners={},
        n=n,
        kind="analytic",
        diagnostics={
            "n_pre": float(n_pre),
            "n_post": float(n_post),
            "hac_lags": float(lags),
            "r2": r2,
        },
        _point=value,
        _std_error=std_error,
        design=design,
        at=at,
        level_change=level_change,
        level_change_se=level_change_se,
        slope_change=slope_change,
        slope_change_se=slope_change_se,
        times=tuple(times_list),
        observed=tuple(float(v) for v in y),
        counterfactual=tuple(float(v) for v in counterfactual),
        fitted=tuple(float(v) for v in fitted),
    )

policy_tree(result_or_scores, covariates=None, *, depth=2)

Fit a shallow doubly-robust optimal-policy tree of depth depth.

result_or_scores is a :class:~ergodic.inference.CATEResult (its cached cross-fit scores are read) or a :class:~ergodic.inference._scores.DRScores directly. covariates names the split covariates (a name, a list of names, or None for every covariate carried by the scores). The tree maximizes the doubly-robust welfare W(pi) = mean[(2 pi(X) - 1) Gamma_i]; at depth <= 2 the search is exhaustive over thresholds and so globally optimal, and a deeper tree uses a greedy recursive best-split (documented, not guaranteed optimal).

Returns a frozen :class:PolicyTree carrying the rule, the estimated policy value with a standard error from the per-unit influence terms, and the regret against treating everyone / no one.

Source code in src/ergodic/inference/explain/_policy.py

def policy_tree(
    result_or_scores: CATEResult | DRScores,
    covariates: str | tuple[str, ...] | list[str] | None = None,
    *,
    depth: int = 2,
) -> PolicyTree:
    """Fit a shallow doubly-robust optimal-policy tree of depth ``depth``.

    ``result_or_scores`` is a :class:`~ergodic.inference.CATEResult` (its cached
    cross-fit scores are read) or a :class:`~ergodic.inference._scores.DRScores`
    directly. ``covariates`` names the split covariates (a name, a list of names, or
    ``None`` for every covariate carried by the scores). The tree maximizes the
    doubly-robust welfare ``W(pi) = mean[(2 pi(X) - 1) Gamma_i]``; at ``depth <= 2``
    the search is exhaustive over thresholds and so globally optimal, and a deeper
    tree uses a greedy recursive best-split (documented, not guaranteed optimal).

    Returns a frozen :class:`PolicyTree` carrying the rule, the estimated policy
    value with a standard error from the per-unit influence terms, and the regret
    against treating everyone / no one.
    """
    if depth < 1:
        raise ValueError(f"depth must be at least 1, got {depth}")
    scores, full_design, all_names, treatment, outcome = _extract(result_or_scores)
    n = scores.shape[0]
    if n < 2:
        raise ValueError(f"policy_tree needs at least 2 units, got {n}")
    design, names = _select_columns(full_design, all_names, covariates)

    exhaustive = depth <= _EXHAUSTIVE_DEPTH
    if exhaustive:
        _, root = _best_split(scores, design, names, depth)
    else:
        _, root = _greedy_split(scores, design, names, depth)

    actions = np.asarray([root.predict(row) for row in design], dtype=np.float64)
    influence = (2.0 * actions - 1.0) * scores
    value = float(np.mean(influence))
    std_error = float(np.std(influence, ddof=1) / np.sqrt(n)) if n > 1 else float("nan")
    treat_all = float(np.mean(scores))
    treat_none = float(np.mean(-scores))
    return PolicyTree(
        root=root,
        covariates=names,
        n=n,
        depth=depth,
        value=value,
        std_error=std_error,
        treat_all_value=treat_all,
        treat_none_value=treat_none,
        exhaustive=exhaustive,
        treatment=treatment,
        outcome=outcome,
        _scores=scores,
        _design=design,
    )

refute_effect(graph, data, treatment, outcome, *, method='placebo_treatment', strategy=None, learner=None, treated=None, estimand_value=None, n_simulations=100, subset_fraction=0.8, random_state=None, **kwargs)

Stress-test the effect of treatment on outcome under a perturbation.

Runs the same identify-then-estimate path as :func:estimate_effect, so pass the same strategy, learner, and treated you used for the estimate. estimand_value skips re-computing the effect on the real data when you already have it.

The three methods:

• "placebo_treatment" shuffles the treatment column, breaking its link to everything, and re-estimates. The placebo effect should sit near zero. passed when the mean placebo effect is under 20 percent of the real one; p_value is the share of placebo runs whose effect matches or beats the real one in size.
• "random_common_cause" adds an independent random covariate as a parent of both treatment and outcome and re-estimates. The effect should not move. passed when the mean stays within 10 percent of the original.
• "data_subset" re-estimates on a random subset_fraction of the rows. The effect should hold. passed when the original sits within two standard deviations of the subset spread.

Flexible or cross-fit estimators make each run a full fit, so lower n_simulations (or pick a faster strategy) when the estimator is heavy.

Source code in src/ergodic/inference/_refute.py

def refute_effect(
    graph: MixedGraph,
    data: TabularData,
    treatment: str,
    outcome: str,
    *,
    method: str = "placebo_treatment",
    strategy: str | None = None,
    learner: Any = None,
    treated: Any | None = None,
    estimand_value: float | None = None,
    n_simulations: int = 100,
    subset_fraction: float = 0.8,
    random_state: int | None = None,
    **kwargs: Any,
) -> Refutation:
    """Stress-test the effect of ``treatment`` on ``outcome`` under a perturbation.

    Runs the same identify-then-estimate path as :func:`estimate_effect`, so pass
    the same ``strategy``, ``learner``, and ``treated`` you used for the estimate.
    ``estimand_value`` skips re-computing the effect on the real data when you
    already have it.

    The three methods:

    - ``"placebo_treatment"`` shuffles the treatment column, breaking its link to
      everything, and re-estimates. The placebo effect should sit near zero.
      ``passed`` when the mean placebo effect is under 20 percent of the real one;
      ``p_value`` is the share of placebo runs whose effect matches or beats the
      real one in size.
    - ``"random_common_cause"`` adds an independent random covariate as a parent of
      both treatment and outcome and re-estimates. The effect should not move.
      ``passed`` when the mean stays within 10 percent of the original.
    - ``"data_subset"`` re-estimates on a random ``subset_fraction`` of the rows.
      The effect should hold. ``passed`` when the original sits within two standard
      deviations of the subset spread.

    Flexible or cross-fit estimators make each run a full fit, so lower
    ``n_simulations`` (or pick a faster ``strategy``) when the estimator is heavy.
    """
    if method not in _METHODS:
        raise ValueError(f"unknown refutation method {method!r}; use one of {', '.join(_METHODS)}")
    if n_simulations < 2:
        raise ValueError(f"n_simulations must be at least 2, got {n_simulations}")
    rng = np.random.default_rng(random_state)

    def run(g: MixedGraph, d: TabularData) -> float:
        return estimate_effect(
            g,
            d,
            treatment,
            outcome,
            strategy=strategy,
            learner=learner,
            treated=treated,
            random_state=random_state,
            **kwargs,
        ).value

    if estimand_value is None:
        estimand_value = run(graph, data)

    if method == "placebo_treatment":
        base = data.to_pandas()
        column = base[treatment].to_numpy()
        effects = np.array(
            [
                run(graph, _rebuilt(data, base.assign(**{treatment: rng.permutation(column)})))
                for _ in range(n_simulations)
            ]
        )
        refuted_value = float(effects.mean())
        refuted_std = float(effects.std(ddof=1))
        p_value: float | None = float(
            (1 + np.sum(np.abs(effects) >= abs(estimand_value))) / (1 + n_simulations)
        )
        passed = abs(refuted_value) <= 0.2 * abs(estimand_value)

    elif method == "random_common_cause":
        noise = "_random_common_cause"
        if noise in data.names:
            raise ValueError(f"the data already has a column named {noise!r}")
        augmented = (
            graph.add_edge(f"{noise} -> {treatment}")
            .add_edge(f"{noise} -> {outcome}")
            .as_kind(graph.kind)
        )
        base = data.to_pandas()
        effects = np.array(
            [
                run(
                    augmented,
                    _rebuilt(data, base.assign(**{noise: rng.standard_normal(len(base))})),
                )
                for _ in range(n_simulations)
            ]
        )
        refuted_value = float(effects.mean())
        refuted_std = float(effects.std(ddof=1))
        p_value = None
        passed = abs(refuted_value - estimand_value) <= 0.1 * abs(estimand_value) + 1e-9

    else:  # data_subset
        n = data.n_rows
        size = max(2, round(subset_fraction * n))
        effects = np.array(
            [
                run(graph, data.take(np.sort(rng.choice(n, size=size, replace=False))))
                for _ in range(n_simulations)
            ]
        )
        refuted_value = float(effects.mean())
        refuted_std = float(effects.std(ddof=1))
        p_value = None
        passed = abs(estimand_value - refuted_value) <= 2.0 * refuted_std + 1e-9

    return Refutation(
        method=method,
        estimand_value=float(estimand_value),
        refuted_value=refuted_value,
        refuted_std=refuted_std,
        p_value=p_value,
        passed=passed,
        n_simulations=n_simulations,
        diagnostics={"effect_min": float(effects.min()), "effect_max": float(effects.max())},
    )

select_cate(graph, data, treatment, outcome, *, candidates=None, learner=None, treated=None, random_state=None, n_boot=1000, folds=5, nu=0.5, q_iters=200)

Fit and rank candidate CATE learners by the doubly-robust tau-risk.

Identifies the effect (the same back-door gate as :func:~ergodic.inference.cate), fits each candidate in candidates (default the Phase-1 registered meta-learners dr_learner/r_learner/x_learner/t_learner) through the cate facade, and scores each against one shared cross-fit AIPW signal Gamma by the DR tau-risk mean[(Gamma - tau_hat)^2]. Returns a :class:SelectionReport with the ranked candidates (best risk first), the selected fitted :class:~ergodic.inference.CATEResult, and a Q-aggregation convex-ensemble result.

learner fills every candidate's nuisance slots; nu is the Q-aggregation mix (the penalty pulling weight toward lower-risk models). The shared Gamma is computed once with :func:~ergodic.inference.dr_scores, so the comparison is fair across candidates and each candidate's own out-of-fold tau_hat keeps the ranking honest.

Source code in src/ergodic/inference/_select.py

def select_cate(
    graph: MixedGraph,
    data: TabularData,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    candidates: Sequence[str] | None = None,
    learner: Any = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    folds: int = 5,
    nu: float = 0.5,
    q_iters: int = 200,
) -> SelectionReport:
    """Fit and rank candidate CATE learners by the doubly-robust tau-risk.

    Identifies the effect (the same back-door gate as :func:`~ergodic.inference.cate`),
    fits each candidate in ``candidates`` (default the Phase-1 registered meta-learners
    ``dr_learner``/``r_learner``/``x_learner``/``t_learner``) through the ``cate``
    facade, and scores each against one shared cross-fit AIPW signal ``Gamma`` by the
    DR tau-risk ``mean[(Gamma - tau_hat)^2]``. Returns a :class:`SelectionReport` with
    the ranked candidates (best risk first), the selected fitted
    :class:`~ergodic.inference.CATEResult`, and a Q-aggregation convex-ensemble result.

    ``learner`` fills every candidate's nuisance slots; ``nu`` is the Q-aggregation
    mix (the penalty pulling weight toward lower-risk models). The shared ``Gamma`` is
    computed once with :func:`~ergodic.inference.dr_scores`, so the comparison is fair
    across candidates and each candidate's own out-of-fold ``tau_hat`` keeps the
    ranking honest.
    """
    from ..identification import identify_effect
    from ._estimate import cate
    from ._scores import dr_scores

    names = tuple(candidates) if candidates is not None else _DEFAULT_CANDIDATES
    if not names:
        raise ValueError("candidates is empty; pass at least one CATE method name")

    identification = identify_effect(graph, treatment, outcome)
    if not identification.identifiable:
        raise ValueError(
            f"effect of {treatment} on {outcome} is not identifiable: {identification.estimand}"
        )
    if identification.strategy != "adjustment":
        raise ValueError(
            "select_cate needs an adjustment-identified (back-door) effect; the effect "
            f"was identified by {identification.strategy!r}, which has no CATE meta-learner"
        )

    # One shared orthogonal signal: the comparison currency and the aggregation target.
    shared = dr_scores(
        identification,
        data,
        learners=learner,
        folds=folds,
        treated=treated,
        random_state=random_state,
    )
    gamma = np.asarray(shared.scores, dtype=np.float64)

    fitted: dict[str, CATEResult] = {}
    tau_hats: dict[str, npt.NDArray[np.float64]] = {}
    scored: list[CandidateScore] = []
    for method in names:
        result = cate(
            graph,
            data,
            treatment,
            outcome,
            method=method,
            learner=learner,
            treated=treated,
            random_state=random_state,
            n_boot=n_boot,
        )
        # the candidate's out-of-fold proxy ranking when cached (the honest choice),
        # else its in-sample fit. Scored against the shared Gamma.
        cached = result._scores
        if cached is not None and cached.tau_hat is not None:
            tau = np.asarray(cached.tau_hat, dtype=np.float64)
        else:
            tau = np.asarray(result.values, dtype=np.float64)
        fitted[method] = result
        tau_hats[method] = tau
        scored.append(CandidateScore(method=method, tau_risk=_tau_risk(gamma, tau)))

    ranking = tuple(sorted(scored, key=lambda s: (s.tau_risk, s.method)))
    selected_method = ranking[0].method
    selected = fitted[selected_method]

    # Q-aggregation: convex DR-loss ensemble with the Q penalty over the candidates.
    agg_methods = tuple(names)
    design = np.column_stack([tau_hats[m] for m in agg_methods])
    per_model_risk = np.asarray([_tau_risk(gamma, tau_hats[m]) for m in agg_methods])
    weights = _frank_wolfe_simplex(design, gamma, per_model_risk, nu=nu, iters=q_iters)
    agg_tau = np.asarray(design @ weights, dtype=np.float64)
    aggregate = _aggregate_result(selected, agg_tau, shared)
    return SelectionReport(
        ranking=ranking,
        selected_method=selected_method,
        selected=selected,
        aggregate=aggregate,
        aggregate_methods=agg_methods,
        weights=weights,
        aggregate_tau_risk=_tau_risk(gamma, agg_tau),
    )

synthetic_control(data, outcome, *, treatment=None, treated=None, at=None, donors=None)

Fit a synthetic control for one treated unit in a panel.

Name the treated unit in exactly one of two ways. Pass treatment, an absorbing 0/1 column (once a unit turns 1 it stays 1) that marks exactly one unit as treated; its first treated period becomes at. Or pass treated and at directly. donors defaults to every other unit.

The donor weights minimize the squared pre-period tracking error over the simplex (nonnegative, summing to one) by projected gradient descent with the Duchi projection; the run is deterministic. The headline value is the ATT, the mean of observed - synthetic over the post periods (t >= at), and pre_rmspe reports the fit it rests on.

Inference is by placebo permutation, with the treated unit excluded from every placebo run: each donor is treated as pseudo-treated, fit on the remaining donors, and scored by its post-to-pre RMSPE ratio (skipped when its pre RMSPE falls below 1e-12). p_value is (1 + count(ratio_j >= ratio_treated)) / (1 + n_placebos). There is no sampling standard error for one treated unit, so std_error stays None and interval reads the placebo gaps instead.

Validation: at must be one of the panel's periods with at least 2 pre periods and 1 post period; the treated unit must be observed at every period; donors missing any period are dropped with a warning and at least 2 complete donors must remain; a NaN outcome among the used rows raises.

Source code in src/ergodic/inference/estimators/_synthetic_control.py

def synthetic_control(
    data: PanelData,
    outcome: str,
    *,
    treatment: str | None = None,
    treated: Any | None = None,
    at: Any | None = None,
    donors: Iterable[Any] | None = None,
) -> SyntheticControlEffect:
    """Fit a synthetic control for one treated unit in a panel.

    Name the treated unit in exactly one of two ways. Pass ``treatment``, an
    absorbing 0/1 column (once a unit turns 1 it stays 1) that marks exactly
    one unit as treated; its first treated period becomes ``at``. Or pass
    ``treated`` and ``at`` directly. ``donors`` defaults to every other unit.

    The donor weights minimize the squared pre-period tracking error over the
    simplex (nonnegative, summing to one) by projected gradient descent with
    the Duchi projection; the run is deterministic. The headline value is the
    ATT, the mean of ``observed - synthetic`` over the post periods
    (``t >= at``), and ``pre_rmspe`` reports the fit it rests on.

    Inference is by placebo permutation, with the treated unit excluded from
    every placebo run: each donor is treated as pseudo-treated, fit on the
    remaining donors, and scored by its post-to-pre RMSPE ratio (skipped when
    its pre RMSPE falls below 1e-12). ``p_value`` is
    ``(1 + count(ratio_j >= ratio_treated)) / (1 + n_placebos)``. There is no
    sampling standard error for one treated unit, so ``std_error`` stays
    ``None`` and ``interval`` reads the placebo gaps instead.

    Validation: ``at`` must be one of the panel's periods with at least 2 pre
    periods and 1 post period; the treated unit must be observed at every
    period; donors missing any period are dropped with a warning and at least
    2 complete donors must remain; a NaN outcome among the used rows raises.
    """
    if outcome not in data.variables:
        raise ValueError(
            f"outcome {outcome!r} is not a variable of the panel; have {sorted(data.variables)}"
        )
    pair_given = treated is not None or at is not None
    if treatment is not None and pair_given:
        raise ValueError("pass exactly one of treatment= or the (treated=, at=) pair, got both")
    if treatment is None and (treated is None or at is None):
        raise ValueError(
            "pass exactly one of treatment= or the (treated=, at=) pair; "
            f"got treated={treated!r} and at={at!r}"
        )
    if treatment is not None:
        treated, at = _treated_from_column(data, treatment)

    entities = data.entities
    periods = data.periods
    if treated not in entities:
        raise ValueError(f"treated unit {treated!r} is not an entity of the panel")
    if at not in periods:
        raise ValueError(f"at={at!r} is not one of the panel's periods")
    at_idx = periods.index(at)
    n_pre = at_idx
    n_post = len(periods) - at_idx
    if n_pre < 2:
        raise ValueError(
            f"need at least 2 pre periods before at={at!r}, got {n_pre}; "
            "the weights have nothing to track"
        )

    if donors is None:
        donor_list = [unit for unit in entities if unit != treated]
    else:
        donor_list = sorted(set(donors))
        unknown = [donor for donor in donor_list if donor not in entities]
        if unknown:
            raise ValueError(f"unknown donor units: {unknown}")
        if treated in donor_list:
            raise ValueError(f"the treated unit {treated!r} cannot be one of its own donors")

    units_raw = data.column(data.entity).tolist()
    times_raw = data.column(data.time).tolist()
    outcome_values = data.matrix([outcome])[:, 0]
    paths: dict[Any, dict[Any, float]] = {}
    for unit, period, value in zip(units_raw, times_raw, outcome_values.tolist(), strict=True):
        paths.setdefault(unit, {})[period] = float(value)

    treated_path = paths[treated]
    n_missing = sum(period not in treated_path for period in periods)
    if n_missing:
        raise ValueError(
            f"the treated unit {treated!r} is missing {n_missing} of {len(periods)} periods; "
            "synthetic control needs its full path"
        )
    kept: list[Any] = []
    dropped: list[Any] = []
    for donor in donor_list:
        if all(period in paths[donor] for period in periods):
            kept.append(donor)
        else:
            dropped.append(donor)
    if dropped:
        warnings.warn(
            f"dropped {len(dropped)} donor(s) not observed at every period: {dropped[:5]}",
            UserWarning,
            stacklevel=2,
        )
    if len(kept) < 2:
        raise ValueError(f"need at least 2 donors observed at every period, got {len(kept)}")

    y1 = np.asarray([treated_path[period] for period in periods], dtype=np.float64)
    y0 = np.asarray(
        [[paths[donor][period] for donor in kept] for period in periods], dtype=np.float64
    )
    if np.isnan(y1).any() or np.isnan(y0).any():
        raise ValueError(
            f"outcome {outcome!r} has missing values among the used rows; "
            "synthetic control needs complete paths"
        )

    w = _simplex_weights(y1[:at_idx], y0[:at_idx])
    synthetic_path = np.asarray(y0 @ w, dtype=np.float64)
    gaps = y1 - synthetic_path
    att = float(np.mean(gaps[at_idx:]))
    pre_rmspe = float(np.sqrt(np.mean(gaps[:at_idx] ** 2)))
    post_rmspe = float(np.sqrt(np.mean(gaps[at_idx:] ** 2)))

    ratios: list[float] = []
    n_skipped = 0
    placebo_gaps: dict[Any, tuple[float, ...]] = {}
    for j, donor in enumerate(kept):
        pool = [i for i in range(len(kept)) if i != j]
        w_j = _simplex_weights(y0[:at_idx, j], y0[:at_idx][:, pool])
        gaps_j = np.asarray(y0[:, j] - y0[:, pool] @ w_j, dtype=np.float64)
        placebo_gaps[donor] = tuple(float(g) for g in gaps_j[at_idx:])
        pre_j = float(np.sqrt(np.mean(gaps_j[:at_idx] ** 2)))
        post_j = float(np.sqrt(np.mean(gaps_j[at_idx:] ** 2)))
        if pre_j < 1e-12:
            n_skipped += 1
            continue
        ratios.append(post_j / pre_j)
    treated_ratio = float("inf") if pre_rmspe < 1e-12 else post_rmspe / pre_rmspe
    p_value: float | None = None
    if ratios:
        p_value = (1 + sum(ratio >= treated_ratio for ratio in ratios)) / (1 + len(ratios))

    design = Design(
        strategy="synthetic_control",
        outcome=outcome,
        treatment=treatment,
        assumptions=(
            "the treated unit's untreated path lies in the convex hull of the donors",
            "no spillovers from the treated unit to donors",
            "a good pre-period fit, judge pre_rmspe",
        ),
        estimand="ATT for the treated unit, post-period gap to its synthetic control",
    )
    return SyntheticControlEffect(
        target=EffectTarget.ATT,
        estimator="synthetic_control",
        learners={},
        n=(1 + len(kept)) * len(periods),
        kind="analytic",
        diagnostics={
            "n_donors": float(len(kept)),
            "n_pre": float(n_pre),
            "n_post": float(n_post),
            "n_placebos": float(len(ratios)),
            "n_dropped_donors": float(len(dropped)),
            "n_placebo_skipped": float(n_skipped),
            "pre_rmspe": pre_rmspe,
            "post_rmspe": post_rmspe,
        },
        _point=att,
        _std_error=None,
        design=design,
        treated_unit=treated,
        at=at,
        weights={donor: float(weight) for donor, weight in zip(kept, w, strict=True)},
        periods=periods,
        observed=tuple(float(v) for v in y1),
        synthetic=tuple(float(v) for v in synthetic_path),
        pre_rmspe=pre_rmspe,
        post_rmspe=post_rmspe,
        p_value=p_value,
        placebo_gaps=placebo_gaps,
    )

Per-unit effects (CATE/ITE) and explainability

The cate facade and the uniform per-unit result. cate(graph, data, X, Y) is identification-gated like estimate_effect and returns a CATEResult; SCM.ite returns an ITEResult. Both share the PerUnitEffect surface (.values, .interval, .ate, __array__). CATEResult.explain returns a Heterogeneity record for Tier-B characterization (BLP, GATES, CLAN, the calibration test, RATE).

ergodic.cate(graph, data, treatment, outcome, *, method='auto', learner=None, treated=None, random_state=None, n_boot=1000, uncertainty=True, **kwargs)

Identify the effect of treatment on outcome, then estimate the CATE.

Identification-gated exactly like :func:estimate_effect: identify the effect, select a CATE-capable estimator for the family, fit it, and return a :class:~ergodic.inference.CATEResult. method names the meta-learner ("dr_learner" by default, the doubly-robust choice that shares the explainability score helper; also "r_learner", "x_learner", "t_learner", "s_learner", "dml"). learner fills the nuisance slots (an alias, a learner, or a slot mapping). The result carries per-unit uncertainty (a posterior with a Bayesian slot, the bootstrap otherwise, the analytic standard error for the linear DR final) and the cached doubly-robust scores, so :meth:~ergodic.inference.CATEResult.explain runs the Tier-B characterization. Pass uncertainty=False to skip that step for a fast point-only result (explain still works off the cached scores).

Source code in src/ergodic/inference/_estimate.py

def cate(
    graph: MixedGraph,
    data: TabularData,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    method: str = "auto",
    learner: Any = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    uncertainty: bool = True,
    **kwargs: Any,
) -> CATEResult:
    """Identify the effect of ``treatment`` on ``outcome``, then estimate the CATE.

    Identification-gated exactly like :func:`estimate_effect`: identify the effect,
    select a CATE-capable estimator for the family, fit it, and return a
    :class:`~ergodic.inference.CATEResult`. ``method`` names the meta-learner
    (``"dr_learner"`` by default, the doubly-robust choice that shares the
    explainability score helper; also ``"r_learner"``, ``"x_learner"``,
    ``"t_learner"``, ``"s_learner"``, ``"dml"``). ``learner`` fills the nuisance
    slots (an alias, a learner, or a slot mapping). The result carries per-unit
    uncertainty (a posterior with a Bayesian slot, the bootstrap otherwise, the
    analytic standard error for the linear DR final) and the cached doubly-robust
    scores, so :meth:`~ergodic.inference.CATEResult.explain` runs the Tier-B
    characterization. Pass ``uncertainty=False`` to skip that step for a fast
    point-only result (``explain`` still works off the cached scores).
    """
    identification = identify_effect(graph, treatment, outcome)
    if not identification.identifiable:
        raise ValueError(
            f"effect of {treatment} on {outcome} is not identifiable: {identification.estimand}"
        )
    estimator_cls = cate_estimator_for(method)
    if identification.strategy not in estimator_cls.consumes:
        raise ValueError(
            f"the CATE estimator {estimator_cls.name!r} consumes "
            f"{sorted(estimator_cls.consumes)}, but the effect was identified by "
            f"{identification.strategy!r}; CATE estimation needs an adjustment-identified "
            "effect (back-door). Front-door, instruments, and the general-ID functional "
            "have no per-unit meta-learner yet"
        )
    from typing import cast

    from ._registry import HasCate

    method_kwargs = {**cate_kwargs_for(method), **kwargs}
    estimator = estimator_cls(
        learner, random_state=random_state, n_boot=n_boot, treated=treated, **method_kwargs
    )
    estimator.fit(identification, data)
    return cast("HasCate", estimator).cate(uncertainty=uncertainty)

ergodic.CATEResult dataclass

Bases: PerUnitEffect

Conditional average treatment effects, with an explainability hook.

A conditional average over the unobserved noise (the inference-pillar object). Carries the :class:~ergodic.identification.Identification it was fit under and the cached doubly-robust scores + covariate design, so :meth:explain projects honest summaries (BLP, GATES, CLAN, the calibration test, RATE) off the same orthogonal signal.

Source code in src/ergodic/inference/_perunit.py

@dataclass(frozen=True, eq=False)
class CATEResult(PerUnitEffect):
    """Conditional average treatment effects, with an explainability hook.

    A conditional average over the unobserved noise (the inference-pillar object).
    Carries the :class:`~ergodic.identification.Identification` it was fit under and
    the cached doubly-robust scores + covariate design, so :meth:`explain` projects
    honest summaries (BLP, GATES, CLAN, the calibration test, RATE) off the same
    orthogonal signal.
    """

    estimand: str = "cate"
    identification: Identification | None = None
    _scores: DRScores | None = field(default=None, repr=False)

    def explain(self) -> Heterogeneity:
        """Tier-B characterization of the heterogeneity, on the doubly-robust scores.

        Builds a :class:`~ergodic.inference.explain.Heterogeneity` record whose
        methods (``blp``, ``gates``, ``clan``, ``calibration_test``, ``rate``) all
        consume the cached cross-fit scores, the covariate design, and the fitted
        per-unit ``tau_hat``. Needs the scores; raises when they were not cached
        (an estimator that does not build them, or a result without them).
        """
        from .explain import Heterogeneity

        if self._scores is None:
            raise RuntimeError(
                "this CATEResult carries no doubly-robust scores to explain; "
                "explainability needs the cross-fit AIPW signal (the cate() facade "
                "and the doubly-robust meta-learners cache it)"
            )
        # the out-of-fold proxy ranking when available (the honesty recipe: the
        # ranking is fixed on data not used for the summary), else the in-sample fit
        tau_hat = self._scores.tau_hat if self._scores.tau_hat is not None else self.values
        return Heterogeneity(
            scores=self._scores,
            tau_hat=np.asarray(tau_hat, dtype=np.float64),
            treatment=self.treatment,
            outcome=self.outcome,
        )

explain()

Tier-B characterization of the heterogeneity, on the doubly-robust scores.

Builds a :class:~ergodic.inference.explain.Heterogeneity record whose methods (blp, gates, clan, calibration_test, rate) all consume the cached cross-fit scores, the covariate design, and the fitted per-unit tau_hat. Needs the scores; raises when they were not cached (an estimator that does not build them, or a result without them).

Source code in src/ergodic/inference/_perunit.py

def explain(self) -> Heterogeneity:
    """Tier-B characterization of the heterogeneity, on the doubly-robust scores.

    Builds a :class:`~ergodic.inference.explain.Heterogeneity` record whose
    methods (``blp``, ``gates``, ``clan``, ``calibration_test``, ``rate``) all
    consume the cached cross-fit scores, the covariate design, and the fitted
    per-unit ``tau_hat``. Needs the scores; raises when they were not cached
    (an estimator that does not build them, or a result without them).
    """
    from .explain import Heterogeneity

    if self._scores is None:
        raise RuntimeError(
            "this CATEResult carries no doubly-robust scores to explain; "
            "explainability needs the cross-fit AIPW signal (the cate() facade "
            "and the doubly-robust meta-learners cache it)"
        )
    # the out-of-fold proxy ranking when available (the honesty recipe: the
    # ranking is fixed on data not used for the summary), else the in-sample fit
    tau_hat = self._scores.tau_hat if self._scores.tau_hat is not None else self.values
    return Heterogeneity(
        scores=self._scores,
        tau_hat=np.asarray(tau_hat, dtype=np.float64),
        treatment=self.treatment,
        outcome=self.outcome,
    )

ergodic.ITEResult dataclass

Bases: PerUnitEffect

Individual treatment effects from a structural model (the twin network).

The individual's own counterfactual contrast (the SCM-pillar object), not a conditional average. Carries the factual outcome column, the contrast framing (levels= or shift=), and a structural provenance note.

:meth:explain is deferred: the Tier-B machinery is defined over the CATE's doubly-robust scores, and the inference-valid descriptive regression of abducted ITEs is a later phase, so it raises with a pointer rather than faking inference.

Source code in src/ergodic/inference/_perunit.py

@dataclass(frozen=True, eq=False)
class ITEResult(PerUnitEffect):
    """Individual treatment effects from a structural model (the twin network).

    The individual's own counterfactual contrast (the SCM-pillar object), not a
    conditional average. Carries the factual outcome column, the contrast framing
    (``levels=`` or ``shift=``), and a structural provenance note.

    :meth:`explain` is deferred: the Tier-B machinery is defined over the CATE's
    doubly-robust scores, and the inference-valid descriptive regression of abducted
    ITEs is a later phase, so it raises with a pointer rather than faking inference.
    """

    estimand: str = "ite"
    framing: str = ""
    note: str = ""
    _factual: npt.NDArray[np.float64] | None = field(default=None, repr=False)

    @property
    def factual(self) -> npt.NDArray[np.float64] | None:
        """The per-unit factual outcome the contrast was taken against, when carried."""
        if self._factual is None:
            return None
        return np.asarray(self._factual, dtype=np.float64)

    def explain(self) -> Heterogeneity:
        """Deferred in phase 1: the Tier-B tools need the CATE's doubly-robust scores."""
        raise NotImplementedError(
            "ITEResult.explain is not available yet: the Tier-B characterization "
            "(BLP, GATES, RATE) is defined over the doubly-robust CATE scores, and the "
            "descriptive regression of structural ITEs is a later phase. Use a CATE "
            "meta-learner (cate(...)) for inference-valid heterogeneity characterization"
        )

factual property

The per-unit factual outcome the contrast was taken against, when carried.

explain()

Deferred in phase 1: the Tier-B tools need the CATE's doubly-robust scores.

Source code in src/ergodic/inference/_perunit.py

def explain(self) -> Heterogeneity:
    """Deferred in phase 1: the Tier-B tools need the CATE's doubly-robust scores."""
    raise NotImplementedError(
        "ITEResult.explain is not available yet: the Tier-B characterization "
        "(BLP, GATES, RATE) is defined over the doubly-robust CATE scores, and the "
        "descriptive regression of structural ITEs is a later phase. Use a CATE "
        "meta-learner (cate(...)) for inference-valid heterogeneity characterization"
    )

ergodic.PerUnitEffect dataclass

Per-unit treatment effects with their uncertainty.

_values is the per-unit point vector (n,), ordered to the input rows. Either _std_error (per-unit analytic SEs) or _draws (S, n) (bootstrap or posterior) carries the uncertainty, and the accessors read whichever is present. index are explicit row ids; estimand is set by the subclass.

Source code in src/ergodic/inference/_perunit.py

@dataclass(frozen=True, eq=False)
class PerUnitEffect:
    """Per-unit treatment effects with their uncertainty.

    ``_values`` is the per-unit point vector ``(n,)``, ordered to the input rows.
    Either ``_std_error`` (per-unit analytic SEs) or ``_draws`` ``(S, n)`` (bootstrap
    or posterior) carries the uncertainty, and the accessors read whichever is
    present. ``index`` are explicit row ids; ``estimand`` is set by the subclass.
    """

    treatment: str
    outcome: str
    covariates: tuple[str, ...]
    n: int
    kind: str  # "point", "analytic", "bootstrap", or "posterior"
    estimand: str = ""
    source: str = ""
    diagnostics: Mapping[str, float] = field(default_factory=dict)
    _values: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty(0))
    _index: npt.NDArray[Any] | None = field(default=None, repr=False)
    _std_error: npt.NDArray[np.float64] | None = field(default=None, repr=False)
    _draws: npt.NDArray[np.float64] | None = field(default=None, repr=False)

    # -- the per-unit surface ------------------------------------------------

    @property
    def values(self) -> npt.NDArray[np.float64]:
        """The per-unit point estimates, ``(n,)``, ordered to the input rows.

        When only draws are present, the posterior or bootstrap mean per unit.
        """
        if self._values.size or self._draws is None:
            return np.asarray(self._values, dtype=np.float64)
        return np.asarray(self._draws.mean(axis=0), dtype=np.float64)

    def __array__(self, dtype: Any = None) -> npt.NDArray[np.float64]:
        """The per-unit point vector, so ``np.asarray(result)`` and arithmetic work."""
        out = self.values
        return out.astype(dtype) if dtype is not None else out

    @property
    def shape(self) -> tuple[int, ...]:
        """The shape of the per-unit point vector, ``(n,)``."""
        return self.values.shape

    def __getitem__(self, item: Any) -> Any:
        """Index into the per-unit point vector (so ``result[i]`` is the unit's effect)."""
        return self.values[item]

    def __iter__(self) -> Any:
        return iter(self.values)

    # Element-wise comparisons against the point vector, so a wrapped result is a
    # drop-in for the bare array the SCM and the meta-learners returned before.
    def __eq__(self, other: Any) -> Any:
        return self.values == np.asarray(other)

    def __ne__(self, other: Any) -> Any:
        return self.values != np.asarray(other)

    def __lt__(self, other: Any) -> Any:
        return self.values < np.asarray(other)

    def __le__(self, other: Any) -> Any:
        return self.values <= np.asarray(other)

    def __gt__(self, other: Any) -> Any:
        return self.values > np.asarray(other)

    def __ge__(self, other: Any) -> Any:
        return self.values >= np.asarray(other)

    __hash__ = None  # type: ignore[assignment]

    @property
    def index(self) -> npt.NDArray[Any]:
        """The row ids for the units, positional ``arange(n)`` when none were given."""
        if self._index is not None:
            return np.asarray(self._index)
        return np.arange(self.n)

    @property
    def std_error(self) -> npt.NDArray[np.float64] | None:
        """Per-unit standard errors, or the per-unit draw standard deviations."""
        if self._std_error is not None:
            return np.asarray(self._std_error, dtype=np.float64)
        if self._draws is not None:
            return np.asarray(self._draws.std(axis=0, ddof=1), dtype=np.float64)
        return None

    def interval(
        self, level: float = 0.9
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        """Per-unit confidence or credible intervals, ``(lo[n], hi[n])`` at ``level``.

        Draw quantiles when draws are present, else the normal interval from the
        per-unit point and standard error. Raises when there is no uncertainty.
        """
        if not 0.0 < level < 1.0:
            raise ValueError(f"level must be in (0, 1), got {level}")
        if self._draws is not None:
            lo = np.quantile(self._draws, (1.0 - level) / 2.0, axis=0)
            hi = np.quantile(self._draws, (1.0 + level) / 2.0, axis=0)
            return np.asarray(lo, dtype=np.float64), np.asarray(hi, dtype=np.float64)
        if self._std_error is not None:
            z = _norm_ppf((1.0 + level) / 2.0)
            point = self.values
            se = np.asarray(self._std_error, dtype=np.float64)
            return point - z * se, point + z * se
        raise ValueError("per-unit effect has no uncertainty to form an interval from")

    def draws(self) -> npt.NDArray[np.float64] | None:
        """The ``(S, n)`` bootstrap or posterior draws, or ``None`` for a point result."""
        if self._draws is None:
            return None
        return np.array(self._draws, copy=True)

    def ate(self) -> EffectEstimate:
        """The population summary: the average per-unit effect, as an ``EffectEstimate``.

        The bridge back to the scalar family. With draws, the per-unit draws average
        to a draw vector of the mean effect (so the population interval reflects the
        joint uncertainty). With per-unit analytic SEs, the point is the mean and the
        standard error combines them assuming independence. With neither, a point.
        """
        target = EffectTarget.CATE if self.estimand == "cate" else EffectTarget.ITE
        learners: dict[str, str] = {}
        if self._draws is not None:
            mean_draws = np.asarray(self._draws.mean(axis=1), dtype=np.float64)
            return EffectEstimate(
                target=target,
                estimator=self.source,
                learners=learners,
                n=self.n,
                kind=self.kind if self.kind in ("bootstrap", "posterior") else "bootstrap",
                diagnostics=dict(self.diagnostics),
                _draws=mean_draws,
            )
        point = float(np.mean(self.values))
        if self._std_error is not None:
            se = np.asarray(self._std_error, dtype=np.float64)
            combined = float(np.sqrt(np.sum(se * se)) / se.shape[0])
            return EffectEstimate(
                target=target,
                estimator=self.source,
                learners=learners,
                n=self.n,
                kind="analytic",
                diagnostics=dict(self.diagnostics),
                _point=point,
                _std_error=combined,
            )
        return EffectEstimate(
            target=target,
            estimator=self.source,
            learners=learners,
            n=self.n,
            kind="point",
            diagnostics=dict(self.diagnostics),
            _point=point,
        )

    def to_frame(self) -> pd.DataFrame:
        """A per-unit table (index, effect, and lo/hi/se when available). Pandas only here."""
        import pandas as pd

        columns: dict[str, npt.NDArray[Any]] = {"effect": self.values}
        se = self.std_error
        if se is not None:
            columns["std_error"] = se
        try:
            lo, hi = self.interval()
            columns["ci_low"] = lo
            columns["ci_high"] = hi
        except ValueError:
            pass
        return pd.DataFrame(columns, index=pd.Index(self.index, name="unit"))

    def __len__(self) -> int:
        return self.n

    def __repr__(self) -> str:
        v = self.values
        head = f"{self.estimand or 'effect'}"
        body = (
            f"ate={float(np.mean(v)):.4g}, "
            f"min={float(np.min(v)):.3g}, median={float(np.median(v)):.3g}, "
            f"max={float(np.max(v)):.3g}"
        )
        name = type(self).__name__
        return f"{name}({self.source or head}, {self.kind}: {body}, n={self.n})"

index property

The row ids for the units, positional arange(n) when none were given.

shape property

The shape of the per-unit point vector, (n,).

std_error property

Per-unit standard errors, or the per-unit draw standard deviations.

values property

The per-unit point estimates, (n,), ordered to the input rows.

When only draws are present, the posterior or bootstrap mean per unit.

__array__(dtype=None)

The per-unit point vector, so np.asarray(result) and arithmetic work.

Source code in src/ergodic/inference/_perunit.py

def __array__(self, dtype: Any = None) -> npt.NDArray[np.float64]:
    """The per-unit point vector, so ``np.asarray(result)`` and arithmetic work."""
    out = self.values
    return out.astype(dtype) if dtype is not None else out

__getitem__(item)

Index into the per-unit point vector (so result[i] is the unit's effect).

Source code in src/ergodic/inference/_perunit.py

def __getitem__(self, item: Any) -> Any:
    """Index into the per-unit point vector (so ``result[i]`` is the unit's effect)."""
    return self.values[item]

ate()

The population summary: the average per-unit effect, as an EffectEstimate.

The bridge back to the scalar family. With draws, the per-unit draws average to a draw vector of the mean effect (so the population interval reflects the joint uncertainty). With per-unit analytic SEs, the point is the mean and the standard error combines them assuming independence. With neither, a point.

Source code in src/ergodic/inference/_perunit.py

def ate(self) -> EffectEstimate:
    """The population summary: the average per-unit effect, as an ``EffectEstimate``.

    The bridge back to the scalar family. With draws, the per-unit draws average
    to a draw vector of the mean effect (so the population interval reflects the
    joint uncertainty). With per-unit analytic SEs, the point is the mean and the
    standard error combines them assuming independence. With neither, a point.
    """
    target = EffectTarget.CATE if self.estimand == "cate" else EffectTarget.ITE
    learners: dict[str, str] = {}
    if self._draws is not None:
        mean_draws = np.asarray(self._draws.mean(axis=1), dtype=np.float64)
        return EffectEstimate(
            target=target,
            estimator=self.source,
            learners=learners,
            n=self.n,
            kind=self.kind if self.kind in ("bootstrap", "posterior") else "bootstrap",
            diagnostics=dict(self.diagnostics),
            _draws=mean_draws,
        )
    point = float(np.mean(self.values))
    if self._std_error is not None:
        se = np.asarray(self._std_error, dtype=np.float64)
        combined = float(np.sqrt(np.sum(se * se)) / se.shape[0])
        return EffectEstimate(
            target=target,
            estimator=self.source,
            learners=learners,
            n=self.n,
            kind="analytic",
            diagnostics=dict(self.diagnostics),
            _point=point,
            _std_error=combined,
        )
    return EffectEstimate(
        target=target,
        estimator=self.source,
        learners=learners,
        n=self.n,
        kind="point",
        diagnostics=dict(self.diagnostics),
        _point=point,
    )

draws()

The (S, n) bootstrap or posterior draws, or None for a point result.

Source code in src/ergodic/inference/_perunit.py

def draws(self) -> npt.NDArray[np.float64] | None:
    """The ``(S, n)`` bootstrap or posterior draws, or ``None`` for a point result."""
    if self._draws is None:
        return None
    return np.array(self._draws, copy=True)

interval(level=0.9)

Per-unit confidence or credible intervals, (lo[n], hi[n]) at level.

Draw quantiles when draws are present, else the normal interval from the per-unit point and standard error. Raises when there is no uncertainty.

Source code in src/ergodic/inference/_perunit.py

def interval(
    self, level: float = 0.9
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
    """Per-unit confidence or credible intervals, ``(lo[n], hi[n])`` at ``level``.

    Draw quantiles when draws are present, else the normal interval from the
    per-unit point and standard error. Raises when there is no uncertainty.
    """
    if not 0.0 < level < 1.0:
        raise ValueError(f"level must be in (0, 1), got {level}")
    if self._draws is not None:
        lo = np.quantile(self._draws, (1.0 - level) / 2.0, axis=0)
        hi = np.quantile(self._draws, (1.0 + level) / 2.0, axis=0)
        return np.asarray(lo, dtype=np.float64), np.asarray(hi, dtype=np.float64)
    if self._std_error is not None:
        z = _norm_ppf((1.0 + level) / 2.0)
        point = self.values
        se = np.asarray(self._std_error, dtype=np.float64)
        return point - z * se, point + z * se
    raise ValueError("per-unit effect has no uncertainty to form an interval from")

to_frame()

A per-unit table (index, effect, and lo/hi/se when available). Pandas only here.

Source code in src/ergodic/inference/_perunit.py

def to_frame(self) -> pd.DataFrame:
    """A per-unit table (index, effect, and lo/hi/se when available). Pandas only here."""
    import pandas as pd

    columns: dict[str, npt.NDArray[Any]] = {"effect": self.values}
    se = self.std_error
    if se is not None:
        columns["std_error"] = se
    try:
        lo, hi = self.interval()
        columns["ci_low"] = lo
        columns["ci_high"] = hi
    except ValueError:
        pass
    return pd.DataFrame(columns, index=pd.Index(self.index, name="unit"))

ergodic.inference.Heterogeneity dataclass

Tier-B heterogeneity characterization on shared doubly-robust scores.

Carries the cross-fit :class:~ergodic.inference._scores.DRScores and the fitted per-unit tau_hat. The honesty register: the rankings and proxies are fixed (computed cross-fit), so the orthogonal summaries below inherit valid inference while the CATE point predictions themselves never carry a confidence interval.

Source code in src/ergodic/inference/explain/_record.py

@dataclass(frozen=True)
class Heterogeneity:
    """Tier-B heterogeneity characterization on shared doubly-robust scores.

    Carries the cross-fit :class:`~ergodic.inference._scores.DRScores` and the fitted
    per-unit ``tau_hat``. The honesty register: the rankings and proxies are fixed
    (computed cross-fit), so the orthogonal summaries below inherit valid inference
    while the CATE point predictions themselves never carry a confidence interval.
    """

    scores: DRScores
    tau_hat: npt.NDArray[np.float64]
    treatment: str
    outcome: str

    def blp(self, *, level: float = 0.9) -> BLPResult:
        """Best Linear Projection of the CATE on the covariates (HC3 inference)."""
        return best_linear_projection(
            self.scores.scores,
            self.scores.covariates,
            self.scores.covariate_names,
            level=level,
        )

    def gates(self, *, n_groups: int = 5, level: float = 0.9) -> GATESResult:
        """Sorted Group ATEs across ``n_groups`` predicted-effect quantiles."""
        return gates(self.scores.scores, self.tau_hat, n_groups=n_groups, level=level)

    def clan(self, *, n_groups: int = 5, level: float = 0.9) -> CLANResult:
        """Covariate-mean differences between the most- and least-affected groups."""
        return clan(
            self.scores.covariates,
            self.scores.covariate_names,
            self.tau_hat,
            n_groups=n_groups,
            level=level,
        )

    def calibration_test(self) -> CalibrationResult:
        """The omnibus calibration / heterogeneity test (one-sided HTE p-value)."""
        return calibration_test(self.scores.scores, self.tau_hat)

    def rate(
        self, metric: str = "autoc", *, n_boot: int = 200, random_state: int | None = None
    ) -> RATEResult:
        """The Rank-Weighted ATE of the ``tau_hat`` ranking: ``"autoc"`` or ``"qini"``."""
        return rate(
            self.scores.scores,
            self.tau_hat,
            metric=metric,
            n_boot=n_boot,
            random_state=random_state,
        )

    def policy_tree(
        self,
        covariates: str | tuple[str, ...] | list[str] | None = None,
        *,
        depth: int = 2,
    ) -> PolicyTree:
        """The depth-limited, DR-welfare-optimal policy tree off these scores.

        A convenience that forwards the shared doubly-robust scores to
        :func:`~ergodic.inference.policy_tree`. ``covariates`` selects the split
        columns by name (``None`` uses all of them). The "act" companion to
        :meth:`rate`: where ``rate`` scores a ranking, this returns an interpretable
        treat-or-not rule and its estimated policy value.
        """
        from ._policy import policy_tree as _policy_tree

        return _policy_tree(self.scores, covariates, depth=depth)

blp(*, level=0.9)

Best Linear Projection of the CATE on the covariates (HC3 inference).

Source code in src/ergodic/inference/explain/_record.py

def blp(self, *, level: float = 0.9) -> BLPResult:
    """Best Linear Projection of the CATE on the covariates (HC3 inference)."""
    return best_linear_projection(
        self.scores.scores,
        self.scores.covariates,
        self.scores.covariate_names,
        level=level,
    )

calibration_test()

The omnibus calibration / heterogeneity test (one-sided HTE p-value).

Source code in src/ergodic/inference/explain/_record.py

def calibration_test(self) -> CalibrationResult:
    """The omnibus calibration / heterogeneity test (one-sided HTE p-value)."""
    return calibration_test(self.scores.scores, self.tau_hat)

clan(*, n_groups=5, level=0.9)

Covariate-mean differences between the most- and least-affected groups.

Source code in src/ergodic/inference/explain/_record.py

def clan(self, *, n_groups: int = 5, level: float = 0.9) -> CLANResult:
    """Covariate-mean differences between the most- and least-affected groups."""
    return clan(
        self.scores.covariates,
        self.scores.covariate_names,
        self.tau_hat,
        n_groups=n_groups,
        level=level,
    )

gates(*, n_groups=5, level=0.9)

Sorted Group ATEs across n_groups predicted-effect quantiles.

Source code in src/ergodic/inference/explain/_record.py

def gates(self, *, n_groups: int = 5, level: float = 0.9) -> GATESResult:
    """Sorted Group ATEs across ``n_groups`` predicted-effect quantiles."""
    return gates(self.scores.scores, self.tau_hat, n_groups=n_groups, level=level)

policy_tree(covariates=None, *, depth=2)

The depth-limited, DR-welfare-optimal policy tree off these scores.

A convenience that forwards the shared doubly-robust scores to :func:~ergodic.inference.policy_tree. covariates selects the split columns by name (None uses all of them). The "act" companion to :meth:rate: where rate scores a ranking, this returns an interpretable treat-or-not rule and its estimated policy value.

Source code in src/ergodic/inference/explain/_record.py

def policy_tree(
    self,
    covariates: str | tuple[str, ...] | list[str] | None = None,
    *,
    depth: int = 2,
) -> PolicyTree:
    """The depth-limited, DR-welfare-optimal policy tree off these scores.

    A convenience that forwards the shared doubly-robust scores to
    :func:`~ergodic.inference.policy_tree`. ``covariates`` selects the split
    columns by name (``None`` uses all of them). The "act" companion to
    :meth:`rate`: where ``rate`` scores a ranking, this returns an interpretable
    treat-or-not rule and its estimated policy value.
    """
    from ._policy import policy_tree as _policy_tree

    return _policy_tree(self.scores, covariates, depth=depth)

rate(metric='autoc', *, n_boot=200, random_state=None)

The Rank-Weighted ATE of the tau_hat ranking: "autoc" or "qini".

Source code in src/ergodic/inference/explain/_record.py

def rate(
    self, metric: str = "autoc", *, n_boot: int = 200, random_state: int | None = None
) -> RATEResult:
    """The Rank-Weighted ATE of the ``tau_hat`` ranking: ``"autoc"`` or ``"qini"``."""
    return rate(
        self.scores.scores,
        self.tau_hat,
        metric=metric,
        n_boot=n_boot,
        random_state=random_state,
    )

Causal forest, policy trees, and CATE model selection

The honest causal forest (GRF) is its own CATE learner, reachable through the cate facade (method="causal_forest" or "grf") or directly as CausalForest; it returns a CATEResult with native per-unit infinitesimal-jackknife intervals, and CausalForest(...).fit(...).variable_importance() returns a VariableImportance driver signal. policy_tree (also CATEResult.explain().policy_tree(...)) learns an interpretable, depth-limited, doubly-robust-optimal treat/don't PolicyTree. select_cate ranks candidate CATE learners by the doubly-robust tau-risk and returns a SelectionReport with the selected result and a Q-aggregation ensemble.

ergodic.CausalForest

Bases: _MetaLearnerBase

An honest causal forest (GRF) for the conditional average treatment effect.

The forest is the CATE learner. The outcome and propensity nuisance slots take any learner (they residualize Y and W on the covariates before the forest grows); the forest machinery, honest splitting and the kernel readout, is native. The ATE is the average per-unit effect, and cate returns a :class:~ergodic.inference.CATEResult with per-unit infinitesimal-jackknife standard errors and intervals.

Hyperparameters: n_trees (forest size), min_leaf (the minimum estimate-half rows per leaf), max_depth, honest (split/estimate sample-splitting, on by default), subsample (the per-tree subsample fraction), and random_state. Binary treatment in this phase; a continuous treatment via the same local moment is deferred and raises.

Source code in src/ergodic/inference/estimators/_forest.py

class CausalForest(_MetaLearnerBase):
    """An honest causal forest (GRF) for the conditional average treatment effect.

    The forest is the CATE learner. The ``outcome`` and ``propensity`` nuisance slots
    take any learner (they residualize ``Y`` and ``W`` on the covariates before the
    forest grows); the forest machinery, honest splitting and the kernel readout, is
    native. The ATE is the average per-unit effect, and ``cate`` returns a
    :class:`~ergodic.inference.CATEResult` with per-unit infinitesimal-jackknife
    standard errors and intervals.

    Hyperparameters: ``n_trees`` (forest size), ``min_leaf`` (the minimum estimate-half
    rows per leaf), ``max_depth``, ``honest`` (split/estimate sample-splitting, on by
    default), ``subsample`` (the per-tree subsample fraction), and ``random_state``.
    Binary treatment in this phase; a continuous treatment via the same local moment
    is deferred and raises.
    """

    name = "causal_forest"
    slots = (
        Slot("outcome", Task.REGRESSION),
        Slot("propensity", Task.CLASSIFICATION),
    )
    targets = frozenset({EffectTarget.ATE})

    def __init__(
        self,
        learners: Any = None,
        *,
        n_trees: int = 200,
        min_leaf: int = 5,
        max_depth: int = 8,
        honest: bool = True,
        subsample: float = 0.5,
        max_features: float | int | None = None,
        alpha: float = 0.05,
        n_folds: int = 5,
        clip: float = 1e-3,
        **kwargs: Any,
    ) -> None:
        super().__init__(learners, **kwargs)
        if not 0.0 < subsample <= 1.0:
            raise ValueError(f"subsample must be in (0, 1], got {subsample}")
        if n_trees < 1:
            raise ValueError(f"n_trees must be at least 1, got {n_trees}")
        if min_leaf < 1:
            raise ValueError(f"min_leaf must be at least 1, got {min_leaf}")
        if not 0.0 <= alpha < 0.5:
            raise ValueError(f"alpha must be in [0, 0.5), got {alpha}")
        self.n_trees = n_trees
        self.min_leaf = min_leaf
        self.max_depth = max_depth
        self.honest = honest
        self.subsample = subsample
        self.max_features = max_features
        self.alpha = alpha
        self.n_folds = n_folds
        self.clip = clip

    # -- residualization (the forest R-learner front end) --------------------

    def _residualize(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        """Cross-fit ``m(x)`` and ``e(x)`` and return the residuals ``W - e``, ``Y - m``.

        Reuses the shared cross-fit nuisance routine: the outcome model carries the
        treatment as a feature (so ``m = e * mu1 + (1 - e) * mu0``), and the propensity
        is the treatment model. The residuals are what the forest splits on.
        """
        a = arrays["A"]
        y = arrays["Y"]
        z = arrays["Z"]
        mu0, mu1, e = _crossfit_nuisances(
            a,
            y,
            z,
            make_outcome=make["outcome"],
            make_propensity=make.get("propensity") or (lambda: LogisticRegression()),
            n_folds=self.n_folds,
            clip=self.clip,
            random_state=self._random_state,
        )
        m = e * mu1 + (1.0 - e) * mu0
        return a - e, y - m

    def _max_features_for(self, n_features: int) -> int:
        if self.max_features is None:
            # GRF default: ceil(sqrt(p) + 20), capped at p
            return min(n_features, int(np.ceil(np.sqrt(n_features) + 20)))
        if isinstance(self.max_features, int):
            return max(1, min(n_features, self.max_features))
        return max(1, min(n_features, int(np.ceil(self.max_features * n_features))))

    # -- the meta-learner hooks ----------------------------------------------

    def _fit_cate(
        self, arrays: dict[str, npt.NDArray[np.float64]], make: dict[str, Callable[[], Learner]]
    ) -> _ForestModel:
        w_res, y_res = self._residualize(arrays, make)
        x = arrays["Z"]
        n, n_features = x.shape
        rng = np.random.default_rng(self._random_state)
        sub_n = max(2 * self.min_leaf, int(np.ceil(self.subsample * n)))
        sub_n = min(sub_n, n)
        n_try = self._max_features_for(n_features)
        trees: list[_Tree] = []
        importance = np.zeros(n_features, dtype=np.float64)
        for _ in range(self.n_trees):
            in_bag = np.asarray(rng.choice(n, size=sub_n, replace=False), dtype=np.int64)
            if self.honest and in_bag.shape[0] >= 2:
                shuffled = rng.permutation(in_bag)
                half = shuffled.shape[0] // 2
                split_rows = shuffled[:half]
                estimate_rows = shuffled[half:]
            else:
                split_rows = in_bag
                estimate_rows = in_bag
            builder = _TreeBuilder(
                x,
                w_res,
                y_res,
                min_leaf=self.min_leaf,
                max_depth=self.max_depth,
                n_features_try=n_try,
                alpha=self.alpha,
                rng=rng,
            )
            root = builder.build(split_rows, estimate_rows)
            trees.append(_Tree(root=root, in_bag=in_bag, split_importance=builder.importance))
            importance += builder.importance
        total = float(importance.sum())
        if total > 0:
            importance = importance / total
        return _ForestModel(
            trees=trees,
            x_train=x,
            w_res=w_res,
            y_res=y_res,
            importance=importance,
            covariate_names=self._covariates,
        )

    def _predict_cate(
        self, model: _ForestModel, z: npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64]:
        return model.predict(z)

    def _analytic(
        self,
        arrays: dict[str, npt.NDArray[np.float64]],
        target: EffectTarget,
        make: dict[str, Callable[[], Learner]],
    ) -> tuple[float, npt.NDArray[np.float64]] | None:
        """The doubly-robust AIPW ATE with its influence-function standard error.

        GRF averages per-unit AIPW scores for a doubly-robust, efficient ATE rather
        than bootstrapping the forest. The score is
        ``Gamma_i = tau_hat(X_i) + (W_i - e_i) / (e_i (1 - e_i)) (Y_i - mu_hat_i)``
        with ``mu_hat_i = W_i mu1_i + (1 - W_i) mu0_i``; its mean is the ATE and its
        scaled standard deviation the standard error. Reusing this for ``effect()``
        keeps the scalar ATE path off the (prohibitively slow) forest-refit bootstrap.
        Computed only on the fit arrays (``_analytic`` on a resample returns ``None``).
        """
        if target is not EffectTarget.ATE:
            return None
        model = getattr(self, "_fitted_cate", None)
        if model is None or arrays["A"].shape[0] != model.x_train.shape[0]:
            return None
        a, y, z = arrays["A"], arrays["Y"], arrays["Z"]
        try:
            mu0, mu1, e = _crossfit_nuisances(
                a,
                y,
                z,
                make_outcome=make["outcome"],
                make_propensity=make.get("propensity") or (lambda: LogisticRegression()),
                n_folds=self.n_folds,
                clip=self.clip,
                random_state=self._random_state,
            )
        except (ValueError, np.linalg.LinAlgError):
            return None
        tau_hat = model.predict(z)
        mu_hat = a * mu1 + (1.0 - a) * mu0
        gamma = tau_hat + (a - e) / (e * (1.0 - e)) * (y - mu_hat)
        return float(np.mean(gamma)), np.asarray(gamma, dtype=np.float64)

    # -- the per-unit result, with the infinitesimal-jackknife variance ------

    def cate(self, data: TabularData | None = None, *, uncertainty: bool = True) -> CATEResult:
        """Per-unit CATE with infinitesimal-jackknife standard errors and intervals.

        Overrides the base meta-learner ``cate`` (which would bootstrap-refit the
        whole forest). The point is the forest-kernel local effect; the per-unit
        variance is the infinitesimal jackknife over the trees (the
        bootstrap-of-little-bags grouping), valid because the forest is honest and
        subsampled. The result carries the cached doubly-robust scores for
        :meth:`~ergodic.inference.CATEResult.explain`.

        ``uncertainty=False`` returns the point effects with ``kind="point"`` and no
        standard error, skipping the infinitesimal-jackknife pass over the trees.
        """
        from .._perunit import CATEResult

        model = self._cate_model_forest()
        z = self._cate_covariates(data)
        point = np.asarray(model.predict(z), dtype=np.float64)
        n = point.shape[0]
        std_error = self._ij_standard_error(model, z, point) if uncertainty else None
        kind = "analytic" if uncertainty else "point"
        scores = self._cate_scores(data)
        return CATEResult(
            treatment=self._treatment_name,
            outcome=self._outcome_name,
            covariates=self._covariates,
            n=n,
            kind=kind,
            source=self.name,
            diagnostics=dict(self._state.diagnostics),
            identification=self._identification,
            _values=point,
            _index=self._cate_index(data, n),
            _std_error=std_error,
            _scores=scores,
        )

    def _cate_model_forest(self) -> _ForestModel:
        model = getattr(self, "_fitted_cate", None)
        if model is None:
            raise RuntimeError(f"{self.name} is not fitted; call fit first")
        return model  # type: ignore[no-any-return]

    def _ij_standard_error(
        self,
        model: _ForestModel,
        z: npt.NDArray[np.float64],
        point: npt.NDArray[np.float64],
    ) -> npt.NDArray[np.float64]:
        """Per-unit infinitesimal-jackknife standard errors over the trees.

        The IJ variance (Wager-Athey Eq. 8, the bootstrap-of-little-bags form)
        is ``V(x) = ((n - 1)/n)(n/(n - s))^2 sum_i Cov_b[t_b(x), N_ib]^2``, where
        ``t_b(x)`` is tree ``b``'s prediction at ``x``, ``N_ib`` is the inclusion count
        of training row ``i`` in tree ``b`` (here 0/1 under subsampling), and ``s`` is
        the subsample size. A Monte-Carlo correction subtracts the within-tree
        sampling noise so the estimate stays non-negative and is not inflated by a
        finite tree count. Honesty and subsampling are what make this valid.
        """
        n_trees = len(model.trees)
        n_query = z.shape[0]
        n_train = model.x_train.shape[0]
        if n_trees < 2 or n_train < 2:
            return np.zeros(n_query, dtype=np.float64)
        # per-tree predictions, (B, n_query)
        preds = np.empty((n_trees, n_query), dtype=np.float64)
        inbag = np.zeros((n_trees, n_train), dtype=np.float64)
        for b, tree in enumerate(model.trees):
            preds[b] = model.tree_predict(tree, z)
            inbag[b, tree.in_bag] = 1.0
        s = float(inbag.sum(axis=1).mean())
        pred_centered = preds - preds.mean(axis=0, keepdims=True)
        inbag_centered = inbag - inbag.mean(axis=0, keepdims=True)
        # Cov_b[t_b(x), N_ib] over trees, shape (n_query, n_train)
        cov = (pred_centered.T @ inbag_centered) / n_trees
        raw = np.sum(cov * cov, axis=1)
        # bootstrap-of-little-bags Monte-Carlo bias correction
        pred_var = pred_centered.var(axis=0)  # (n_query,)
        inbag_var = float(inbag_centered.var())
        correction = n_train * inbag_var * pred_var / n_trees
        var_unscaled = np.clip(raw - correction, 0.0, np.inf)
        if s >= n_train:
            scale = (n_train - 1.0) / n_train
        else:
            scale = ((n_train - 1.0) / n_train) * (n_train / (n_train - s)) ** 2
        var = scale * var_unscaled
        return np.asarray(np.sqrt(np.clip(var, 0.0, np.inf)), dtype=np.float64)

    # -- variable importance -------------------------------------------------

    def variable_importance(self) -> VariableImportance:
        """Depth-weighted, normalized split frequencies over the covariates.

        A driver signal: which covariates the forest splits on, weighted toward the
        shallow splits that move the effect most. Available after ``fit``.
        """
        model = self._cate_model_forest()
        return VariableImportance(names=model.covariate_names, values=model.importance)

cate(data=None, *, uncertainty=True)

Per-unit CATE with infinitesimal-jackknife standard errors and intervals.

Overrides the base meta-learner cate (which would bootstrap-refit the whole forest). The point is the forest-kernel local effect; the per-unit variance is the infinitesimal jackknife over the trees (the bootstrap-of-little-bags grouping), valid because the forest is honest and subsampled. The result carries the cached doubly-robust scores for :meth:~ergodic.inference.CATEResult.explain.

uncertainty=False returns the point effects with kind="point" and no standard error, skipping the infinitesimal-jackknife pass over the trees.

Source code in src/ergodic/inference/estimators/_forest.py

def cate(self, data: TabularData | None = None, *, uncertainty: bool = True) -> CATEResult:
    """Per-unit CATE with infinitesimal-jackknife standard errors and intervals.

    Overrides the base meta-learner ``cate`` (which would bootstrap-refit the
    whole forest). The point is the forest-kernel local effect; the per-unit
    variance is the infinitesimal jackknife over the trees (the
    bootstrap-of-little-bags grouping), valid because the forest is honest and
    subsampled. The result carries the cached doubly-robust scores for
    :meth:`~ergodic.inference.CATEResult.explain`.

    ``uncertainty=False`` returns the point effects with ``kind="point"`` and no
    standard error, skipping the infinitesimal-jackknife pass over the trees.
    """
    from .._perunit import CATEResult

    model = self._cate_model_forest()
    z = self._cate_covariates(data)
    point = np.asarray(model.predict(z), dtype=np.float64)
    n = point.shape[0]
    std_error = self._ij_standard_error(model, z, point) if uncertainty else None
    kind = "analytic" if uncertainty else "point"
    scores = self._cate_scores(data)
    return CATEResult(
        treatment=self._treatment_name,
        outcome=self._outcome_name,
        covariates=self._covariates,
        n=n,
        kind=kind,
        source=self.name,
        diagnostics=dict(self._state.diagnostics),
        identification=self._identification,
        _values=point,
        _index=self._cate_index(data, n),
        _std_error=std_error,
        _scores=scores,
    )

variable_importance()

Depth-weighted, normalized split frequencies over the covariates.

A driver signal: which covariates the forest splits on, weighted toward the shallow splits that move the effect most. Available after fit.

Source code in src/ergodic/inference/estimators/_forest.py

def variable_importance(self) -> VariableImportance:
    """Depth-weighted, normalized split frequencies over the covariates.

    A driver signal: which covariates the forest splits on, weighted toward the
    shallow splits that move the effect most. Available after ``fit``.
    """
    model = self._cate_model_forest()
    return VariableImportance(names=model.covariate_names, values=model.importance)

ergodic.inference.VariableImportance dataclass

Depth-weighted, normalized split frequencies: a driver signal over covariates.

names are the covariate columns and values the matching importances, which sum to one (or are all zero for a forest that never split). A split at depth d contributes 1 / (d + 1), so shallow splits, the ones that move the effect most, weigh more, following the GRF variable-importance convention.

Source code in src/ergodic/inference/estimators/_forest.py

@dataclass(frozen=True)
class VariableImportance:
    """Depth-weighted, normalized split frequencies: a driver signal over covariates.

    ``names`` are the covariate columns and ``values`` the matching importances, which
    sum to one (or are all zero for a forest that never split). A split at depth ``d``
    contributes ``1 / (d + 1)``, so shallow splits, the ones that move the effect most,
    weigh more, following the GRF variable-importance convention.
    """

    names: tuple[str, ...]
    values: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty(0))

    def top(self, k: int | None = None) -> list[tuple[str, float]]:
        """The covariates by descending importance (all of them when ``k`` is ``None``)."""
        order = np.argsort(self.values)[::-1]
        pairs = [(self.names[i], float(self.values[i])) for i in order]
        return pairs if k is None else pairs[:k]

    def to_dict(self) -> dict[str, float]:
        return {name: float(value) for name, value in zip(self.names, self.values, strict=True)}

top(k=None)

The covariates by descending importance (all of them when k is None).

Source code in src/ergodic/inference/estimators/_forest.py

def top(self, k: int | None = None) -> list[tuple[str, float]]:
    """The covariates by descending importance (all of them when ``k`` is ``None``)."""
    order = np.argsort(self.values)[::-1]
    pairs = [(self.names[i], float(self.values[i])) for i in order]
    return pairs if k is None else pairs[:k]

ergodic.policy_tree(result_or_scores, covariates=None, *, depth=2)

Fit a shallow doubly-robust optimal-policy tree of depth depth.

result_or_scores is a :class:~ergodic.inference.CATEResult (its cached cross-fit scores are read) or a :class:~ergodic.inference._scores.DRScores directly. covariates names the split covariates (a name, a list of names, or None for every covariate carried by the scores). The tree maximizes the doubly-robust welfare W(pi) = mean[(2 pi(X) - 1) Gamma_i]; at depth <= 2 the search is exhaustive over thresholds and so globally optimal, and a deeper tree uses a greedy recursive best-split (documented, not guaranteed optimal).

Returns a frozen :class:PolicyTree carrying the rule, the estimated policy value with a standard error from the per-unit influence terms, and the regret against treating everyone / no one.

Source code in src/ergodic/inference/explain/_policy.py

def policy_tree(
    result_or_scores: CATEResult | DRScores,
    covariates: str | tuple[str, ...] | list[str] | None = None,
    *,
    depth: int = 2,
) -> PolicyTree:
    """Fit a shallow doubly-robust optimal-policy tree of depth ``depth``.

    ``result_or_scores`` is a :class:`~ergodic.inference.CATEResult` (its cached
    cross-fit scores are read) or a :class:`~ergodic.inference._scores.DRScores`
    directly. ``covariates`` names the split covariates (a name, a list of names, or
    ``None`` for every covariate carried by the scores). The tree maximizes the
    doubly-robust welfare ``W(pi) = mean[(2 pi(X) - 1) Gamma_i]``; at ``depth <= 2``
    the search is exhaustive over thresholds and so globally optimal, and a deeper
    tree uses a greedy recursive best-split (documented, not guaranteed optimal).

    Returns a frozen :class:`PolicyTree` carrying the rule, the estimated policy
    value with a standard error from the per-unit influence terms, and the regret
    against treating everyone / no one.
    """
    if depth < 1:
        raise ValueError(f"depth must be at least 1, got {depth}")
    scores, full_design, all_names, treatment, outcome = _extract(result_or_scores)
    n = scores.shape[0]
    if n < 2:
        raise ValueError(f"policy_tree needs at least 2 units, got {n}")
    design, names = _select_columns(full_design, all_names, covariates)

    exhaustive = depth <= _EXHAUSTIVE_DEPTH
    if exhaustive:
        _, root = _best_split(scores, design, names, depth)
    else:
        _, root = _greedy_split(scores, design, names, depth)

    actions = np.asarray([root.predict(row) for row in design], dtype=np.float64)
    influence = (2.0 * actions - 1.0) * scores
    value = float(np.mean(influence))
    std_error = float(np.std(influence, ddof=1) / np.sqrt(n)) if n > 1 else float("nan")
    treat_all = float(np.mean(scores))
    treat_none = float(np.mean(-scores))
    return PolicyTree(
        root=root,
        covariates=names,
        n=n,
        depth=depth,
        value=value,
        std_error=std_error,
        treat_all_value=treat_all,
        treat_none_value=treat_none,
        exhaustive=exhaustive,
        treatment=treatment,
        outcome=outcome,
        _scores=scores,
        _design=design,
    )

ergodic.PolicyTree dataclass

A shallow doubly-robust optimal-policy tree with its value and regret.

root is the learned tree over the named covariates; :meth:predict and :meth:assign evaluate it. value/std_error are the estimated policy value W(pi_hat) and its standard error from the per-unit influence terms; treat_all_value and treat_none_value are the two trivial baselines, and :meth:regret is the value gap against either. depth is the requested depth and exhaustive records whether the globally optimal search was used.

Source code in src/ergodic/inference/explain/_policy.py

@dataclass(frozen=True)
class PolicyTree:
    """A shallow doubly-robust optimal-policy tree with its value and regret.

    ``root`` is the learned tree over the named ``covariates``; :meth:`predict` and
    :meth:`assign` evaluate it. ``value``/``std_error`` are the estimated policy
    value ``W(pi_hat)`` and its standard error from the per-unit influence terms;
    ``treat_all_value`` and ``treat_none_value`` are the two trivial baselines, and
    :meth:`regret` is the value gap against either. ``depth`` is the requested depth
    and ``exhaustive`` records whether the globally optimal search was used.
    """

    root: PolicyNode
    covariates: tuple[str, ...]
    n: int
    depth: int
    value: float
    std_error: float
    treat_all_value: float
    treat_none_value: float
    exhaustive: bool
    treatment: str = ""
    outcome: str = ""
    _scores: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty(0), repr=False)
    _design: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty((0, 0)), repr=False)

    # -- evaluation ----------------------------------------------------------

    def predict(self, covariates: npt.NDArray[np.float64]) -> npt.NDArray[np.int64]:
        """The per-row action (0/1) for a covariate matrix aligned to ``covariates``."""
        x = np.asarray(covariates, dtype=np.float64)
        if x.ndim == 1:
            x = x.reshape(1, -1)
        if x.shape[1] != len(self.covariates):
            raise ValueError(
                f"covariate matrix has {x.shape[1]} columns, expected {len(self.covariates)}"
            )
        return np.asarray([self.root.predict(row) for row in x], dtype=np.int64)

    def assign(self) -> npt.NDArray[np.int64]:
        """The fitted policy's action for each unit it was learned on."""
        return self.predict(self._design)

    def interval(self, level: float = 0.9) -> tuple[float, float]:
        """A normal confidence interval for the policy value at ``level``."""
        from .._stats import _norm_ppf

        if not 0.0 < level < 1.0:
            raise ValueError(f"level must be in (0, 1), got {level}")
        z = _norm_ppf((1.0 + level) / 2.0)
        return self.value - z * self.std_error, self.value + z * self.std_error

    def regret(self, baseline: str = "treat_all") -> float:
        """The learned policy's welfare GAIN over a trivial baseline (the baseline's regret).

        ``baseline`` is ``"treat_all"``, ``"treat_none"``, or ``"best_trivial"`` (the
        better of the two). Returns ``value - baseline_value``: how much welfare the
        learned rule adds over treating everyone / no one. The trivial constant
        policies are depth-0 leaves inside the tree's class, so the optimal tree is at
        least as good and this gain is non-negative up to sampling noise (it is the
        baseline's regret against the learned optimum, not the optimum's own regret,
        which would need the unknown oracle value).
        """
        if baseline == "treat_all":
            ref = self.treat_all_value
        elif baseline == "treat_none":
            ref = self.treat_none_value
        elif baseline == "best_trivial":
            ref = max(self.treat_all_value, self.treat_none_value)
        else:
            raise ValueError(
                f"baseline must be 'treat_all', 'treat_none', or 'best_trivial', got {baseline!r}"
            )
        return self.value - ref

    def leaves(self) -> list[PolicyNode]:
        """The leaf nodes, left to right (each carries its action, count, and value)."""
        out: list[PolicyNode] = []

        def walk(node: PolicyNode) -> None:
            if node.is_leaf:
                out.append(node)
                return
            assert node.left is not None and node.right is not None
            walk(node.left)
            walk(node.right)

        walk(self.root)
        return out

    def describe(self) -> str:
        """A readable ``if-else`` rendering of the tree (the interpretable rule)."""
        lines: list[str] = []

        def walk(node: PolicyNode, indent: str) -> None:
            if node.is_leaf:
                lines.append(f"{indent}-> treat={node.action} (n={node.n})")
                return
            feat = node.feature
            thr = node.threshold
            assert thr is not None
            lines.append(f"{indent}if {feat} <= {thr:.4g}:")
            assert node.left is not None and node.right is not None
            walk(node.left, indent + "  ")
            lines.append(f"{indent}else:  # {feat} > {thr:.4g}")
            walk(node.right, indent + "  ")

        walk(self.root, "")
        return "\n".join(lines)

assign()

The fitted policy's action for each unit it was learned on.

Source code in src/ergodic/inference/explain/_policy.py

def assign(self) -> npt.NDArray[np.int64]:
    """The fitted policy's action for each unit it was learned on."""
    return self.predict(self._design)

describe()

A readable if-else rendering of the tree (the interpretable rule).

Source code in src/ergodic/inference/explain/_policy.py

def describe(self) -> str:
    """A readable ``if-else`` rendering of the tree (the interpretable rule)."""
    lines: list[str] = []

    def walk(node: PolicyNode, indent: str) -> None:
        if node.is_leaf:
            lines.append(f"{indent}-> treat={node.action} (n={node.n})")
            return
        feat = node.feature
        thr = node.threshold
        assert thr is not None
        lines.append(f"{indent}if {feat} <= {thr:.4g}:")
        assert node.left is not None and node.right is not None
        walk(node.left, indent + "  ")
        lines.append(f"{indent}else:  # {feat} > {thr:.4g}")
        walk(node.right, indent + "  ")

    walk(self.root, "")
    return "\n".join(lines)

interval(level=0.9)

A normal confidence interval for the policy value at level.

Source code in src/ergodic/inference/explain/_policy.py

def interval(self, level: float = 0.9) -> tuple[float, float]:
    """A normal confidence interval for the policy value at ``level``."""
    from .._stats import _norm_ppf

    if not 0.0 < level < 1.0:
        raise ValueError(f"level must be in (0, 1), got {level}")
    z = _norm_ppf((1.0 + level) / 2.0)
    return self.value - z * self.std_error, self.value + z * self.std_error

leaves()

The leaf nodes, left to right (each carries its action, count, and value).

Source code in src/ergodic/inference/explain/_policy.py

def leaves(self) -> list[PolicyNode]:
    """The leaf nodes, left to right (each carries its action, count, and value)."""
    out: list[PolicyNode] = []

    def walk(node: PolicyNode) -> None:
        if node.is_leaf:
            out.append(node)
            return
        assert node.left is not None and node.right is not None
        walk(node.left)
        walk(node.right)

    walk(self.root)
    return out

predict(covariates)

The per-row action (0/1) for a covariate matrix aligned to covariates.

Source code in src/ergodic/inference/explain/_policy.py

def predict(self, covariates: npt.NDArray[np.float64]) -> npt.NDArray[np.int64]:
    """The per-row action (0/1) for a covariate matrix aligned to ``covariates``."""
    x = np.asarray(covariates, dtype=np.float64)
    if x.ndim == 1:
        x = x.reshape(1, -1)
    if x.shape[1] != len(self.covariates):
        raise ValueError(
            f"covariate matrix has {x.shape[1]} columns, expected {len(self.covariates)}"
        )
    return np.asarray([self.root.predict(row) for row in x], dtype=np.int64)

regret(baseline='treat_all')

The learned policy's welfare GAIN over a trivial baseline (the baseline's regret).

baseline is "treat_all", "treat_none", or "best_trivial" (the better of the two). Returns value - baseline_value: how much welfare the learned rule adds over treating everyone / no one. The trivial constant policies are depth-0 leaves inside the tree's class, so the optimal tree is at least as good and this gain is non-negative up to sampling noise (it is the baseline's regret against the learned optimum, not the optimum's own regret, which would need the unknown oracle value).

Source code in src/ergodic/inference/explain/_policy.py

def regret(self, baseline: str = "treat_all") -> float:
    """The learned policy's welfare GAIN over a trivial baseline (the baseline's regret).

    ``baseline`` is ``"treat_all"``, ``"treat_none"``, or ``"best_trivial"`` (the
    better of the two). Returns ``value - baseline_value``: how much welfare the
    learned rule adds over treating everyone / no one. The trivial constant
    policies are depth-0 leaves inside the tree's class, so the optimal tree is at
    least as good and this gain is non-negative up to sampling noise (it is the
    baseline's regret against the learned optimum, not the optimum's own regret,
    which would need the unknown oracle value).
    """
    if baseline == "treat_all":
        ref = self.treat_all_value
    elif baseline == "treat_none":
        ref = self.treat_none_value
    elif baseline == "best_trivial":
        ref = max(self.treat_all_value, self.treat_none_value)
    else:
        raise ValueError(
            f"baseline must be 'treat_all', 'treat_none', or 'best_trivial', got {baseline!r}"
        )
    return self.value - ref

ergodic.select_cate(graph, data, treatment, outcome, *, candidates=None, learner=None, treated=None, random_state=None, n_boot=1000, folds=5, nu=0.5, q_iters=200)

Fit and rank candidate CATE learners by the doubly-robust tau-risk.

Identifies the effect (the same back-door gate as :func:~ergodic.inference.cate), fits each candidate in candidates (default the Phase-1 registered meta-learners dr_learner/r_learner/x_learner/t_learner) through the cate facade, and scores each against one shared cross-fit AIPW signal Gamma by the DR tau-risk mean[(Gamma - tau_hat)^2]. Returns a :class:SelectionReport with the ranked candidates (best risk first), the selected fitted :class:~ergodic.inference.CATEResult, and a Q-aggregation convex-ensemble result.

learner fills every candidate's nuisance slots; nu is the Q-aggregation mix (the penalty pulling weight toward lower-risk models). The shared Gamma is computed once with :func:~ergodic.inference.dr_scores, so the comparison is fair across candidates and each candidate's own out-of-fold tau_hat keeps the ranking honest.

Source code in src/ergodic/inference/_select.py

def select_cate(
    graph: MixedGraph,
    data: TabularData,
    treatment: str | Iterable[str],
    outcome: str | Iterable[str],
    *,
    candidates: Sequence[str] | None = None,
    learner: Any = None,
    treated: Any | None = None,
    random_state: int | None = None,
    n_boot: int = 1000,
    folds: int = 5,
    nu: float = 0.5,
    q_iters: int = 200,
) -> SelectionReport:
    """Fit and rank candidate CATE learners by the doubly-robust tau-risk.

    Identifies the effect (the same back-door gate as :func:`~ergodic.inference.cate`),
    fits each candidate in ``candidates`` (default the Phase-1 registered meta-learners
    ``dr_learner``/``r_learner``/``x_learner``/``t_learner``) through the ``cate``
    facade, and scores each against one shared cross-fit AIPW signal ``Gamma`` by the
    DR tau-risk ``mean[(Gamma - tau_hat)^2]``. Returns a :class:`SelectionReport` with
    the ranked candidates (best risk first), the selected fitted
    :class:`~ergodic.inference.CATEResult`, and a Q-aggregation convex-ensemble result.

    ``learner`` fills every candidate's nuisance slots; ``nu`` is the Q-aggregation
    mix (the penalty pulling weight toward lower-risk models). The shared ``Gamma`` is
    computed once with :func:`~ergodic.inference.dr_scores`, so the comparison is fair
    across candidates and each candidate's own out-of-fold ``tau_hat`` keeps the
    ranking honest.
    """
    from ..identification import identify_effect
    from ._estimate import cate
    from ._scores import dr_scores

    names = tuple(candidates) if candidates is not None else _DEFAULT_CANDIDATES
    if not names:
        raise ValueError("candidates is empty; pass at least one CATE method name")

    identification = identify_effect(graph, treatment, outcome)
    if not identification.identifiable:
        raise ValueError(
            f"effect of {treatment} on {outcome} is not identifiable: {identification.estimand}"
        )
    if identification.strategy != "adjustment":
        raise ValueError(
            "select_cate needs an adjustment-identified (back-door) effect; the effect "
            f"was identified by {identification.strategy!r}, which has no CATE meta-learner"
        )

    # One shared orthogonal signal: the comparison currency and the aggregation target.
    shared = dr_scores(
        identification,
        data,
        learners=learner,
        folds=folds,
        treated=treated,
        random_state=random_state,
    )
    gamma = np.asarray(shared.scores, dtype=np.float64)

    fitted: dict[str, CATEResult] = {}
    tau_hats: dict[str, npt.NDArray[np.float64]] = {}
    scored: list[CandidateScore] = []
    for method in names:
        result = cate(
            graph,
            data,
            treatment,
            outcome,
            method=method,
            learner=learner,
            treated=treated,
            random_state=random_state,
            n_boot=n_boot,
        )
        # the candidate's out-of-fold proxy ranking when cached (the honest choice),
        # else its in-sample fit. Scored against the shared Gamma.
        cached = result._scores
        if cached is not None and cached.tau_hat is not None:
            tau = np.asarray(cached.tau_hat, dtype=np.float64)
        else:
            tau = np.asarray(result.values, dtype=np.float64)
        fitted[method] = result
        tau_hats[method] = tau
        scored.append(CandidateScore(method=method, tau_risk=_tau_risk(gamma, tau)))

    ranking = tuple(sorted(scored, key=lambda s: (s.tau_risk, s.method)))
    selected_method = ranking[0].method
    selected = fitted[selected_method]

    # Q-aggregation: convex DR-loss ensemble with the Q penalty over the candidates.
    agg_methods = tuple(names)
    design = np.column_stack([tau_hats[m] for m in agg_methods])
    per_model_risk = np.asarray([_tau_risk(gamma, tau_hats[m]) for m in agg_methods])
    weights = _frank_wolfe_simplex(design, gamma, per_model_risk, nu=nu, iters=q_iters)
    agg_tau = np.asarray(design @ weights, dtype=np.float64)
    aggregate = _aggregate_result(selected, agg_tau, shared)
    return SelectionReport(
        ranking=ranking,
        selected_method=selected_method,
        selected=selected,
        aggregate=aggregate,
        aggregate_methods=agg_methods,
        weights=weights,
        aggregate_tau_risk=_tau_risk(gamma, agg_tau),
    )

ergodic.SelectionReport dataclass

The ranked CATE model-selection result with the chosen and aggregated models.

ranking is the candidates sorted by ascending DR tau-risk (best first). selected is the fitted :class:~ergodic.inference.CATEResult of the lowest risk; selected_method names it. aggregate is the Q-aggregation ensemble result and weights its convex weights (aligned to aggregate_methods), with aggregate_tau_risk the ensemble's DR tau-risk.

Source code in src/ergodic/inference/_select.py

@dataclass(frozen=True)
class SelectionReport:
    """The ranked CATE model-selection result with the chosen and aggregated models.

    ``ranking`` is the candidates sorted by ascending DR tau-risk (best first).
    ``selected`` is the fitted :class:`~ergodic.inference.CATEResult` of the lowest
    risk; ``selected_method`` names it. ``aggregate`` is the Q-aggregation ensemble
    result and ``weights`` its convex weights (aligned to ``aggregate_methods``),
    with ``aggregate_tau_risk`` the ensemble's DR tau-risk.
    """

    ranking: tuple[CandidateScore, ...]
    selected_method: str
    selected: CATEResult
    aggregate: CATEResult
    aggregate_methods: tuple[str, ...]
    weights: npt.NDArray[np.float64] = field(default_factory=lambda: np.empty(0))
    aggregate_tau_risk: float = float("nan")

    def risk(self, method: str) -> float:
        """The DR tau-risk of a candidate by name (raises when it was not scored)."""
        for score in self.ranking:
            if score.method == method:
                return score.tau_risk
        raise KeyError(f"{method!r} is not among the scored candidates")

    @property
    def best_single_tau_risk(self) -> float:
        """The DR tau-risk of the best single candidate (the first in the ranking)."""
        return self.ranking[0].tau_risk

    def frame(self) -> dict[str, float]:
        """A method-to-risk mapping (the aggregate included), handy for tables/tests."""
        out = {score.method: score.tau_risk for score in self.ranking}
        out["q_aggregate"] = self.aggregate_tau_risk
        return out

best_single_tau_risk property

The DR tau-risk of the best single candidate (the first in the ranking).

frame()

A method-to-risk mapping (the aggregate included), handy for tables/tests.

Source code in src/ergodic/inference/_select.py

def frame(self) -> dict[str, float]:
    """A method-to-risk mapping (the aggregate included), handy for tables/tests."""
    out = {score.method: score.tau_risk for score in self.ranking}
    out["q_aggregate"] = self.aggregate_tau_risk
    return out

risk(method)

The DR tau-risk of a candidate by name (raises when it was not scored).

Source code in src/ergodic/inference/_select.py

def risk(self, method: str) -> float:
    """The DR tau-risk of a candidate by name (raises when it was not scored)."""
    for score in self.ranking:
        if score.method == method:
            return score.tau_risk
    raise KeyError(f"{method!r} is not among the scored candidates")

Process intelligence

Event-log mining: the map, the causal case table, the model layer, and decision-point analysis.

ergodic.process

Process mining and intelligence: from event logs to causal questions.

The descriptive layer recovers the daily workflow over an :class:~ergodic.data.EventLog: :func:process_map builds the directly-follows map as a value object (frequencies, gap statistics, and a simplify that records what it dropped), :func:performance reports cycle times and, when a lifecycle column exists, the waiting and service split, :func:variant_table ranks the trace variants, and the filter vocabulary (:func:filter_time, :func:filter_variants, :func:filter_endpoints, :func:filter_activities, :func:filter_follows, :func:filter_performance) composes, with each filter saying whether it drops events or cases.

The causal layer is the point. :func:case_table encodes the log into one row per case with timestamp-gated features, the named process confounders (workload, calendar), and an auto-built :class:~ergodic.knowledge.DomainKnowledge whose tiers forbid post-treatment adjustment mechanically; its output feeds :func:~ergodic.discovery.discover, :func:~ergodic.identification.identify_effect, and :func:~ergodic.inference.estimate_effect unchanged. :func:kpi_panel aggregates the log into a :class:~ergodic.data.PanelData (or a single series), so a process change becomes a natural experiment for :func:~ergodic.inference.did, :func:~ergodic.inference.event_study, :func:~ergodic.inference.synthetic_control, and :func:~ergodic.inference.interrupted_time_series.

The model layer goes past the map. :func:discover_model is the inductive miner: it cuts the log into a block-structured :class:ProcessTree (build one by hand with :func:seq, :func:xor, :func:par, :func:loop, and :data:TAU) whose :meth:~ProcessTree.to_net gives a :class:WorkflowNet that is sound by construction, with :func:is_sound checking any net for real. :func:conformance replays the log against a model, token replay for the fast fitness and optimal alignments for the per-variant deviations. And :func:decision_points with :func:decision_table make the model causal: an exclusive split is a decision, the branch taken is a treatment, the features known at the decision moment are the honest adjusters, and positivity ships as a first-class report instead of a footnote.

END = '__end__' module-attribute

Target of the synthetic arc out of each case's last activity.

START = '__start__' module-attribute

Source of the synthetic arc into each case's first activity.

TAU = ProcessTree() module-attribute

The silent step: a leaf that takes no time and leaves no event.

ActivityStats dataclass

Per-activity timing: event count, mean wait before it, mean service inside it.

count is how many events carry the activity. mean_wait is None when no event of the activity has anything earlier in its case to wait on; mean_service is None on complete-only logs, where one timestamp per activity cannot split waiting from working, and for activities with no matched start/complete pair.

Source code in src/ergodic/process/_performance.py

@dataclass(frozen=True)
class ActivityStats:
    """Per-activity timing: event count, mean wait before it, mean service inside it.

    ``count`` is how many events carry the activity. ``mean_wait`` is None
    when no event of the activity has anything earlier in its case to wait
    on; ``mean_service`` is None on complete-only logs, where one timestamp
    per activity cannot split waiting from working, and for activities with
    no matched start/complete pair.
    """

    activity: str
    count: int
    mean_wait: float | None
    mean_service: float | None

ArcStats dataclass

Traversal count and gap timing for one arc of the map.

mean_gap and median_gap are seconds between the two events. A START or END arc has no second event, so both are NaN there (the one shared math.nan object, which keeps identically built maps equal).

Source code in src/ergodic/process/_map.py

@dataclass(frozen=True)
class ArcStats:
    """Traversal count and gap timing for one arc of the map.

    ``mean_gap`` and ``median_gap`` are seconds between the two events. A
    START or END arc has no second event, so both are NaN there (the one
    shared ``math.nan`` object, which keeps identically built maps equal).
    """

    count: int
    mean_gap: float
    median_gap: float

CaseTable dataclass

An event log flattened to one row per case, ready for the causal stack.

data holds the covariates (sorted), the treatment, and the outcome, one row per kept case in case order. knowledge carries the tier order covariates, then treatment, then outcome. The drop counters are zero outside the landmark modes.

Source code in src/ergodic/process/_case_table.py

@dataclass(frozen=True)
class CaseTable:
    """An event log flattened to one row per case, ready for the causal stack.

    ``data`` holds the covariates (sorted), the treatment, and the outcome, one
    row per kept case in case order. ``knowledge`` carries the tier order
    covariates, then treatment, then outcome. The drop counters are zero
    outside the landmark modes.
    """

    data: TabularData
    knowledge: DomainKnowledge
    treatment: str
    outcome: str
    covariates: tuple[str, ...]
    when: str
    n_cases: int
    n_dropped_short: int = 0
    n_dropped_treated: int = 0

    def __repr__(self) -> str:
        dropped = ""
        if self.n_dropped_short or self.n_dropped_treated:
            dropped = (
                f", dropped_short={self.n_dropped_short}, dropped_treated={self.n_dropped_treated}"
            )
        return (
            f"CaseTable(cases={self.n_cases}, treatment={self.treatment!r}, "
            f"outcome={self.outcome!r}, covariates={len(self.covariates)}, "
            f"when={self.when!r}{dropped})"
        )

ConformanceResult dataclass

How well a log fits a model, by token replay and by alignments.

fitness is the token-replay log fitness in Rozinat's form: the missing/consumed and remaining/produced sums aggregate over all cases before the two halves average, so frequent variants weigh more. alignment_fitness is the case-weighted mean of the per-variant alignment fitness. A case is fitting when its variant aligns at cost zero; n_fitting and share_fitting count those. variants holds one :class:VariantConformance per distinct trace, most frequent first, ties by sequence.

Source code in src/ergodic/process/_conformance.py

@dataclass(frozen=True)
class ConformanceResult:
    """How well a log fits a model, by token replay and by alignments.

    ``fitness`` is the token-replay log fitness in Rozinat's form: the
    missing/consumed and remaining/produced sums aggregate over all cases
    before the two halves average, so frequent variants weigh more.
    ``alignment_fitness`` is the case-weighted mean of the per-variant
    alignment fitness. A case is fitting when its variant aligns at cost
    zero; ``n_fitting`` and ``share_fitting`` count those. ``variants``
    holds one :class:`VariantConformance` per distinct trace, most frequent
    first, ties by sequence.
    """

    fitness: float
    alignment_fitness: float
    n_cases: int
    n_variants: int
    n_fitting: int
    share_fitting: float
    variants: tuple[VariantConformance, ...]

    def __repr__(self) -> str:
        return (
            f"ConformanceResult(fitness={self.fitness:.3f}, "
            f"alignment_fitness={self.alignment_fitness:.3f}, "
            f"fitting={self.n_fitting}/{self.n_cases} cases)"
        )

DecisionPoint dataclass

One xor split of a process tree, with the branch each case took.

path locates the xor node in the :meth:ProcessTree.nodes convention. The branch tuples line up positionally: branches[j] names the j-th child (a single visible leaf by its label, any other child by its sorted alphabet joined with +, a silent child as TAU), alphabets[j] is its activity alphabet, and counts[j] and shares[j] say how many assigned cases took it. A case is assigned the branch owning its first event inside the split's combined alphabet; events in two or more branches make it ambiguous (n_ambiguous, excluded from the counts). A case touching no branch counts as the silent branch when the split has one and is always reached; otherwise it lands in n_skipped, because the log cannot tell a silent branch from a split never reached.

Source code in src/ergodic/process/_decisions.py

@dataclass(frozen=True)
class DecisionPoint:
    """One xor split of a process tree, with the branch each case took.

    ``path`` locates the xor node in the :meth:`ProcessTree.nodes` convention.
    The branch tuples line up positionally: ``branches[j]`` names the j-th
    child (a single visible leaf by its label, any other child by its sorted
    alphabet joined with ``+``, a silent child as ``TAU``), ``alphabets[j]``
    is its activity alphabet, and ``counts[j]`` and ``shares[j]`` say how many
    assigned cases took it. A case is assigned the branch owning its first
    event inside the split's combined alphabet; events in two or more branches
    make it ambiguous (``n_ambiguous``, excluded from the counts). A case
    touching no branch counts as the silent branch when the split has one and
    is always reached; otherwise it lands in ``n_skipped``, because the log
    cannot tell a silent branch from a split never reached.
    """

    path: tuple[int, ...]
    branches: tuple[str, ...]
    alphabets: tuple[tuple[str, ...], ...]
    counts: tuple[int, ...]
    shares: tuple[float, ...]
    n_cases: int
    n_skipped: int
    n_ambiguous: int

    def __repr__(self) -> str:
        body = ", ".join(
            f"{name}={count}" for name, count in zip(self.branches, self.counts, strict=True)
        )
        extra = ""
        if self.n_skipped or self.n_ambiguous:
            extra = f", skipped={self.n_skipped}, ambiguous={self.n_ambiguous}"
        return f"DecisionPoint(path={self.path}, {body}{extra})"

DecisionTable dataclass

A decision point encoded for the causal stack, one row per assigned case.

data holds the sorted covariates, the categorical branch treatment, and the outcome, one row per unambiguously assigned case in case order. knowledge carries the tier order covariates, then branch, then outcome. point is the underlying :class:DecisionPoint, positivity the coverage diagnostic, and n_dropped_ambiguous counts the cases excluded for touching two or more branches.

Source code in src/ergodic/process/_decisions.py

@dataclass(frozen=True)
class DecisionTable:
    """A decision point encoded for the causal stack, one row per assigned case.

    ``data`` holds the sorted covariates, the categorical ``branch`` treatment,
    and the outcome, one row per unambiguously assigned case in case order.
    ``knowledge`` carries the tier order covariates, then branch, then outcome.
    ``point`` is the underlying :class:`DecisionPoint`, ``positivity`` the
    coverage diagnostic, and ``n_dropped_ambiguous`` counts the cases excluded
    for touching two or more branches.
    """

    data: TabularData
    knowledge: DomainKnowledge
    outcome: str
    covariates: tuple[str, ...]
    point: DecisionPoint
    positivity: PositivityReport
    n_cases: int
    n_dropped_ambiguous: int
    treatment: str = "branch"

    def __repr__(self) -> str:
        dropped = ""
        if self.n_dropped_ambiguous:
            dropped = f", dropped_ambiguous={self.n_dropped_ambiguous}"
        flag = ", positivity flagged" if self.positivity.flagged else ""
        return (
            f"DecisionTable(cases={self.n_cases}, branches={self.point.branches}, "
            f"outcome={self.outcome!r}, covariates={len(self.covariates)}{dropped}{flag})"
        )

PositivityReport dataclass

Per-branch coverage of a decision table: the positivity diagnostic.

counts and shares are per branch among the included cases, and min_share is the smallest share. For a two-branch split, overlap_share is the share of cases whose in-sample propensity (a logistic regression of the branch on the covariate matrix) lies inside [0.05, 0.95]; it is None for wider splits, and 1.0 when there are no covariates, since the propensity is then the base rate for every case and the covariates cannot starve either branch.

flagged is True when min_share falls below 0.05 (one branch holds under five percent of the cases), any branch holds fewer than 10 cases (too few for a stable branch mean), or overlap_share exists and falls below 0.8 (for over a fifth of the cases the covariates nearly decide the branch, so the comparison there rests on extrapolation). A flag is a warning that the branches barely overlap, not a verdict that an estimate is wrong.

Source code in src/ergodic/process/_decisions.py

@dataclass(frozen=True)
class PositivityReport:
    """Per-branch coverage of a decision table: the positivity diagnostic.

    ``counts`` and ``shares`` are per branch among the included cases, and
    ``min_share`` is the smallest share. For a two-branch split,
    ``overlap_share`` is the share of cases whose in-sample propensity (a
    logistic regression of the branch on the covariate matrix) lies inside
    [0.05, 0.95]; it is None for wider splits, and 1.0 when there are no
    covariates, since the propensity is then the base rate for every case and
    the covariates cannot starve either branch.

    ``flagged`` is True when ``min_share`` falls below 0.05 (one branch holds
    under five percent of the cases), any branch holds fewer than 10 cases
    (too few for a stable branch mean), or ``overlap_share`` exists and falls
    below 0.8 (for over a fifth of the cases the covariates nearly decide the
    branch, so the comparison there rests on extrapolation). A flag is a
    warning that the branches barely overlap, not a verdict that an estimate
    is wrong.
    """

    branches: tuple[str, ...]
    counts: tuple[int, ...]
    shares: tuple[float, ...]
    min_share: float
    overlap_share: float | None
    flagged: bool

ProcessMap dataclass

A directly-follows map: activities, timed arcs, and filter provenance.

activities lists the real activities, sorted. arcs maps each (source, target) pair to its :class:ArcStats and iterates in sorted key order; it holds one (START, a) arc per first activity and one (b, END) arc per last, with per-case counts aggregated. A freshly built map has dropped_arcs and dropped_share at zero; they record what :meth:simplify has removed and accumulate over repeated calls.

Source code in src/ergodic/process/_map.py

@dataclass(frozen=True)
class ProcessMap:
    """A directly-follows map: activities, timed arcs, and filter provenance.

    ``activities`` lists the real activities, sorted. ``arcs`` maps each
    ``(source, target)`` pair to its :class:`ArcStats` and iterates in sorted
    key order; it holds one ``(START, a)`` arc per first activity and one
    ``(b, END)`` arc per last, with per-case counts aggregated. A freshly
    built map has ``dropped_arcs`` and ``dropped_share`` at zero; they record
    what :meth:`simplify` has removed and accumulate over repeated calls.
    """

    activities: tuple[str, ...]
    arcs: Mapping[tuple[str, str], ArcStats]
    n_cases: int
    n_events: int
    dropped_arcs: int = 0
    dropped_share: float = 0.0

    def frequency(self, a: str, b: str) -> int:
        """How many times the arc ``a -> b`` was traversed, 0 when it is absent."""
        stats = self.arcs.get((a, b))
        return 0 if stats is None else stats.count

    def simplify(self, keep: float = 0.8) -> ProcessMap:
        """A thinner map keeping the busiest arcs that cover ``keep`` of the traversals.

        ``keep`` must lie in (0, 1]. START and END arcs always stay, so every
        case still enters and leaves the map. Real arcs are ranked by count,
        heaviest first (ties broken by the sorted pair), and the smallest
        prefix whose counts reach ``keep`` times the total real traversals
        survives; the rest are dropped. Every activity stays, isolated ones
        included.

        Filtering is visible, never silent: ``dropped_arcs`` and
        ``dropped_share`` on the returned map say how many arcs are gone and
        what share of the original log's real traversals they carried, and
        both accumulate across repeated calls.
        """
        if not 0.0 < keep <= 1.0:
            raise ValueError(f"keep must be in (0, 1], got {keep}")
        real = [(pair, stats) for pair, stats in self.arcs.items() if not _is_boundary(pair)]
        total = sum(stats.count for _, stats in real)
        if total == 0:
            return self
        ranked = sorted(real, key=lambda item: (-item[1].count, item[0]))
        threshold = keep * total
        covered = 0
        n_keep = 0
        for _, stats in ranked:
            n_keep += 1
            covered += stats.count
            if covered >= threshold:
                break
        kept = {pair for pair, _ in ranked[:n_keep]}
        arcs = {
            pair: stats for pair, stats in self.arcs.items() if _is_boundary(pair) or pair in kept
        }
        # The current real traversals are (1 - dropped_share) of the original
        # log's, so the newly dropped share is measured against that original.
        original_total = total / (1.0 - self.dropped_share)
        return ProcessMap(
            activities=self.activities,
            arcs=arcs,
            n_cases=self.n_cases,
            n_events=self.n_events,
            dropped_arcs=self.dropped_arcs + (len(ranked) - n_keep),
            dropped_share=self.dropped_share + (total - covered) / original_total,
        )

    def to_dict(self) -> dict[str, Any]:
        """The map as plain data: the activities and arcs as lists, plus the counts."""
        return {
            "activities": list(self.activities),
            "arcs": [
                {
                    "source": pair[0],
                    "target": pair[1],
                    "count": stats.count,
                    "mean_gap": stats.mean_gap,
                    "median_gap": stats.median_gap,
                }
                for pair, stats in self.arcs.items()
            ],
            "n_cases": self.n_cases,
            "n_events": self.n_events,
            "dropped_arcs": self.dropped_arcs,
            "dropped_share": self.dropped_share,
        }

    def __repr__(self) -> str:
        note = ""
        if self.dropped_arcs:
            note = f", dropped {self.dropped_arcs} arcs ({self.dropped_share:.1%} of traversals)"
        return (
            f"ProcessMap(activities={len(self.activities)}, arcs={len(self.arcs)}, "
            f"cases={self.n_cases}{note})"
        )

frequency(a, b)

How many times the arc a -> b was traversed, 0 when it is absent.

Source code in src/ergodic/process/_map.py

def frequency(self, a: str, b: str) -> int:
    """How many times the arc ``a -> b`` was traversed, 0 when it is absent."""
    stats = self.arcs.get((a, b))
    return 0 if stats is None else stats.count

simplify(keep=0.8)

A thinner map keeping the busiest arcs that cover keep of the traversals.

keep must lie in (0, 1]. START and END arcs always stay, so every case still enters and leaves the map. Real arcs are ranked by count, heaviest first (ties broken by the sorted pair), and the smallest prefix whose counts reach keep times the total real traversals survives; the rest are dropped. Every activity stays, isolated ones included.

Filtering is visible, never silent: dropped_arcs and dropped_share on the returned map say how many arcs are gone and what share of the original log's real traversals they carried, and both accumulate across repeated calls.

Source code in src/ergodic/process/_map.py

def simplify(self, keep: float = 0.8) -> ProcessMap:
    """A thinner map keeping the busiest arcs that cover ``keep`` of the traversals.

    ``keep`` must lie in (0, 1]. START and END arcs always stay, so every
    case still enters and leaves the map. Real arcs are ranked by count,
    heaviest first (ties broken by the sorted pair), and the smallest
    prefix whose counts reach ``keep`` times the total real traversals
    survives; the rest are dropped. Every activity stays, isolated ones
    included.

    Filtering is visible, never silent: ``dropped_arcs`` and
    ``dropped_share`` on the returned map say how many arcs are gone and
    what share of the original log's real traversals they carried, and
    both accumulate across repeated calls.
    """
    if not 0.0 < keep <= 1.0:
        raise ValueError(f"keep must be in (0, 1], got {keep}")
    real = [(pair, stats) for pair, stats in self.arcs.items() if not _is_boundary(pair)]
    total = sum(stats.count for _, stats in real)
    if total == 0:
        return self
    ranked = sorted(real, key=lambda item: (-item[1].count, item[0]))
    threshold = keep * total
    covered = 0
    n_keep = 0
    for _, stats in ranked:
        n_keep += 1
        covered += stats.count
        if covered >= threshold:
            break
    kept = {pair for pair, _ in ranked[:n_keep]}
    arcs = {
        pair: stats for pair, stats in self.arcs.items() if _is_boundary(pair) or pair in kept
    }
    # The current real traversals are (1 - dropped_share) of the original
    # log's, so the newly dropped share is measured against that original.
    original_total = total / (1.0 - self.dropped_share)
    return ProcessMap(
        activities=self.activities,
        arcs=arcs,
        n_cases=self.n_cases,
        n_events=self.n_events,
        dropped_arcs=self.dropped_arcs + (len(ranked) - n_keep),
        dropped_share=self.dropped_share + (total - covered) / original_total,
    )

to_dict()

The map as plain data: the activities and arcs as lists, plus the counts.

Source code in src/ergodic/process/_map.py

def to_dict(self) -> dict[str, Any]:
    """The map as plain data: the activities and arcs as lists, plus the counts."""
    return {
        "activities": list(self.activities),
        "arcs": [
            {
                "source": pair[0],
                "target": pair[1],
                "count": stats.count,
                "mean_gap": stats.mean_gap,
                "median_gap": stats.median_gap,
            }
            for pair, stats in self.arcs.items()
        ],
        "n_cases": self.n_cases,
        "n_events": self.n_events,
        "dropped_arcs": self.dropped_arcs,
        "dropped_share": self.dropped_share,
    }

ProcessModel dataclass

A mined process model and the provenance of everything removed.

tree is the block-structured model and net is its workflow net, tree.to_net(), sound by construction. The rest records what the miner filtered or dropped on the way (filtering is visible, never silent):

• fallthroughs: paths, in :meth:ProcessTree.nodes convention (child indices from the root), of the subtrees where no cut was valid and the miner fell through to the flower model, which accepts its activities in any order. Empty when every subtree came from a real cut.
• n_filtered_arcs: directly-follows edges removed by the infrequent filter, summed over all recursions. Zero when noise is 0.
• n_dropped_events: events discarded while projecting traces through an exclusive-choice cut, where each trace goes whole to one branch and its events outside that branch are dropped.
• dropped_share: n_dropped_events over the log's total events.
• noise: the filter level the model was mined with.

Source code in src/ergodic/process/_mine.py

@dataclass(frozen=True)
class ProcessModel:
    """A mined process model and the provenance of everything removed.

    ``tree`` is the block-structured model and ``net`` is its workflow net,
    ``tree.to_net()``, sound by construction. The rest records what the miner
    filtered or dropped on the way (filtering is visible, never silent):

    - ``fallthroughs``: paths, in :meth:`ProcessTree.nodes` convention (child
      indices from the root), of the subtrees where no cut was valid and the
      miner fell through to the flower model, which accepts its activities in
      any order. Empty when every subtree came from a real cut.
    - ``n_filtered_arcs``: directly-follows edges removed by the infrequent
      filter, summed over all recursions. Zero when ``noise`` is 0.
    - ``n_dropped_events``: events discarded while projecting traces through
      an exclusive-choice cut, where each trace goes whole to one branch and
      its events outside that branch are dropped.
    - ``dropped_share``: ``n_dropped_events`` over the log's total events.
    - ``noise``: the filter level the model was mined with.
    """

    tree: ProcessTree
    net: WorkflowNet
    fallthroughs: tuple[tuple[int, ...], ...]
    n_filtered_arcs: int
    n_dropped_events: int
    dropped_share: float
    noise: float

    def __repr__(self) -> str:
        extras = []
        if self.fallthroughs:
            extras.append(f"flowers={list(self.fallthroughs)}")
        if self.n_filtered_arcs:
            extras.append(f"filtered_arcs={self.n_filtered_arcs}")
        if self.n_dropped_events:
            extras.append(f"dropped_events={self.n_dropped_events} ({self.dropped_share:.1%})")
        note = f", {', '.join(extras)}" if extras else ""
        return f"ProcessModel(tree={self.tree!r}, noise={self.noise}{note})"

ProcessPerformance dataclass

Case cycle times plus the per-activity timing table.

Durations are each case's last event minus its first, in seconds. activities is sorted by activity name. lifecycle says whether a waiting/service split was possible, and n_unmatched counts lifecycle rows that found no partner.

Source code in src/ergodic/process/_performance.py

@dataclass(frozen=True)
class ProcessPerformance:
    """Case cycle times plus the per-activity timing table.

    Durations are each case's last event minus its first, in seconds.
    ``activities`` is sorted by activity name. ``lifecycle`` says whether a
    waiting/service split was possible, and ``n_unmatched`` counts lifecycle
    rows that found no partner.
    """

    n_cases: int
    n_events: int
    mean_duration: float
    median_duration: float
    p90_duration: float
    activities: tuple[ActivityStats, ...]
    lifecycle: bool
    n_unmatched: int = 0

ProcessTree dataclass

One node of a process tree.

A leaf has no children and carries an activity label (or None, which is the silent step :data:TAU). An operator node has op set, no label, and at least two children; loop takes exactly two, the do body and the redo body. The children of xor and par are stored sorted by their activity alphabets (a silent child first), so equal models are equal objects regardless of how they were assembled.

Source code in src/ergodic/process/_tree.py

@dataclass(frozen=True)
class ProcessTree:
    """One node of a process tree.

    A leaf has no children and carries an activity ``label`` (or ``None``,
    which is the silent step :data:`TAU`). An operator node has ``op`` set,
    no label, and at least two children; ``loop`` takes exactly two, the do
    body and the redo body. The children of ``xor`` and ``par`` are stored
    sorted by their activity alphabets (a silent child first), so equal models
    are equal objects regardless of how they were assembled.
    """

    op: str | None = None
    label: str | None = None
    children: tuple[ProcessTree, ...] = field(default=())

    def __post_init__(self) -> None:
        if self.op is None:
            if self.children:
                raise ValueError("a leaf has no children")
            return
        if self.op not in _OPERATORS:
            raise ValueError(f"unknown operator {self.op!r}; expected one of {_OPERATORS}")
        if self.label is not None:
            raise ValueError("an operator node has no label")
        if self.op == "loop":
            if len(self.children) != 2:
                raise ValueError("loop takes exactly two children: the do body and the redo body")
        elif len(self.children) < 2:
            raise ValueError(f"{self.op} takes at least two children")
        if self.op in ("xor", "par"):
            ordered = tuple(sorted(self.children, key=lambda child: child.activities()))
            object.__setattr__(self, "children", ordered)

    @property
    def is_leaf(self) -> bool:
        """Whether this node is a leaf (an activity or the silent step)."""
        return self.op is None

    @property
    def is_tau(self) -> bool:
        """Whether this node is the silent step."""
        return self.op is None and self.label is None

    def activities(self) -> tuple[str, ...]:
        """The activity labels in this subtree, sorted and unique."""
        if self.op is None:
            return () if self.label is None else (self.label,)
        seen: set[str] = set()
        for child in self.children:
            seen.update(child.activities())
        return tuple(sorted(seen))

    def nodes(self) -> tuple[tuple[tuple[int, ...], ProcessTree], ...]:
        """Every node with its path from the root, depth first.

        A path is the tuple of child indices leading to the node; the root's
        path is ``()``.
        """
        out: list[tuple[tuple[int, ...], ProcessTree]] = []

        def walk(node: ProcessTree, path: tuple[int, ...]) -> None:
            out.append((path, node))
            for i, child in enumerate(node.children):
                walk(child, (*path, i))

        walk(self, ())
        return tuple(out)

    def subtree(self, path: tuple[int, ...]) -> ProcessTree:
        """The node at ``path``, where each entry picks a child by index."""
        node = self
        for depth, i in enumerate(path):
            if not 0 <= i < len(node.children):
                raise ValueError(f"no child {i} at path {path[:depth]} of {node!r}")
            node = node.children[i]
        return node

    def to_net(self) -> WorkflowNet:
        """The workflow net of this tree, sound by construction."""
        from ._net import net_of_tree

        return net_of_tree(self)

    def __repr__(self) -> str:
        if self.op is None:
            return "TAU" if self.label is None else repr(self.label)
        return f"{self.op}({', '.join(repr(child) for child in self.children)})"

is_leaf property

Whether this node is a leaf (an activity or the silent step).

is_tau property

Whether this node is the silent step.

activities()

The activity labels in this subtree, sorted and unique.

Source code in src/ergodic/process/_tree.py

def activities(self) -> tuple[str, ...]:
    """The activity labels in this subtree, sorted and unique."""
    if self.op is None:
        return () if self.label is None else (self.label,)
    seen: set[str] = set()
    for child in self.children:
        seen.update(child.activities())
    return tuple(sorted(seen))

nodes()

Every node with its path from the root, depth first.

A path is the tuple of child indices leading to the node; the root's path is ().

Source code in src/ergodic/process/_tree.py

def nodes(self) -> tuple[tuple[tuple[int, ...], ProcessTree], ...]:
    """Every node with its path from the root, depth first.

    A path is the tuple of child indices leading to the node; the root's
    path is ``()``.
    """
    out: list[tuple[tuple[int, ...], ProcessTree]] = []

    def walk(node: ProcessTree, path: tuple[int, ...]) -> None:
        out.append((path, node))
        for i, child in enumerate(node.children):
            walk(child, (*path, i))

    walk(self, ())
    return tuple(out)

subtree(path)

The node at path, where each entry picks a child by index.

Source code in src/ergodic/process/_tree.py

def subtree(self, path: tuple[int, ...]) -> ProcessTree:
    """The node at ``path``, where each entry picks a child by index."""
    node = self
    for depth, i in enumerate(path):
        if not 0 <= i < len(node.children):
            raise ValueError(f"no child {i} at path {path[:depth]} of {node!r}")
        node = node.children[i]
    return node

to_net()

The workflow net of this tree, sound by construction.

Source code in src/ergodic/process/_tree.py

def to_net(self) -> WorkflowNet:
    """The workflow net of this tree, sound by construction."""
    from ._net import net_of_tree

    return net_of_tree(self)

Variant dataclass

One distinct activity sequence and its case statistics.

Source code in src/ergodic/process/_variants.py

@dataclass(frozen=True)
class Variant:
    """One distinct activity sequence and its case statistics."""

    sequence: tuple[str, ...]
    n_cases: int
    share: float
    mean_duration: float
    median_duration: float

VariantConformance dataclass

Replay and alignment results for one distinct trace.

sequence is the trace and n_cases how many cases follow it. The token-replay counters come next: missing tokens had to be added, remaining were left behind, and consumed and produced count every token the replay moved, the initial token and the final take from the sink included. replay_fitness is Rozinat's half-and-half of the two ratios. cost is the optimal alignment cost, worst_cost the reference cost of deleting the whole trace and running the model's cheapest visible path instead, and alignment_fitness is 1 - cost / worst_cost. alignment lists the moves: ("sync", a) where trace and model agree, ("log", a) for an event the model cannot mirror, ("model", a) for a required visible model step the trace skipped. Silent model steps cost nothing and are omitted from the list.

Source code in src/ergodic/process/_conformance.py

@dataclass(frozen=True)
class VariantConformance:
    """Replay and alignment results for one distinct trace.

    ``sequence`` is the trace and ``n_cases`` how many cases follow it. The
    token-replay counters come next: ``missing`` tokens had to be added,
    ``remaining`` were left behind, and ``consumed`` and ``produced`` count
    every token the replay moved, the initial token and the final take from
    the sink included. ``replay_fitness`` is Rozinat's half-and-half of the
    two ratios. ``cost`` is the optimal alignment cost, ``worst_cost`` the
    reference cost of deleting the whole trace and running the model's
    cheapest visible path instead, and ``alignment_fitness`` is
    ``1 - cost / worst_cost``. ``alignment`` lists the moves: ``("sync", a)``
    where trace and model agree, ``("log", a)`` for an event the model cannot
    mirror, ``("model", a)`` for a required visible model step the trace
    skipped. Silent model steps cost nothing and are omitted from the list.
    """

    sequence: tuple[str, ...]
    n_cases: int
    replay_fitness: float
    missing: int
    remaining: int
    consumed: int
    produced: int
    cost: int
    worst_cost: int
    alignment_fitness: float
    alignment: tuple[tuple[str, str], ...]

WorkflowNet dataclass

A workflow net: places, labeled or silent transitions, and arcs.

transitions maps ids to labels as (id, label) pairs, with None for a silent transition; ids and place names live in one namespace and must not collide. Arcs connect a place to a transition or a transition to a place, once each (these are weight-one nets). The constructor sorts the fields into canonical order and validates the workflow-net shape: a single source with no inputs, a single sink with no outputs, and every node on some directed path from source to sink.

Source code in src/ergodic/process/_net.py

@dataclass(frozen=True)
class WorkflowNet:
    """A workflow net: places, labeled or silent transitions, and arcs.

    ``transitions`` maps ids to labels as ``(id, label)`` pairs, with ``None``
    for a silent transition; ids and place names live in one namespace and
    must not collide. Arcs connect a place to a transition or a transition to
    a place, once each (these are weight-one nets). The constructor sorts the
    fields into canonical order and validates the workflow-net shape: a single
    ``source`` with no inputs, a single ``sink`` with no outputs, and every
    node on some directed path from source to sink.
    """

    places: tuple[str, ...]
    transitions: tuple[tuple[str, str | None], ...]
    arcs: tuple[tuple[str, str], ...]
    source: str
    sink: str

    def __post_init__(self) -> None:
        object.__setattr__(self, "places", tuple(sorted(self.places)))
        object.__setattr__(self, "transitions", tuple(sorted(self.transitions)))
        object.__setattr__(self, "arcs", tuple(sorted(self.arcs)))
        place_set = set(self.places)
        tids = [tid for tid, _ in self.transitions]
        tid_set = set(tids)
        if len(place_set) != len(self.places):
            raise ValueError("duplicate place names")
        if len(tid_set) != len(tids):
            raise ValueError("duplicate transition ids")
        shared = sorted(place_set & tid_set)
        if shared:
            raise ValueError(f"names used as both place and transition: {shared}")
        for name, role in ((self.source, "source"), (self.sink, "sink")):
            if name not in place_set:
                raise ValueError(f"{role} {name!r} is not a place")
        if self.source == self.sink:
            raise ValueError("source and sink must be different places")
        if len(set(self.arcs)) != len(self.arcs):
            raise ValueError("duplicate arcs (these are weight-one nets)")
        for a, b in self.arcs:
            ok = (a in place_set and b in tid_set) or (a in tid_set and b in place_set)
            if not ok:
                raise ValueError(
                    f"arc ({a!r}, {b!r}) must connect a place and a transition, "
                    "and both ends must exist"
                )
            if b == self.source:
                raise ValueError(f"source {self.source!r} cannot have incoming arcs")
            if a == self.sink:
                raise ValueError(f"sink {self.sink!r} cannot have outgoing arcs")
        forward: dict[str, list[str]] = {}
        backward: dict[str, list[str]] = {}
        for a, b in self.arcs:
            forward.setdefault(a, []).append(b)
            backward.setdefault(b, []).append(a)
        from_source = _reach(self.source, forward)
        to_sink = _reach(self.sink, backward)
        stranded = sorted((place_set | tid_set) - (from_source & to_sink))
        if stranded:
            raise ValueError(
                f"every node must lie on a path from source to sink; these do not: {stranded}"
            )

    @property
    def initial_marking(self) -> Marking:
        """One token on the source place."""
        return ((self.source, 1),)

    @property
    def final_marking(self) -> Marking:
        """One token on the sink place."""
        return ((self.sink, 1),)

    def index(self) -> NetIndex:
        """The pre/post maps for the firing rule; build once per algorithm run."""
        pre: dict[str, tuple[str, ...]] = {}
        post: dict[str, tuple[str, ...]] = {}
        labels = dict(self.transitions)
        pre_lists: dict[str, list[str]] = {tid: [] for tid in labels}
        post_lists: dict[str, list[str]] = {tid: [] for tid in labels}
        for a, b in self.arcs:
            if b in labels:
                pre_lists[b].append(a)
            else:
                post_lists[a].append(b)
        for tid in labels:
            pre[tid] = tuple(sorted(pre_lists[tid]))
            post[tid] = tuple(sorted(post_lists[tid]))
        return NetIndex(pre=pre, post=post, labels=labels)

    def __repr__(self) -> str:
        n_silent = sum(1 for _, label in self.transitions if label is None)
        return (
            f"WorkflowNet(places={len(self.places)}, transitions={len(self.transitions)} "
            f"({n_silent} silent), arcs={len(self.arcs)})"
        )

final_marking property

One token on the sink place.

initial_marking property

One token on the source place.

index()

The pre/post maps for the firing rule; build once per algorithm run.

Source code in src/ergodic/process/_net.py

def index(self) -> NetIndex:
    """The pre/post maps for the firing rule; build once per algorithm run."""
    pre: dict[str, tuple[str, ...]] = {}
    post: dict[str, tuple[str, ...]] = {}
    labels = dict(self.transitions)
    pre_lists: dict[str, list[str]] = {tid: [] for tid in labels}
    post_lists: dict[str, list[str]] = {tid: [] for tid in labels}
    for a, b in self.arcs:
        if b in labels:
            pre_lists[b].append(a)
        else:
            post_lists[a].append(b)
    for tid in labels:
        pre[tid] = tuple(sorted(pre_lists[tid]))
        post[tid] = tuple(sorted(post_lists[tid]))
    return NetIndex(pre=pre, post=post, labels=labels)

case_table(log, treatment, outcome='cycle_time', *, covariates='auto', when='anytime')

Encode an event log as one row per case: covariates, treatment, outcome.

Parameters:

Name	Type	Description	Default
log	EventLog	The event log to flatten.	required
treatment	str	An activity name (the treatment is whether it occurred, per when) or a case-constant 0/1 column (when="anytime" only, since a column carries no timing). A name that is both raises.	required
outcome	str	"cycle_time" (case duration in seconds), "n_events" (events per case), or a case-constant column. The two derived names are reserved: a log column by either name raises.	'cycle_time'
covariates	str | Iterable[str]	"auto" takes every case-constant variable column plus the mode's derived features; an iterable names columns (validated case-constant) and/or derived features. A derived covariate that is constant across the kept cases is dropped in silence.	'auto'
when	str | tuple[str, int | float]	"anytime", ("after_events", k) with k >= 1, or ("after_elapsed", x) with x > 0 seconds. The landmark modes keep only cases that reached the landmark untreated, so their estimand is the effect in that population.	'anytime'

Returns:

Name	Type	Description
A	CaseTable	class:CaseTable: the :class:~ergodic.data.TabularData (columns
	CaseTable	in the order sorted covariates, treatment, outcome; one row per kept
	CaseTable	case, in case order), the tiered :class:~ergodic.DomainKnowledge,
	CaseTable	and the drop counters of the landmark modes.

Source code in src/ergodic/process/_case_table.py

def case_table(
    log: EventLog,
    treatment: str,
    outcome: str = "cycle_time",
    *,
    covariates: str | Iterable[str] = "auto",
    when: str | tuple[str, int | float] = "anytime",
) -> CaseTable:
    """Encode an event log as one row per case: covariates, treatment, outcome.

    Args:
        log: The event log to flatten.
        treatment: An activity name (the treatment is whether it occurred, per
            ``when``) or a case-constant 0/1 column (``when="anytime"`` only,
            since a column carries no timing). A name that is both raises.
        outcome: ``"cycle_time"`` (case duration in seconds), ``"n_events"``
            (events per case), or a case-constant column. The two derived
            names are reserved: a log column by either name raises.
        covariates: ``"auto"`` takes every case-constant variable column plus
            the mode's derived features; an iterable names columns (validated
            case-constant) and/or derived features. A derived covariate that
            is constant across the kept cases is dropped in silence.
        when: ``"anytime"``, ``("after_events", k)`` with ``k >= 1``, or
            ``("after_elapsed", x)`` with ``x > 0`` seconds. The landmark
            modes keep only cases that reached the landmark untreated, so
            their estimand is the effect in that population.

    Returns:
        A :class:`CaseTable`: the :class:`~ergodic.data.TabularData` (columns
        in the order sorted covariates, treatment, outcome; one row per kept
        case, in case order), the tiered :class:`~ergodic.DomainKnowledge`,
        and the drop counters of the landmark modes.
    """
    mode, parameter = _parse_when(when)
    if treatment == outcome:
        raise ValueError(f"treatment and outcome must differ; both are {treatment!r}")
    _, starts = case_slices(log)
    _, first, last = case_bounds(log)
    n_all = int(first.shape[0])
    if n_all == 0:
        raise ValueError("the event log has no cases")
    stamps = time_seconds(log.column(log.timestamp))
    outcome_all, outcome_type = _outcome_values(log, outcome, first, last, starts)
    source = _treatment_source(log, treatment)

    derived: dict[str, npt.NDArray[np.float64]] = {
        "workload_at_start": open_cases(first, last, first)
    }
    if log.schema[log.timestamp].is_datetime:
        days = np.floor(first / _SECONDS_PER_DAY)
        derived["start_weekday"] = np.mod(days + 3.0, 7.0)  # the epoch was a Thursday
        derived["start_hour"] = np.floor((first - days * _SECONDS_PER_DAY) / _SECONDS_PER_HOUR)

    keep = np.ones(n_all, dtype=bool)
    n_dropped_short = 0
    n_dropped_treated = 0
    if source == "column":
        if mode != "anytime":
            raise ValueError(
                f"treatment column {treatment!r} carries no event timing, so a landmark "
                "mode cannot tell whether it preceded the landmark; use when='anytime' "
                "or name a treatment activity"
            )
        treated = _binary_treatment(case_attribute(log, treatment), treatment)
    else:
        hit = np.asarray(log.column(log.activity) == treatment, dtype=bool)
        if mode == "anytime":
            treated = np.logical_or.reduceat(hit, starts[:-1]).astype(np.int64)
        elif mode == "after_events":
            assert isinstance(parameter, int)  # narrowed by _parse_when
            keep, treated, extra, n_dropped_short, n_dropped_treated = _after_events(
                log, hit, starts, stamps, first, last, parameter
            )
            derived.update(extra)
        else:
            assert isinstance(parameter, float)  # narrowed by _parse_when
            keep, treated, extra, n_dropped_short, n_dropped_treated = _after_elapsed(
                log, hit, starts, stamps, first, last, parameter
            )
            derived.update(extra)

    clash = sorted(set(derived) & set(log.schema.names))
    if clash:
        raise ValueError(
            f"columns {clash} collide with the derived feature names of when={mode!r}; "
            "rename the columns so the covariates are unambiguous"
        )

    if isinstance(covariates, str) and covariates == "auto":
        attributes = _auto_attributes(log, treatment, outcome)
        chosen_derived = dict(derived)
    else:
        requested = [covariates] if isinstance(covariates, str) else list(covariates)
        attributes, chosen_derived = _named_covariates(
            log, requested, derived, treatment, outcome, mode
        )

    if not bool(keep.any()):
        raise ValueError(
            f"no cases remain at the landmark: {n_dropped_short} never reached it "
            f"and {n_dropped_treated} were already treated"
        )

    columns: dict[str, npt.NDArray[Any]] = {}
    types: dict[str, str | Ordinal] = {}
    for name in attributes:
        columns[name] = attributes[name][keep]
        types[name] = _attribute_type(log.schema[name])
    for name, values in chosen_derived.items():
        kept_values = values[keep]
        if bool(np.all(kept_values == kept_values[0])):
            continue  # constant across the kept cases: useless and harmful downstream
        columns[name] = kept_values
        types[name] = "continuous"

    covariate_names = sorted(columns)
    ordered: dict[str, npt.NDArray[Any]] = {name: columns[name] for name in covariate_names}
    ordered[treatment] = treated[keep]
    types[treatment] = "categorical"
    ordered[outcome] = outcome_all[keep]
    types[outcome] = outcome_type
    data = tabular(ordered, types=types)

    tiers: list[list[str]] = []
    if covariate_names:
        tiers.append(list(covariate_names))
    tiers.extend(([treatment], [outcome]))
    knowledge = DomainKnowledge().with_tiers(tiers)

    return CaseTable(
        data=data,
        knowledge=knowledge,
        treatment=treatment,
        outcome=outcome,
        covariates=tuple(covariate_names),
        when=mode,
        n_cases=int(keep.sum()),
        n_dropped_short=n_dropped_short,
        n_dropped_treated=n_dropped_treated,
    )

conformance(log, model, *, max_states=100000)

Check log against model by token replay and by optimal alignments.

A mined :class:~ergodic.process.ProcessModel contributes its net, and a process tree converts to its workflow net up front. Each distinct variant replays and aligns once, then weights by its case count: fitness aggregates the replay token sums over all cases (Rozinat's form), alignment_fitness is the case-weighted mean of the per-variant alignment fitness, and a case is fitting when its variant aligns at cost zero. max_states caps every state search (the silent closures inside replay, the alignment search, and the model's cheapest-run search); past the cap a ValueError raises rather than returning a partial answer. A model that cannot reach its final marking at all raises too, since no alignment exists for it.

Source code in src/ergodic/process/_conformance.py

def conformance(
    log: EventLog,
    model: ProcessModel | ProcessTree | WorkflowNet,
    *,
    max_states: int = 100_000,
) -> ConformanceResult:
    """Check ``log`` against ``model`` by token replay and by optimal alignments.

    A mined :class:`~ergodic.process.ProcessModel` contributes its net, and a
    process tree converts to its workflow net up front. Each distinct
    variant replays and aligns once, then weights by its case count:
    ``fitness`` aggregates the replay token sums over all cases (Rozinat's
    form), ``alignment_fitness`` is the case-weighted mean of the per-variant
    alignment fitness, and a case is fitting when its variant aligns at cost
    zero. ``max_states`` caps every state search (the silent closures inside
    replay, the alignment search, and the model's cheapest-run search); past
    the cap a ``ValueError`` raises rather than returning a partial answer.
    A model that cannot reach its final marking at all raises too, since no
    alignment exists for it.
    """
    if isinstance(model, ProcessModel):
        net = model.net
    elif isinstance(model, ProcessTree):
        net = model.to_net()
    elif isinstance(model, WorkflowNet):
        net = model
    else:
        raise TypeError(
            f"model must be a ProcessModel, a ProcessTree, or a WorkflowNet, "
            f"got {type(model).__name__}"
        )
    counts = log.variants()
    if not counts:
        raise ValueError("conformance needs a log with at least one case")
    index = net.index()
    silents = tuple(tid for tid, label in index.labels.items() if label is None)
    by_label = _by_label(index)
    min_visible = _model_path_cost(net, index, max_states)

    records: list[VariantConformance] = []
    for sequence, n_cases in sorted(counts.items(), key=lambda item: (-item[1], item[0])):
        missing, remaining, consumed, produced = _replay_variant(
            net, index, by_label, silents, sequence, max_states
        )
        cost, moves = _align_variant(net, index, by_label, sequence, max_states)
        worst_cost = len(sequence) + min_visible
        records.append(
            VariantConformance(
                sequence=sequence,
                n_cases=n_cases,
                replay_fitness=0.5 * (1.0 - missing / consumed)
                + 0.5 * (1.0 - remaining / produced),
                missing=missing,
                remaining=remaining,
                consumed=consumed,
                produced=produced,
                cost=cost,
                worst_cost=worst_cost,
                alignment_fitness=_alignment_fitness(cost, worst_cost),
                alignment=moves,
            )
        )

    total_cases = sum(record.n_cases for record in records)
    total_missing = sum(record.n_cases * record.missing for record in records)
    total_consumed = sum(record.n_cases * record.consumed for record in records)
    total_remaining = sum(record.n_cases * record.remaining for record in records)
    total_produced = sum(record.n_cases * record.produced for record in records)
    n_fitting = sum(record.n_cases for record in records if record.cost == 0)
    return ConformanceResult(
        fitness=0.5 * (1.0 - total_missing / total_consumed)
        + 0.5 * (1.0 - total_remaining / total_produced),
        alignment_fitness=sum(r.n_cases * r.alignment_fitness for r in records) / total_cases,
        n_cases=total_cases,
        n_variants=len(records),
        n_fitting=n_fitting,
        share_fitting=n_fitting / total_cases,
        variants=tuple(records),
    )

decision_points(log, model)

Locate every xor split of model and the branch each case of log took.

model is a process tree or a mined :class:~ergodic.process.ProcessModel (its tree is read). Returns one :class:DecisionPoint per xor node, in depth-first order. The tree must carry each activity label in exactly one leaf (a mined tree does, by construction); a duplicated label raises, since branch membership would then be ambiguous by design rather than by data.

Source code in src/ergodic/process/_decisions.py

def decision_points(log: EventLog, model: ProcessTree | ProcessModel) -> tuple[DecisionPoint, ...]:
    """Locate every xor split of ``model`` and the branch each case of ``log`` took.

    ``model`` is a process tree or a mined :class:`~ergodic.process.ProcessModel`
    (its tree is read). Returns one :class:`DecisionPoint` per xor node, in
    depth-first order. The tree must carry each activity label in exactly one
    leaf (a mined tree does, by construction); a duplicated label raises, since
    branch membership would then be ambiguous by design rather than by data.
    """
    tree = model.tree if isinstance(model, ProcessModel) else model
    _check_unique_labels(tree)
    return tuple(_point_at(log, tree, path)[0] for path in _xor_paths(tree))

decision_table(log, model, *, point=None, outcome='remaining_time', covariates='auto')

Encode one xor split of model as a causal table: covariates, branch, outcome.

The landmark is per case the time of its first event inside the split's combined alphabet, the moment the case reveals its branch. Only cases assigned exactly one branch enter the table; cases touching two or more branches are dropped and counted in n_dropped_ambiguous. A split with a silent (tau or empty-alphabet) branch raises: taking it leaves no event, so the decision moment is unobservable without alignment-based replay.

Parameters:

Name	Type	Description	Default
log	EventLog	The event log the tree describes.	required
model	ProcessTree | ProcessModel	A process tree with each activity in exactly one leaf, or a mined :class:~ergodic.process.ProcessModel (its tree is read).	required
point	tuple[int, ...] | DecisionPoint | None	The xor node to analyze: its path, a :class:DecisionPoint, or None when the tree has exactly one xor node (more than one raises, listing the paths).	None
outcome	str	"remaining_time" (seconds from the landmark to the case's last event; the default, because time already spent before the decision is not an effect of the branch), "cycle_time", "n_events", or a case-constant column. The three derived names are reserved: a log column by any of them raises.	'remaining_time'
covariates	str | Sequence[str]	"auto" takes every case-constant variable column plus the derived features, each measured strictly before the landmark: events and elapsed seconds before it (n_events_prefix, prefix_activities, elapsed_at_landmark), the open-case workload at the start and at the landmark, and the calendar features when the timestamps are datetimes. A sequence restricts to the named columns and derived features. A derived covariate that is constant across the kept cases is dropped in silence.	'auto'

Returns:

Name	Type	Description
A	DecisionTable	class:DecisionTable: the :class:~ergodic.data.TabularData
	DecisionTable	(columns in the order sorted covariates, branch, outcome; one row
	DecisionTable	per kept case, in case order), the tiered
	DecisionTable	class:~ergodic.DomainKnowledge, the :class:DecisionPoint, and the
	DecisionTable	class:PositivityReport.

Source code in src/ergodic/process/_decisions.py

def decision_table(
    log: EventLog,
    model: ProcessTree | ProcessModel,
    *,
    point: tuple[int, ...] | DecisionPoint | None = None,
    outcome: str = "remaining_time",
    covariates: str | Sequence[str] = "auto",
) -> DecisionTable:
    """Encode one xor split of ``model`` as a causal table: covariates, branch, outcome.

    The landmark is per case the time of its first event inside the split's
    combined alphabet, the moment the case reveals its branch. Only cases
    assigned exactly one branch enter the table; cases touching two or more
    branches are dropped and counted in ``n_dropped_ambiguous``. A split with
    a silent (tau or empty-alphabet) branch raises: taking it leaves no event,
    so the decision moment is unobservable without alignment-based replay.

    Args:
        log: The event log the tree describes.
        model: A process tree with each activity in exactly one leaf, or a
            mined :class:`~ergodic.process.ProcessModel` (its tree is read).
        point: The xor node to analyze: its path, a :class:`DecisionPoint`, or
            None when the tree has exactly one xor node (more than one raises,
            listing the paths).
        outcome: ``"remaining_time"`` (seconds from the landmark to the case's
            last event; the default, because time already spent before the
            decision is not an effect of the branch), ``"cycle_time"``,
            ``"n_events"``, or a case-constant column. The three derived names
            are reserved: a log column by any of them raises.
        covariates: ``"auto"`` takes every case-constant variable column plus
            the derived features, each measured strictly before the landmark:
            events and elapsed seconds before it (``n_events_prefix``,
            ``prefix_activities``, ``elapsed_at_landmark``), the open-case
            workload at the start and at the landmark, and the calendar
            features when the timestamps are datetimes. A sequence restricts
            to the named columns and derived features. A derived covariate
            that is constant across the kept cases is dropped in silence.

    Returns:
        A :class:`DecisionTable`: the :class:`~ergodic.data.TabularData`
        (columns in the order sorted covariates, ``branch``, outcome; one row
        per kept case, in case order), the tiered
        :class:`~ergodic.DomainKnowledge`, the :class:`DecisionPoint`, and the
        :class:`PositivityReport`.
    """
    tree = model.tree if isinstance(model, ProcessModel) else model
    _check_unique_labels(tree)
    path = _resolve_path(tree, point)
    node = tree.subtree(path)
    if any(not child.activities() for child in node.children):
        raise ValueError(
            f"the split at {path} has a silent branch: taking it leaves no event, so "
            "the decision moment is unobservable without alignment-based replay; "
            "alignment-based extraction is future work"
        )
    if "branch" in log.schema.names:
        raise ValueError(
            "column 'branch' collides with the treatment column of the decision table; "
            "rename it so the treatment is unambiguous"
        )
    if outcome == "branch":
        raise ValueError("outcome 'branch' is the treatment column; pick another outcome")

    record, branch_idx, hit_row = _point_at(log, tree, path)
    keep = branch_idx >= 0
    if not bool(keep.any()):
        raise ValueError(
            f"no case takes the split at {path} unambiguously: {record.n_skipped} touch "
            f"no branch and {record.n_ambiguous} touch more than one"
        )

    _, starts = case_slices(log)
    _, first, last = case_bounds(log)
    stamps = time_seconds(log.column(log.timestamp))
    landmark = first.copy()
    landmark[keep] = stamps[hit_row[keep]]

    outcome_all, outcome_type = _outcome_per_case(log, outcome, first, last, landmark, starts)

    activities = log.column(log.activity)
    n_all = int(first.shape[0])
    n_prefix = np.zeros(n_all, dtype=np.float64)
    distinct = np.zeros(n_all, dtype=np.float64)
    for i in np.flatnonzero(keep):
        s, r = int(starts[i]), int(hit_row[i])
        n_prefix[i] = float(r - s)
        distinct[i] = float(len(set(activities[s:r].tolist())))

    derived: dict[str, npt.NDArray[np.float64]] = {
        "workload_at_start": open_cases(first, last, first),
        "elapsed_at_landmark": landmark - first,
        "n_events_prefix": n_prefix,
        "prefix_activities": distinct,
        "workload_at_landmark": open_cases(first, last, landmark),
    }
    if log.schema[log.timestamp].is_datetime:
        days = np.floor(first / _SECONDS_PER_DAY)
        derived["start_weekday"] = np.mod(days + 3.0, 7.0)  # the epoch was a Thursday
        derived["start_hour"] = np.floor((first - days * _SECONDS_PER_DAY) / _SECONDS_PER_HOUR)

    clash = sorted(set(derived) & set(log.schema.names))
    if clash:
        raise ValueError(
            f"columns {clash} collide with the derived covariate names of the decision "
            "table; rename the columns so the covariates are unambiguous"
        )

    if isinstance(covariates, str) and covariates == "auto":
        attributes = _auto_attributes(log, "branch", outcome)
        chosen_derived = dict(derived)
    else:
        requested = [covariates] if isinstance(covariates, str) else list(covariates)
        attributes, chosen_derived = _named_covariates(
            log, requested, derived, "branch", outcome, "decision"
        )

    columns: dict[str, npt.NDArray[Any]] = {}
    types: dict[str, str | Ordinal] = {}
    for name in attributes:
        columns[name] = attributes[name][keep]
        types[name] = _attribute_type(log.schema[name])
    for name, values in chosen_derived.items():
        kept_values = values[keep]
        if bool(np.all(kept_values == kept_values[0])):
            continue  # constant across the kept cases: useless and harmful downstream
        columns[name] = kept_values
        types[name] = "continuous"

    covariate_names = sorted(columns)
    ordered: dict[str, npt.NDArray[Any]] = {name: columns[name] for name in covariate_names}
    ordered["branch"] = np.array([record.branches[int(j)] for j in branch_idx[keep]])
    types["branch"] = "categorical"
    ordered[outcome] = outcome_all[keep]
    types[outcome] = outcome_type
    data = tabular(ordered, types=types)

    tiers: list[list[str]] = []
    if covariate_names:
        tiers.append(list(covariate_names))
    tiers.extend((["branch"], [outcome]))
    knowledge = DomainKnowledge().with_tiers(tiers)

    return DecisionTable(
        data=data,
        knowledge=knowledge,
        outcome=outcome,
        covariates=tuple(covariate_names),
        point=record,
        positivity=_positivity(data, tuple(covariate_names), record.branches),
        n_cases=int(keep.sum()),
        n_dropped_ambiguous=record.n_ambiguous,
    )

discover_model(log, *, noise=0.0)

Mine a block-structured process model from log.

The miner works on the log's variants, weighting each distinct trace by how many cases follow it, so a large log mines as fast as its variant table. noise must lie in [0, 1): at 0 the log is taken literally, and a larger value drops, at every recursion and before cut detection, each directly-follows edge carrying less than noise times the heaviest edge leaving the same activity. The returned :class:ProcessModel carries the tree, its workflow net, and the provenance counters.

Source code in src/ergodic/process/_mine.py

def discover_model(log: EventLog, *, noise: float = 0.0) -> ProcessModel:
    """Mine a block-structured process model from ``log``.

    The miner works on the log's variants, weighting each distinct trace by
    how many cases follow it, so a large log mines as fast as its variant
    table. ``noise`` must lie in ``[0, 1)``: at 0 the log is taken literally,
    and a larger value drops, at every recursion and before cut detection,
    each directly-follows edge carrying less than ``noise`` times the
    heaviest edge leaving the same activity. The returned
    :class:`ProcessModel` carries the tree, its workflow net, and the
    provenance counters.
    """
    if not 0.0 <= noise < 1.0:
        raise ValueError(f"noise must lie in [0.0, 1.0), got {noise}")
    variants = {trace: int(count) for trace, count in sorted(log.variants().items())}
    state = _MinerState()
    tree = _mine(variants, noise, state)
    fallthroughs = tuple(
        path for path, node in tree.nodes() if any(node is flower for flower in state.flowers)
    )
    share = 0.0 if log.n_rows == 0 else state.n_dropped_events / log.n_rows
    return ProcessModel(
        tree=tree,
        net=tree.to_net(),
        fallthroughs=fallthroughs,
        n_filtered_arcs=state.n_filtered_arcs,
        n_dropped_events=state.n_dropped_events,
        dropped_share=share,
        noise=noise,
    )

filter_activities(log, *, keep=None, mandatory=None, forbidden=None)

Drops cases: by mandatory and forbidden activities; keep then drops events.

A case is dropped when it misses any mandatory activity or contains any forbidden one, judged on its full trace. When keep is given, the surviving cases then lose every event whose activity is not in it, and a case left with no events vanishes. At least one argument is required, and an activity name the log has never seen raises.

Source code in src/ergodic/process/_filters.py

def filter_activities(
    log: EventLog,
    *,
    keep: Iterable[str] | None = None,
    mandatory: Iterable[str] | None = None,
    forbidden: Iterable[str] | None = None,
) -> EventLog:
    """Drops cases: by mandatory and forbidden activities; ``keep`` then drops events.

    A case is dropped when it misses any ``mandatory`` activity or contains
    any ``forbidden`` one, judged on its full trace. When ``keep`` is given,
    the surviving cases then lose every event whose activity is not in it,
    and a case left with no events vanishes. At least one argument is
    required, and an activity name the log has never seen raises.
    """
    keep_set = None if keep is None else set(keep)
    need = None if mandatory is None else set(mandatory)
    ban = None if forbidden is None else set(forbidden)
    if keep_set is None and need is None and ban is None:
        raise ValueError("give at least one of keep, mandatory, or forbidden")
    _check_known(log, (keep_set or set()) | (need or set()) | (ban or set()))
    acts = log.column(log.activity)
    cases, starts = case_slices(log)
    keep_case = np.ones(len(cases), dtype=bool)
    if need is not None or ban is not None:
        for i in range(len(cases)):
            block = set(acts[starts[i] : starts[i + 1]].tolist())
            keep_case[i] = (need is None or need <= block) and (ban is None or not ban & block)
    rows = _expand(log, keep_case)
    if keep_set is not None:
        rows &= np.isin(acts, sorted(keep_set))
    return _finish(log, rows)

filter_endpoints(log, *, start_in=None, end_in=None)

Drops cases: keep cases that start and end in the given activity sets.

start_in constrains a trace's first activity and end_in its last; a None set leaves that end unconstrained. The standard way to drop incomplete cases.

Source code in src/ergodic/process/_filters.py

def filter_endpoints(
    log: EventLog,
    *,
    start_in: Iterable[str] | None = None,
    end_in: Iterable[str] | None = None,
) -> EventLog:
    """Drops cases: keep cases that start and end in the given activity sets.

    ``start_in`` constrains a trace's first activity and ``end_in`` its last;
    a ``None`` set leaves that end unconstrained. The standard way to drop
    incomplete cases.
    """
    acts = log.column(log.activity)
    _, starts = case_slices(log)
    keep_case = np.ones(len(starts) - 1, dtype=bool)
    if start_in is not None:
        keep_case &= np.isin(acts[starts[:-1]], sorted(set(start_in)))
    if end_in is not None:
        keep_case &= np.isin(acts[starts[1:] - 1], sorted(set(end_in)))
    return _finish(log, _expand(log, keep_case))

filter_follows(log, a, b, *, directly=True, within=None)

Drops cases: keep cases where activity b follows activity a.

With directly (the default) b must be the very next event after an a; otherwise any later event counts. within bounds the gap in seconds from the a event to the b event, so a pair further apart does not count. Both activities must be known to the log.

Source code in src/ergodic/process/_filters.py

def filter_follows(
    log: EventLog,
    a: str,
    b: str,
    *,
    directly: bool = True,
    within: float | None = None,
) -> EventLog:
    """Drops cases: keep cases where activity ``b`` follows activity ``a``.

    With ``directly`` (the default) ``b`` must be the very next event after
    an ``a``; otherwise any later event counts. ``within`` bounds the gap in
    seconds from the ``a`` event to the ``b`` event, so a pair further apart
    does not count. Both activities must be known to the log.
    """
    _check_known(log, (a, b))
    acts = log.column(log.activity)
    stamps = time_seconds(log.column(log.timestamp))
    cases, starts = case_slices(log)
    keep_case = np.zeros(len(cases), dtype=bool)
    for i in range(len(cases)):
        lo, hi = int(starts[i]), int(starts[i + 1])
        keep_case[i] = _follows_in(acts[lo:hi], stamps[lo:hi], a, b, directly, within)
    return _finish(log, _expand(log, keep_case))

filter_performance(log, *, min_duration=None, max_duration=None, min_events=None, max_events=None)

Drops cases: by case duration in seconds (first to last event) and event count.

Bounds are inclusive and a None bound is open; give at least one.

Source code in src/ergodic/process/_filters.py

def filter_performance(
    log: EventLog,
    *,
    min_duration: float | None = None,
    max_duration: float | None = None,
    min_events: int | None = None,
    max_events: int | None = None,
) -> EventLog:
    """Drops cases: by case duration in seconds (first to last event) and event count.

    Bounds are inclusive and a ``None`` bound is open; give at least one.
    """
    if all(bound is None for bound in (min_duration, max_duration, min_events, max_events)):
        raise ValueError("give at least one bound")
    _, first, last = case_bounds(log)
    _, starts = case_slices(log)
    duration = last - first
    counts = np.diff(starts)
    keep_case = np.ones(duration.size, dtype=bool)
    if min_duration is not None:
        keep_case &= duration >= min_duration
    if max_duration is not None:
        keep_case &= duration <= max_duration
    if min_events is not None:
        keep_case &= counts >= min_events
    if max_events is not None:
        keep_case &= counts <= max_events
    return _finish(log, _expand(log, keep_case))

filter_time(log, *, start=None, end=None, mode='intersecting')

Drops cases: keep cases meeting the window on event times; mode "trim" drops events.

start and end accept whatever compares to the log's stamps, a datetime-like or a number, and a None bound is open. Three modes: "contained" keeps cases whose every event falls inside [start, end], "intersecting" (the default) keeps cases with at least one event inside, and "trim" keeps only the events inside, so a case losing all its events vanishes.

Source code in src/ergodic/process/_filters.py

def filter_time(
    log: EventLog,
    *,
    start: Any | None = None,
    end: Any | None = None,
    mode: str = "intersecting",
) -> EventLog:
    """Drops cases: keep cases meeting the window on event times; mode "trim" drops events.

    ``start`` and ``end`` accept whatever compares to the log's stamps, a
    datetime-like or a number, and a ``None`` bound is open. Three modes:
    "contained" keeps cases whose every event falls inside ``[start, end]``,
    "intersecting" (the default) keeps cases with at least one event inside,
    and "trim" keeps only the events inside, so a case losing all its events
    vanishes.
    """
    if mode not in ("contained", "intersecting", "trim"):
        raise ValueError(f"unknown mode {mode!r}; use 'contained', 'intersecting', or 'trim'")
    stamps = time_seconds(log.column(log.timestamp))
    inside = np.ones(log.n_rows, dtype=bool)
    if start is not None:
        inside &= stamps >= _bound_seconds(start)
    if end is not None:
        inside &= stamps <= _bound_seconds(end)
    if mode == "trim":
        return _finish(log, inside)
    _, starts = case_slices(log)
    reduce = np.logical_and.reduceat if mode == "contained" else np.logical_or.reduceat
    keep_case: npt.NDArray[np.bool_] = reduce(inside, starts[:-1])
    return _finish(log, _expand(log, keep_case))

filter_variants(log, *, top=None, min_cases=None)

Drops cases: keep cases whose full activity sequence is a frequent variant.

Give exactly one selector. top keeps the top most frequent variants (count descending, ties by sequence, the :func:variant_table ranking); min_cases keeps every variant with at least that many cases.

Source code in src/ergodic/process/_filters.py

def filter_variants(
    log: EventLog, *, top: int | None = None, min_cases: int | None = None
) -> EventLog:
    """Drops cases: keep cases whose full activity sequence is a frequent variant.

    Give exactly one selector. ``top`` keeps the ``top`` most frequent
    variants (count descending, ties by sequence, the :func:`variant_table`
    ranking); ``min_cases`` keeps every variant with at least that many cases.
    """
    if top is not None and min_cases is not None:
        raise ValueError("give exactly one of top or min_cases, not both")
    if top is not None:
        kept = {variant.sequence for variant in variant_table(log, top=top)}
    elif min_cases is not None:
        kept = {seq for seq, count in log.variants().items() if count >= min_cases}
    else:
        raise ValueError("give exactly one of top or min_cases")
    traces = log.traces()
    cases, _ = case_slices(log)
    keep_case = np.asarray([traces[case_id] in kept for case_id in cases], dtype=bool)
    return _finish(log, _expand(log, keep_case))

is_sound(net, *, max_states=100000)

Whether net is sound in van der Aalst's classical sense.

Three conditions, all decided on the real state space. Option to complete: from every marking reachable from the initial marking, the final marking is still reachable. Proper completion: every reachable marking that puts a token on the sink is exactly the final marking, so no run finishes with tokens left behind. No dead transitions: every transition fires on some reachable path.

The check explores markings breadth-first from :attr:~ergodic.process.WorkflowNet.initial_marking, recording which marking fires which transition into which marking. Proper completion is then read off the explored markings, dead transitions off the union of fired transition ids, and option to complete by walking the recorded state graph backward from the final marking: every explored marking must reach it, which also settles whether the final marking is reachable at all.

A sound net is bounded, so the exploration terminates. When the number of distinct markings passes max_states, the net is either unbounded (and so unsound) or genuinely this big, and the function raises a ValueError rather than guess; raise the cap if the net really is this big.

Source code in src/ergodic/process/_soundness.py

def is_sound(net: WorkflowNet, *, max_states: int = 100_000) -> bool:
    """Whether ``net`` is sound in van der Aalst's classical sense.

    Three conditions, all decided on the real state space. Option to
    complete: from every marking reachable from the initial marking, the
    final marking is still reachable. Proper completion: every reachable
    marking that puts a token on the sink is exactly the final marking, so
    no run finishes with tokens left behind. No dead transitions: every
    transition fires on some reachable path.

    The check explores markings breadth-first from
    :attr:`~ergodic.process.WorkflowNet.initial_marking`, recording which
    marking fires which transition into which marking. Proper completion is
    then read off the explored markings, dead transitions off the union of
    fired transition ids, and option to complete by walking the recorded
    state graph backward from the final marking: every explored marking must
    reach it, which also settles whether the final marking is reachable at
    all.

    A sound net is bounded, so the exploration terminates. When the number
    of distinct markings passes ``max_states``, the net is either unbounded
    (and so unsound) or genuinely this big, and the function raises a
    ``ValueError`` rather than guess; raise the cap if the net really is
    this big.
    """
    index = net.index()
    final = net.final_marking
    sink = net.sink

    # Forward exploration: the state graph as marking -> [(tid, next marking)].
    edges: dict[Marking, list[tuple[str, Marking]]] = {net.initial_marking: []}
    queue: deque[Marking] = deque([net.initial_marking])
    fired: set[str] = set()
    while queue:
        state = queue.popleft()
        for tid in index.enabled(state):
            following = index.fire(state, tid)
            fired.add(tid)
            edges[state].append((tid, following))
            if following not in edges:
                edges[following] = []
                queue.append(following)
                if len(edges) > max_states:
                    raise ValueError(
                        f"the reachable state space exceeded max_states={max_states}; "
                        "a sound workflow net is bounded, so either this net is "
                        "unbounded or it really is this big: raise max_states to "
                        "keep exploring"
                    )

    # Proper completion: a token on the sink means the run is over, exactly.
    for state in edges:
        if state != final and any(place == sink for place, _ in state):
            return False

    # No dead transitions: every transition fired somewhere above.
    if fired != set(index.labels):
        return False

    # Option to complete: every explored marking reaches the final marking,
    # read backward over the recorded state graph.
    if final not in edges:
        return False
    predecessors: dict[Marking, set[Marking]] = {state: set() for state in edges}
    for state, outgoing in edges.items():
        for _tid, following in outgoing:
            predecessors[following].add(state)
    completes = {final}
    frontier = [final]
    while frontier:
        state = frontier.pop()
        for prev in predecessors[state]:
            if prev not in completes:
                completes.add(prev)
                frontier.append(prev)
    return len(completes) == len(edges)

kpi_panel(log, period, *, unit=None, kpis=('cycle_time',))

Aggregate an event log into per-period KPI rows for the panel estimators.

Each case belongs to the bucket of its first event and, when unit is given, to its unit, read from a case-constant column of the log (a column that varies within a case raises). Every aggregate in a row is over the cases that start in that unit-period cell. A cell where no case starts is simply absent; the panel designs handle the unbalanced grid.

period picks the bucketing. With datetime stamps pass "day", "week", or "month"; the bucket labels are np.datetime64 values: the stamp's calendar day (datetime64[D]), the Monday of the stamp's week (datetime64[D]), or the stamp's calendar month (datetime64[M]). With numeric stamps pass a positive width w; a case starting at t lands in the bucket labeled floor(t / w) * w, in the stamps' own units. A calendar name on numeric stamps raises TypeError, as does a numeric width on datetime stamps, and a width at or below zero raises ValueError.

Each name in kpis becomes one float column, named after itself, in the order given. Three names are derived from the log: "cycle_time" is the mean case duration among the cell's cases, in seconds for datetime stamps and in the stamps' own units otherwise; "n_cases" counts the cases that start in the cell; "n_events" is the mean number of events per case. Any other name reads a case-constant numeric or boolean column of the log and aggregates it by its mean in the cell, so a 0/1 case flag that switches on a period boundary comes out as an exact 0/1 column, ready to be a did treatment; a non-numeric column raises ValueError. The three derived names are reserved: when the log also carries a column by a requested derived name, the request is ambiguous and raises ValueError.

The output columns are the unit column (named after the source column), a "period" column holding the bucket labels, then the kpi columns. A kpi named like the unit column or "period" collides with those index columns and raises, as does a unit column itself named "period". With unit given the result is a :class:~ergodic.data.PanelData keyed by (unit, "period"), the shape did, event_study, and synthetic_control consume. With unit=None the whole log is one series and the result is a :class:~ergodic.data.TimeSeriesData with time="period", so interrupted_time_series consumes it directly.

Two caveats. A case is attributed to the bucket where it starts, so a case that runs long spills its effect into earlier buckets: its whole duration counts where it began, even when it finished much later. And the aggregation assumes the log is complete, with no still-open cases; a censored log drops exactly the long cases near its end, which biases late-bucket cycle times downward.

Source code in src/ergodic/process/_panel.py

def kpi_panel(
    log: EventLog,
    period: str | float,
    *,
    unit: str | None = None,
    kpis: Iterable[str] = ("cycle_time",),
) -> PanelData | TimeSeriesData:
    """Aggregate an event log into per-period KPI rows for the panel estimators.

    Each case belongs to the bucket of its first event and, when ``unit`` is
    given, to its unit, read from a case-constant column of the log (a column
    that varies within a case raises). Every aggregate in a row is over the
    cases that start in that unit-period cell. A cell where no case starts is
    simply absent; the panel designs handle the unbalanced grid.

    ``period`` picks the bucketing. With datetime stamps pass ``"day"``,
    ``"week"``, or ``"month"``; the bucket labels are ``np.datetime64``
    values: the stamp's calendar day (``datetime64[D]``), the Monday of the
    stamp's week (``datetime64[D]``), or the stamp's calendar month
    (``datetime64[M]``). With numeric stamps pass a positive width ``w``; a
    case starting at ``t`` lands in the bucket labeled ``floor(t / w) * w``,
    in the stamps' own units. A calendar name on numeric stamps raises
    ``TypeError``, as does a numeric width on datetime stamps, and a width at
    or below zero raises ``ValueError``.

    Each name in ``kpis`` becomes one float column, named after itself, in
    the order given. Three names are derived from the log: ``"cycle_time"``
    is the mean case duration among the cell's cases, in seconds for datetime
    stamps and in the stamps' own units otherwise; ``"n_cases"`` counts the
    cases that start in the cell; ``"n_events"`` is the mean number of events
    per case. Any other name reads a case-constant numeric or boolean column
    of the log and aggregates it by its mean in the cell, so a 0/1 case flag
    that switches on a period boundary comes out as an exact 0/1 column,
    ready to be a ``did`` treatment; a non-numeric column raises
    ``ValueError``. The three derived names are reserved: when the log also
    carries a column by a requested derived name, the request is ambiguous
    and raises ``ValueError``.

    The output columns are the unit column (named after the source column),
    a ``"period"`` column holding the bucket labels, then the kpi columns. A
    kpi named like the unit column or ``"period"`` collides with those index
    columns and raises, as does a unit column itself named ``"period"``.
    With ``unit`` given the result is a :class:`~ergodic.data.PanelData`
    keyed by ``(unit, "period")``, the shape ``did``, ``event_study``, and
    ``synthetic_control`` consume. With ``unit=None`` the whole log is one
    series and the result is a :class:`~ergodic.data.TimeSeriesData` with
    ``time="period"``, so ``interrupted_time_series`` consumes it directly.

    Two caveats. A case is attributed to the bucket where it starts, so a
    case that runs long spills its effect into earlier buckets: its whole
    duration counts where it began, even when it finished much later. And
    the aggregation assumes the log is complete, with no still-open cases; a
    censored log drops exactly the long cases near its end, which biases
    late-bucket cycle times downward.
    """
    names = [str(name) for name in kpis]
    if len(set(names)) != len(names):
        raise ValueError(f"kpis repeat a name: {names}")
    if unit == "period":
        raise ValueError(
            "unit column 'period' collides with the 'period' column the output adds; "
            "rename the log column"
        )
    for name in names:
        if name == "period":
            raise ValueError("kpi 'period' collides with the output period column; rename it")
        if unit is not None and name == unit:
            raise ValueError(f"kpi {name!r} collides with the output unit column {unit!r}")
        if name in _DERIVED and name in log.schema:
            raise ValueError(
                f"kpi {name!r} is derived from the log, but the log also carries a column "
                f"named {name!r}; the request is ambiguous, so rename the column"
            )

    _, first, last = case_bounds(log)
    labels = _bucket_labels(first, period, log.timestamp, log.schema[log.timestamp].is_datetime)

    label_keys, label_inv = np.unique(labels, return_inverse=True)
    store: dict[str, npt.NDArray[Any]] = {}
    if unit is None:
        cell_of_case = np.asarray(label_inv, dtype=np.int64)
        n_cells = int(label_keys.size)
        store["period"] = label_keys
    else:
        unit_values = case_attribute(log, unit)
        unit_keys, unit_inv = np.unique(unit_values, return_inverse=True)
        stride = max(int(label_keys.size), 1)
        pair = np.asarray(unit_inv, dtype=np.int64) * stride + np.asarray(label_inv, dtype=np.int64)
        cell_ids, inverse = np.unique(pair, return_inverse=True)
        cell_of_case = np.asarray(inverse, dtype=np.int64)
        n_cells = int(cell_ids.size)
        store[unit] = unit_keys[cell_ids // stride]
        store["period"] = label_keys[cell_ids % stride]

    counts = np.bincount(cell_of_case, minlength=n_cells).astype(np.float64)

    def cell_mean(per_case: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        sums = np.bincount(cell_of_case, weights=per_case, minlength=n_cells)
        return np.asarray(sums / counts, dtype=np.float64)

    for name in names:
        if name == "cycle_time":
            store[name] = cell_mean(np.asarray(last - first, dtype=np.float64))
        elif name == "n_cases":
            store[name] = counts.copy()
        elif name == "n_events":
            _, starts = case_slices(log)
            store[name] = cell_mean(np.diff(starts).astype(np.float64))
        else:
            values = case_attribute(log, name)
            if values.dtype.kind not in "biuf":
                raise ValueError(
                    f"kpi column {name!r} has dtype {values.dtype}; a kpi aggregates by its "
                    "cell mean, so the column must be numeric or boolean"
                )
            store[name] = cell_mean(np.asarray(values, dtype=np.float64))

    if unit is None:
        return timeseries(store, time="period")
    return panel(store, entity=unit, time="period")

loop(do, redo)

A loop node: do runs, then any number of redo then do rounds.

Source code in src/ergodic/process/_tree.py

def loop(do: str | ProcessTree, redo: str | ProcessTree) -> ProcessTree:
    """A loop node: ``do`` runs, then any number of ``redo`` then ``do`` rounds."""
    return ProcessTree(op="loop", children=(_coerce(do), _coerce(redo)))

net_of_tree(tree)

The workflow net of a process tree, by structural recursion.

Each subtree is built between an entry and an exit place: a leaf becomes one transition (silent for :data:~ergodic.process.TAU), seq chains its children through fresh places, xor lets every child share the entry and exit, par brackets per-child places with silent split and join transitions, and loop enters through a silent step into a do body whose exit either leaves through a silent step or runs the redo body back to the do entry. The construction is sound for every tree (and :func:~ergodic.process.is_sound will verify it for real).

Source code in src/ergodic/process/_net.py

def net_of_tree(tree: ProcessTree) -> WorkflowNet:
    """The workflow net of a process tree, by structural recursion.

    Each subtree is built between an entry and an exit place: a leaf becomes
    one transition (silent for :data:`~ergodic.process.TAU`), ``seq`` chains
    its children through fresh places, ``xor`` lets every child share the
    entry and exit, ``par`` brackets per-child places with silent split and
    join transitions, and ``loop`` enters through a silent step into a do
    body whose exit either leaves through a silent step or runs the redo body
    back to the do entry. The construction is sound for every tree (and
    :func:`~ergodic.process.is_sound` will verify it for real).
    """
    places: list[str] = ["source", "sink"]
    transitions: list[tuple[str, str | None]] = []
    arcs: list[tuple[str, str]] = []

    def new_place() -> str:
        name = f"p{len(places) - 2}"
        places.append(name)
        return name

    def new_transition(label: str | None) -> str:
        tid = f"t{len(transitions)}"
        transitions.append((tid, label))
        return tid

    def build(node: ProcessTree, entry: str, exit_: str) -> None:
        if node.op is None:
            tid = new_transition(node.label)
            arcs.append((entry, tid))
            arcs.append((tid, exit_))
        elif node.op == "seq":
            current = entry
            for child in node.children[:-1]:
                following = new_place()
                build(child, current, following)
                current = following
            build(node.children[-1], current, exit_)
        elif node.op == "xor":
            for child in node.children:
                build(child, entry, exit_)
        elif node.op == "par":
            split = new_transition(None)
            join = new_transition(None)
            arcs.append((entry, split))
            arcs.append((join, exit_))
            for child in node.children:
                child_in = new_place()
                child_out = new_place()
                arcs.append((split, child_in))
                arcs.append((child_out, join))
                build(child, child_in, child_out)
        else:  # loop
            do_in = new_place()
            do_out = new_place()
            enter = new_transition(None)
            leave = new_transition(None)
            arcs.append((entry, enter))
            arcs.append((enter, do_in))
            arcs.append((do_out, leave))
            arcs.append((leave, exit_))
            build(node.children[0], do_in, do_out)
            build(node.children[1], do_out, do_in)

    build(tree, "source", "sink")
    return WorkflowNet(
        places=tuple(places),
        transitions=tuple(transitions),
        arcs=tuple(arcs),
        source="source",
        sink="sink",
    )

par(*children)

A parallel node: the children interleave. Strings become leaves.

Source code in src/ergodic/process/_tree.py

def par(*children: str | ProcessTree) -> ProcessTree:
    """A parallel node: the children interleave. Strings become leaves."""
    return ProcessTree(op="par", children=tuple(_coerce(child) for child in children))

performance(log, *, lifecycle=None)

Summarize case cycle times and per-activity waiting and service.

Case durations run first event to last, in seconds, summarized as the mean, median, and 90th percentile. Per activity, count is how many events carry it.

With lifecycle=None the log has one timestamp per activity, which cannot split waiting from working: mean_wait is the mean gap from the previous event in the same case (a case's first event contributes no wait), mean_service is None, and the result's lifecycle is False.

With lifecycle="column" the named column must exist (KeyError otherwise). Rows valued "start" and "complete" are paired per case and activity in time order, each start with the next complete of the same case and activity: service is complete minus start, wait is start minus the previous complete event in the case (a start with no earlier complete contributes no wait), and rows that never find a partner are counted in n_unmatched. Other lifecycle values are ignored. Durations still run first event to last over all events.

Source code in src/ergodic/process/_performance.py

def performance(log: EventLog, *, lifecycle: str | None = None) -> ProcessPerformance:
    """Summarize case cycle times and per-activity waiting and service.

    Case durations run first event to last, in seconds, summarized as the
    mean, median, and 90th percentile. Per activity, ``count`` is how many
    events carry it.

    With ``lifecycle=None`` the log has one timestamp per activity, which
    cannot split waiting from working: ``mean_wait`` is the mean gap from the
    previous event in the same case (a case's first event contributes no
    wait), ``mean_service`` is None, and the result's ``lifecycle`` is False.

    With ``lifecycle="column"`` the named column must exist (``KeyError``
    otherwise). Rows valued ``"start"`` and ``"complete"`` are paired per
    case and activity in time order, each start with the next complete of the
    same case and activity: service is complete minus start, wait is start
    minus the previous complete event in the case (a start with no earlier
    complete contributes no wait), and rows that never find a partner are
    counted in ``n_unmatched``. Other lifecycle values are ignored. Durations
    still run first event to last over all events.
    """
    _, first, last = case_bounds(log)
    durations = last - first
    if durations.size:
        mean_duration = float(np.mean(durations))
        median_duration = float(np.median(durations))
        p90_duration = float(np.percentile(durations, 90))
    else:
        mean_duration = median_duration = p90_duration = math.nan

    activities = log.column(log.activity)
    stamps = time_seconds(log.column(log.timestamp))
    cases, starts = case_slices(log)
    counts = Counter(str(name) for name in activities.tolist())

    wait_sums: defaultdict[str, float] = defaultdict(float)
    wait_counts: defaultdict[str, int] = defaultdict(int)
    service_sums: defaultdict[str, float] = defaultdict(float)
    service_counts: defaultdict[str, int] = defaultdict(int)
    n_unmatched = 0

    if lifecycle is None:
        for i in range(cases.size):
            lo, hi = int(starts[i]), int(starts[i + 1])
            for j in range(lo + 1, hi):
                name = str(activities[j])
                wait_sums[name] += float(stamps[j] - stamps[j - 1])
                wait_counts[name] += 1
    else:
        if lifecycle not in log.schema:
            raise KeyError(
                f"lifecycle column {lifecycle!r} not found; have {sorted(log.schema.names)}"
            )
        marks = log.column(lifecycle)
        for i in range(cases.size):
            lo, hi = int(starts[i]), int(starts[i + 1])
            open_starts: defaultdict[str, deque[float]] = defaultdict(deque)
            last_complete: float | None = None
            for j in range(lo, hi):
                mark = str(marks[j])
                name = str(activities[j])
                when = float(stamps[j])
                if mark == "start":
                    if last_complete is not None:
                        wait_sums[name] += when - last_complete
                        wait_counts[name] += 1
                    open_starts[name].append(when)
                elif mark == "complete":
                    if open_starts[name]:
                        service_sums[name] += when - open_starts[name].popleft()
                        service_counts[name] += 1
                    else:
                        n_unmatched += 1
                    last_complete = when
                # other lifecycle values are ignored
            n_unmatched += sum(len(queue) for queue in open_starts.values())

    stats = tuple(
        ActivityStats(
            activity=str(name),
            count=counts[str(name)],
            mean_wait=_mean(wait_sums, wait_counts, str(name)),
            mean_service=_mean(service_sums, service_counts, str(name)),
        )
        for name in log.activities
    )
    return ProcessPerformance(
        n_cases=log.n_cases,
        n_events=log.n_rows,
        mean_duration=mean_duration,
        median_duration=median_duration,
        p90_duration=p90_duration,
        activities=stats,
        lifecycle=lifecycle is not None,
        n_unmatched=n_unmatched,
    )

process_map(log)

Build the directly-follows map of log with counts and gap timing.

Every consecutive event pair within a case traverses the arc between its two activities; each real arc reports its count and the mean and median gap in seconds (datetime and numeric timestamps alike). Each case also traverses (START, first activity) and (last activity, END), whose gaps are NaN. An activity named like a boundary sentinel raises a ValueError, since the boundary arcs would then be ambiguous.

Source code in src/ergodic/process/_map.py

def process_map(log: EventLog) -> ProcessMap:
    """Build the directly-follows map of ``log`` with counts and gap timing.

    Every consecutive event pair within a case traverses the arc between its
    two activities; each real arc reports its count and the mean and median
    gap in seconds (datetime and numeric timestamps alike). Each case also
    traverses ``(START, first activity)`` and ``(last activity, END)``, whose
    gaps are NaN. An activity named like a boundary sentinel raises a
    ``ValueError``, since the boundary arcs would then be ambiguous.
    """
    clash = sorted(set(log.activities) & {START, END})
    if clash:
        raise ValueError(
            f"activities {clash} collide with the reserved boundary names "
            f"{START!r} and {END!r}; rename them before mapping"
        )
    activities = log.column(log.activity)
    stamps = time_seconds(log.column(log.timestamp))
    cases, starts = case_slices(log)

    gap_lists: defaultdict[tuple[str, str], list[float]] = defaultdict(list)
    first_counts: defaultdict[str, int] = defaultdict(int)
    last_counts: defaultdict[str, int] = defaultdict(int)
    for i in range(cases.size):
        lo, hi = int(starts[i]), int(starts[i + 1])
        first_counts[str(activities[lo])] += 1
        last_counts[str(activities[hi - 1])] += 1
        for j in range(lo, hi - 1):
            pair = (str(activities[j]), str(activities[j + 1]))
            gap_lists[pair].append(float(stamps[j + 1] - stamps[j]))

    entries: dict[tuple[str, str], ArcStats] = {}
    for pair, gaps in gap_lists.items():
        arr = np.asarray(gaps, dtype=np.float64)
        entries[pair] = ArcStats(
            count=len(gaps),
            mean_gap=float(np.mean(arr)),
            median_gap=float(np.median(arr)),
        )
    for name, count in first_counts.items():
        entries[(START, name)] = ArcStats(count=count, mean_gap=math.nan, median_gap=math.nan)
    for name, count in last_counts.items():
        entries[(name, END)] = ArcStats(count=count, mean_gap=math.nan, median_gap=math.nan)

    return ProcessMap(
        activities=tuple(str(name) for name in log.activities),
        arcs=dict(sorted(entries.items())),
        n_cases=log.n_cases,
        n_events=log.n_rows,
    )

seq(*children)

A sequence node: the children run left to right. Strings become leaves.

Source code in src/ergodic/process/_tree.py

def seq(*children: str | ProcessTree) -> ProcessTree:
    """A sequence node: the children run left to right. Strings become leaves."""
    return ProcessTree(op="seq", children=tuple(_coerce(child) for child in children))

variant_table(log, *, top=None)

The distinct variants, most frequent first.

Ranked by case count descending, ties by sequence. share is the fraction of all cases on the variant, and the durations are first-to-last event times in seconds over the variant's cases. top truncates the table; None keeps every variant.

Source code in src/ergodic/process/_variants.py

def variant_table(log: EventLog, *, top: int | None = None) -> tuple[Variant, ...]:
    """The distinct variants, most frequent first.

    Ranked by case count descending, ties by sequence. ``share`` is the
    fraction of all cases on the variant, and the durations are first-to-last
    event times in seconds over the variant's cases. ``top`` truncates the
    table; ``None`` keeps every variant.
    """
    if top is not None and top < 1:
        raise ValueError(f"top must be at least 1, got {top}")
    traces = log.traces()
    cases, first, last = case_bounds(log)
    durations = last - first
    by_sequence: dict[tuple[str, ...], list[float]] = {}
    for i, case_id in enumerate(cases):
        by_sequence.setdefault(traces[case_id], []).append(float(durations[i]))
    n_total = len(cases)
    order = sorted(by_sequence, key=lambda seq: (-len(by_sequence[seq]), seq))
    table = tuple(
        Variant(
            sequence=seq,
            n_cases=len(by_sequence[seq]),
            share=len(by_sequence[seq]) / n_total,
            mean_duration=float(np.mean(by_sequence[seq])),
            median_duration=float(np.median(by_sequence[seq])),
        )
        for seq in order
    )
    return table if top is None else table[:top]

xor(*children)

An exclusive choice node: exactly one child runs. Strings become leaves.

Source code in src/ergodic/process/_tree.py

def xor(*children: str | ProcessTree) -> ProcessTree:
    """An exclusive choice node: exactly one child runs. Strings become leaves."""
    return ProcessTree(op="xor", children=tuple(_coerce(child) for child in children))

Structural causal models

The Bayesian SCM: mechanisms, interventions, counterfactuals, and the identification-gated ate.

ergodic.scm.SCM

A Bayesian structural causal model over a DAG or a latent-confounded ADMG.

Build with SCM(graph, data) (mechanisms auto-assigned by schema) or pass mechanisms={node: spec} to override. spec is a built-in name ("linear", "categorical", "binary") or a :class:~ergodic.scm.Mechanism instance. Call :meth:fit, then query.

The graph may be a Markovian DAG, or an ADMG (<-> edges) / a DAG with latent= confounders. A bidirected edge is realized as one explicit per-unit latent common cause: the confounded district fits jointly (the latent absorbs the confounding) so an effect identified despite the confounding (front-door, the general criterion with the latent modeled) computes correctly. A genuinely non-identified effect still returns a :class:~ergodic.scm.Sensitivity through :meth:ate, never a naked point -- identification, not the fit, decides. Selection bias stays out of scope (identification raises on selection marks).

Source code in src/ergodic/scm/_scm.py

class SCM:
    """A Bayesian structural causal model over a DAG or a latent-confounded ADMG.

    Build with ``SCM(graph, data)`` (mechanisms auto-assigned by schema) or pass
    ``mechanisms={node: spec}`` to override. ``spec`` is a built-in name
    (``"linear"``, ``"categorical"``, ``"binary"``) or a
    :class:`~ergodic.scm.Mechanism` instance. Call :meth:`fit`, then query.

    The graph may be a Markovian DAG, or an ADMG (``<->`` edges) / a DAG with
    ``latent=`` confounders. A bidirected edge is realized as one explicit per-unit
    latent common cause: the confounded district fits jointly (the latent absorbs the
    confounding) so an effect identified despite the confounding (front-door, the
    general criterion with the latent modeled) computes correctly. A genuinely
    non-identified effect still returns a :class:`~ergodic.scm.Sensitivity` through
    :meth:`ate`, never a naked point -- identification, not the fit, decides.
    Selection bias stays out of scope (identification raises on selection marks).
    """

    def __init__(
        self,
        graph: MixedGraph,
        data: Dataset,
        *,
        mechanisms: Mapping[str, Any] | None = None,
        draws: int = 1000,
        tune: int = 1000,
        chains: int = 2,
        random_state: int | None = None,
        district_size_cap: int = DISTRICT_SIZE_CAP,
    ) -> None:
        if not graph.is_admg():
            raise ValueError(
                f"SCM fits a DAG or an ADMG (latent-confounded); got kind="
                f"{graph.kind.value}, which is not acyclic-directed-plus-bidirected. "
                "A CPDAG/MAG/PAG is an equivalence class, not a single SCM; pick a "
                "member (dags_of_cpdag / mags_of_pag) or pass a DAG/ADMG"
            )
        # The observed margin: nodes carrying data. A ``latent=`` node has no column
        # and is realized as a ``<->`` over the projection; an ADMG's nodes are all
        # observed. The latent projection turns either into an ADMG over the observed
        # margin, which is what the fit and identification read.
        self.observed: tuple[str, ...] = tuple(sorted(graph.nodes - graph.latent))
        missing = set(self.observed) - set(data.variables)
        if missing:
            raise ValueError(
                f"every observed graph node needs a data column; missing: "
                f"{', '.join(sorted(missing))}"
            )
        self.graph = graph
        self.data = data
        self.draws = draws
        self.tune = tune
        self.chains = chains
        self.random_state = random_state
        self.district_size_cap = district_size_cap

        # The projected ADMG over the observed margin: ``latent=`` nodes become
        # ``<->`` edges; an ADMG that already has them passes through. Parents,
        # districts, and latents are all read off this single object, so a DAG with a
        # latent confounder and the equivalent ADMG fit identically.
        self._projection = latent_projection(graph)
        self._latent = latent_structure(self._projection)
        full_order = topological_order(self._projection)
        self._order: tuple[str, ...] = tuple(n for n in full_order if n in set(self.observed))
        specs = dict(mechanisms or {})
        self.mechanisms: dict[str, Mechanism] = {
            node: self._build_mechanism(specs.get(node), node) for node in self._order
        }
        self._warn_large_districts()
        self._fits: dict[str, FitResult] = {}

    def _build_mechanism(self, spec: Any, node: str) -> Mechanism:
        """Resolve ``node``'s mechanism over the projection, carrying its latents.

        The parents come from the projected ADMG (so a projection-introduced edge is
        honored), and the mechanism is given its incident latents and the subset it
        anchors the sign gauge for. A Markovian node lists no latents and reduces to
        the MVP mechanism exactly.
        """
        from dataclasses import replace

        mech = resolve_mechanism(spec, self._projection, self.data, node)
        lats = self._latent.latents_of(node)
        if lats:
            # only the two built-in mechanisms carry the latent-loading fields; a custom
            # mechanism in a confounded district has no place to put the loadings.
            if not isinstance(mech, GaussianAdditive | CategoricalGumbelMax):
                raise TypeError(
                    f"node {node!r} sits in a confounded district (incident latents "
                    f"{list(lats)}) but its mechanism {type(mech).__name__} is not a "
                    "built-in Gaussian/categorical, so it cannot carry latent loadings; "
                    "a custom mechanism in a bidirected district is post-MVP"
                )
            mech = replace(
                mech,
                latents=lats,
                gauged_latents=self._latent.gauged_of(node),
            )
        return mech

    def _warn_large_districts(self) -> None:
        """Warn (not raise) when a confounded district exceeds the sampling-cost cap.

        Decision 5: a big district samples slower, not wrong, so the cap is advisory.
        """
        import warnings

        for nodes in self._latent.confounded_districts:
            if len(nodes) > self.district_size_cap:
                warnings.warn(
                    f"confounded district {list(nodes)} has {len(nodes)} nodes, over "
                    f"the cap {self.district_size_cap}; the joint fit will sample but "
                    "may be slow. Raise district_size_cap to silence, or screen the "
                    "graph (e.g. a discovery restrict_to) to shrink the district",
                    stacklevel=3,
                )

    # -- fit -----------------------------------------------------------------

    @property
    def fitted(self) -> bool:
        return len(self._fits) == len(self.mechanisms)

    def fit(self, *, parallel: bool = False) -> SCM:
        """Fit every node's mechanism, per node for singletons, jointly per district.

        A Markovian singleton district fits per node (cheap, exact, the MVP path); a
        confounded district (nodes sharing a ``<->``) does not factorize -- the latent
        couples it -- so it fits as one joint NUTS block with the explicit latents,
        the gauge, a non-centered parameterization, and a high ``target_accept`` (the
        Track-D recipe). The composed posterior merges the two the same way it merges
        per-node fits.

        Sequential by default (robust, CI-safe). ``parallel=True`` spreads the
        *singleton* nodes over a process pool with the BLAS/OMP threads pinned in the
        workers; a confounded district's joint fit stays in the main process (its
        chains already parallelize and its model is one block). Returns ``self``.
        """
        seed = 0 if self.random_state is None else self.random_state
        singletons = set(self._latent.singletons)
        single_mechs = {n: m for n, m in self.mechanisms.items() if n in singletons}
        if parallel:
            fits = self._fit_parallel(single_mechs, seed)
        else:
            fits = {
                node: mech.fit(
                    self.data, draws=self.draws, tune=self.tune, chains=self.chains, seed=seed
                )
                for node, mech in single_mechs.items()
            }
        for offset, nodes in enumerate(self._latent.confounded_districts):
            fits.update(self._fit_district(nodes, seed + offset))
        self._fits = fits
        return self

    def _fit_district(self, nodes: tuple[str, ...], seed: int) -> dict[str, FitResult]:
        """Joint fit of one confounded district (its nodes + their incident latents).

        Assembles the observed columns for every district node and every external
        parent (a parent outside the district, e.g. the front-door mediator that is
        its own singleton), names the district's latents, and runs the joint NUTS fit.
        """
        order = tuple(n for n in self._order if n in set(nodes))
        latents = self._district_latents(nodes)
        observed = self._district_observed(order)
        return _compose.fit_district(
            order,
            {n: self.mechanisms[n] for n in order},
            latents,
            observed,
            draws=self.draws,
            tune=self.tune,
            chains=self.chains,
            seed=seed,
        )

    def _district_latents(self, nodes: tuple[str, ...]) -> tuple[str, ...]:
        """The latents incident to a district's nodes, deduplicated and sorted."""
        seen: set[str] = set()
        for node in nodes:
            seen.update(self._latent.latents_of(node))
        return tuple(sorted(seen))

    def _district_observed(self, order: tuple[str, ...]) -> dict[str, Any]:
        """Observed columns for the district fit: each node plus every external parent.

        A node's parents may sit outside the district (the canonical case is the
        front-door mediator, a singleton parent of an in-district node). The joint
        fit conditions on those observed parent columns, so the structural
        coefficient on a clean path is recovered while the latent soaks up the
        confounding on the bidirected one.
        """
        names: set[str] = set(order)
        for node in order:
            names.update(self.mechanisms[node].parents)
        out: dict[str, Any] = {}
        for name in names:
            mech = self.mechanisms.get(name)
            if mech is not None and isinstance(mech, CategoricalGumbelMax):
                # a categorical district node is observed as integer codes (its
                # likelihood is Categorical); the raw float code column also serves
                # as a discrete parent's design (the encoder one-hots it).
                out[name] = np.asarray(self.data.codes(name), dtype=np.int64)
            else:
                out[name] = np.asarray(self.data.matrix([name]).reshape(-1), dtype=np.float64)
        return out

    def _fit_parallel(self, mechanisms: Mapping[str, Mechanism], seed: int) -> dict[str, FitResult]:
        import concurrent.futures
        import multiprocessing as mp

        ctx = mp.get_context("spawn")
        out: dict[str, FitResult] = {}
        if not mechanisms:
            return out
        with concurrent.futures.ProcessPoolExecutor(
            mp_context=ctx, initializer=_pin_threads
        ) as pool:
            futures = {
                pool.submit(
                    _fit_one,
                    mech,
                    self.data,
                    self.draws,
                    self.tune,
                    self.chains,
                    seed,
                ): node
                for node, mech in mechanisms.items()
            }
            for fut in concurrent.futures.as_completed(futures):
                out[futures[fut]] = fut.result()
        return out

    def _require_fit(self) -> None:
        if not self.fitted:
            raise RuntimeError("SCM is not fitted; call .fit() first")

    def _seed(self, offset: int) -> int:
        return (0 if self.random_state is None else self.random_state) + offset

    # -- observational + interventional --------------------------------------

    def sample(self, n: int = 1000) -> SCMResult | CategoricalResult:
        """Observational draws of every node (read the target with :meth:`marginal`).

        Returns the posterior predictive of the generative model over ``n`` units;
        the result carries the draws of the last node in topological order (a sink)
        as a summary. A categorical sink returns a :class:`CategoricalResult` (the
        per-category marginal). Use :meth:`marginal` to read a specific node, or
        :meth:`interventional` / :meth:`counterfactual` for a query on one.
        """
        self._require_fit()
        target = self._order[-1]
        return self._observational(target, n=n, seed=self._seed(1))

    def marginal(self, target: str, n: int = 1000) -> SCMResult | CategoricalResult:
        """Observational posterior-predictive draws of one node.

        A continuous node returns an :class:`SCMResult`; a categorical/binary node a
        :class:`CategoricalResult` (the per-category marginal), exactly as
        :meth:`counterfactual` does for a categorical target.
        """
        self._require_fit()
        self._check_observable(target, "query target")
        return self._observational(target, n=n, seed=self._seed(2))

    def _observational(self, target: str, *, n: int, seed: int) -> SCMResult | CategoricalResult:
        """The shared observational readout: codes-to-marginal for a categorical target."""
        draws = _compose.observational_draws(
            self._order,
            self.mechanisms,
            self._fits,
            target,
            n=n,
            seed=seed,
            latents=self._latent.latents,
        )
        if isinstance(self.mechanisms[target], CategoricalGumbelMax):
            return self._categorical_result("sample", target, draws)
        return SCMResult(
            query="sample", target=target, kind="posterior", identified=True, _draws=draws
        )

    def do(self, intervention: Mapping[str, DoValue]) -> InterventionView:
        """A lightweight intervened view of the SCM for chaining.

        ``scm.do({"X": 1.0}).interventional("Y")`` reads the post-intervention
        distribution. The view shares the fit; it does not refit. A continuous
        do-target takes a float; a categorical/binary do-target takes its integer
        category code.
        """
        for node in intervention:
            self._check_observable(node, "do-target")
            # validate a categorical code early when the SCM is fitted (the resolved
            # class count lives on the fit); an unfitted chainable view still builds,
            # and interventional() re-validates before it samples regardless.
            if self.fitted:
                self._check_do_value(node, intervention[node])
        return InterventionView(self, dict(intervention))

    def interventional(
        self, target: str, do: Mapping[str, DoValue] | None = None, *, n: int = 2000
    ) -> SCMResult | CategoricalResult:
        """``E[Y | do(X = x)]`` and its posterior, through ``pm.do``.

        ``do`` maps each intervened node to its value (broadcast to a length-n
        vector internally, the verified shape gotcha; an integer code for a
        categorical/binary do-target, a float for a continuous one). ``target`` must
        not be an intervened node. A categorical target returns a
        :class:`CategoricalResult` (the post-intervention per-category marginal); a
        continuous target an :class:`SCMResult`. Carries parameter-posterior
        uncertainty via the merged ``idata``.
        """
        self._require_fit()
        self._check_observable(target, "query target")
        intervention = dict(do or {})
        for node in intervention:
            self._check_observable(node, "do-target")
            self._check_do_value(node, intervention[node])
        draws = _compose.interventional_draws(
            self._order,
            self.mechanisms,
            self._fits,
            intervention,
            target,
            n=n,
            seed=self._seed(3),
            latents=self._latent.latents,
        )
        if isinstance(self.mechanisms[target], CategoricalGumbelMax):
            return self._categorical_result("interventional", target, draws)
        return SCMResult(
            query="interventional", target=target, kind="posterior", identified=True, _draws=draws
        )

    # -- counterfactual ------------------------------------------------------

    def counterfactual(
        self,
        units: Dataset | Mapping[str, Any],
        do: Mapping[str, float],
        target: str,
        *,
        use_draws: bool = True,
    ) -> SCMResult | CategoricalResult:
        """The twin-network counterfactual for the given units under ``do``.

        Abduct each unit's structural noise on the factual graph, then predict
        ``target`` on the action branch reusing it. ``units`` is a dataset (or a
        mapping of node name to a column of factual values) holding every node's
        factual value per unit. For invertible mechanisms the counterfactual is
        exact (``use_draws=False`` collapses to the deterministic point; the
        default carries parameter uncertainty). A categorical ``target`` returns a
        :class:`~ergodic.scm.CategoricalResult` (the answer is a distribution).
        """
        self._require_fit()
        self._check_observable(target, "counterfactual target")
        for node in do:
            self._check_observable(node, "do-target")
            self._check_do_value(node, do[node])
        factual = self._factual_values(units)

        mech = self.mechanisms[target]
        if isinstance(mech, CategoricalGumbelMax):
            codes, _ = _compose.counterfactual_categorical(
                self.mechanisms,
                self._fits,
                factual,
                do,
                target,
                parents_of_target=mech.parents,
                seed=self._seed(5),
            )
            # the abducted-noise counterfactual is exact (invertible), so identified.
            return self._categorical_result("counterfactual", target, codes, per_unit=codes)

        if not isinstance(mech, GaussianAdditive):
            raise TypeError(f"no counterfactual for mechanism {type(mech).__name__} on {target!r}")
        # The twin-network counterfactual abducts the exact structural noise, which
        # is only valid for invertible mechanisms. A non-invertible (probabilistic-
        # only) node would need the genuinely wider abduction posterior, which is not
        # in the MVP; refuse rather than reuse the exact residual and mislabel the
        # result. Every shipped mechanism is invertible, so this guards a future
        # mechanism, not today's worlds.
        non_invertible = sorted(n for n, m in self.mechanisms.items() if not _is_invertible(m))
        if non_invertible:
            raise NotImplementedError(
                f"node(s) {non_invertible} use a non-invertible (probabilistic-only) "
                "mechanism; non-invertible-mechanism counterfactuals (the genuinely "
                "wider abduction posterior) are not in the MVP. Every built-in "
                "mechanism is invertible, so this only arises for a custom mechanism"
            )
        draws = _compose.counterfactual_continuous(
            self._order,
            self.mechanisms,
            self._fits,
            factual,
            do,
            target,
            use_draws=use_draws,
            seed=self._seed(6),
        )
        return SCMResult(
            query="counterfactual",
            target=target,
            kind="posterior",
            identified=True,
            _draws=draws,
        )

    # -- individual treatment effects ----------------------------------------

    def ite(
        self,
        units: Dataset | Mapping[str, Any],
        treatment: str,
        outcome: str,
        *,
        levels: tuple[float, float] | None = None,
        shift: float | None = None,
    ) -> ITEResult:
        """Per-unit individual treatment effects, through the twin network.

        For each unit, the exact counterfactual contrast of ``outcome`` under two
        treatment worlds that share *that unit's* abducted noise (the same-person,
        two-worlds twin network). Returns one effect per unit, ordered as ``units`` --
        the structural-model sibling of a CATE estimator's
        :meth:`~ergodic.inference.OutcomeRegression.cate`, which returns a *conditional
        average* over the noise rather than the individual's own counterfactual.

        Two framings of the contrast, and you pick one:

        * an absolute **level** contrast (the default): ``Y | do(X=high) -
          Y | do(X=low)``, both worlds pinning the treatment to a fixed level for every
          unit, the very contrast :meth:`ate` averages. Pass ``levels=(low, high)``; it
          defaults to ``(0.0, 1.0)`` (a unit step, or the two codes of a binary
          treatment). Works for a continuous, binary, or (with explicit codes)
          categorical treatment, and for a continuous or binary outcome -- a binary
          outcome's per-unit effect is its counterfactual flip, in ``{-1, 0, 1}``.
        * a **shift** (a "shock"): ``Y | do(X = x_i + shift) - Y_factual``, perturbing
          each unit's treatment from its *own* factual value, against the world where
          nothing was done. Pass ``shift=delta`` (a continuous treatment and outcome).

        The two coincide for a linear-additive mechanism: the shift contrast is the
        constant ``coef * delta`` for every unit, equal to the matching level step. They
        part only once a mechanism is nonlinear, where the shift is the unit's local
        response at its own baseline. So with the built-in linear-Gaussian mechanisms
        the continuous ITE is constant across units (the coefficient), and genuine
        per-unit spread comes from a **binary or categorical outcome**, whose logistic
        response varies with each unit's baseline.

        Each effect is the exact point (invertible mechanisms). For an individual's
        posterior *interval*, call :meth:`counterfactual` on that one unit per arm with
        ``use_draws=True`` and difference the draws.
        """
        from ..inference import ITEResult

        self._require_fit()
        self._check_observable(treatment, "treatment")
        self._check_observable(outcome, "outcome")
        if levels is not None and shift is not None:
            raise ValueError(
                "pass either levels= or shift=, not both (they are two framings of the "
                "same per-unit contrast)"
            )
        out_mech = self.mechanisms[outcome]
        if isinstance(out_mech, CategoricalGumbelMax):
            n_cat = self._fits[outcome].n_categories or out_mech.n_categories
            if n_cat > 2:
                raise ValueError(
                    f"ite needs a numeric outcome; {outcome!r} has {n_cat} categories, "
                    f"which has no scalar individual effect. Read per-unit counterfactual "
                    f"distributions with counterfactual(units, do, {outcome!r}) instead"
                )
        if shift is not None:
            values = self._ite_shift(units, treatment, outcome, shift)
            framing = f"shift={shift}"
            note = (
                "structural individual treatment effects: per-unit "
                f"Y|do({treatment}=x_i+{shift}) - Y_factual through the twin network "
                "(abduct the unit's noise, predict on the shifted branch)"
            )
        else:
            low, high = self._ate_levels(treatment, levels)
            values = self._ite_levels(units, treatment, outcome, low, high)
            framing = f"levels=({low}, {high})"
            note = (
                "structural individual treatment effects: per-unit "
                f"Y|do({treatment}={high}) - Y|do({treatment}={low}) on shared abducted "
                "noise (the same-person, two-worlds twin network)"
            )
        factual = self._factual_values(units)
        out_col = np.asarray(factual.get(outcome), dtype=np.float64)
        n = int(values.shape[0])
        return ITEResult(
            treatment=treatment,
            outcome=outcome,
            covariates=(),
            n=n,
            kind="point",
            source="scm",
            framing=framing,
            note=note,
            _values=np.asarray(values, dtype=np.float64),
            _index=np.arange(n, dtype=np.int64),
            _factual=out_col,
        )

    def _ite_levels(
        self,
        units: Dataset | Mapping[str, Any],
        treatment: str,
        outcome: str,
        low: float,
        high: float,
    ) -> npt.NDArray[np.float64]:
        """The per-unit ``Y|do(X=high) - Y|do(X=low)`` contrast, on shared abducted noise.

        Both arms route through :meth:`counterfactual` (same units, same fixed seed), so
        each unit's noise is identical across the two worlds and the contrast is a
        coherent individual effect, not two independent draws. A continuous outcome
        differences the per-unit draws; a binary outcome differences the per-unit
        counterfactual codes (the flip, in ``{-1, 0, 1}``).
        """
        high_cf = self.counterfactual(units, {treatment: high}, outcome, use_draws=False)
        low_cf = self.counterfactual(units, {treatment: low}, outcome, use_draws=False)
        if isinstance(high_cf, CategoricalResult):
            assert isinstance(low_cf, CategoricalResult)
            high_u, low_u = high_cf.per_unit(), low_cf.per_unit()
            assert high_u is not None and low_u is not None
            return high_u.astype(np.float64) - low_u.astype(np.float64)
        assert isinstance(low_cf, SCMResult)
        return high_cf.draws() - low_cf.draws()

    def _ite_shift(
        self,
        units: Dataset | Mapping[str, Any],
        treatment: str,
        outcome: str,
        shift: float,
    ) -> npt.NDArray[np.float64]:
        """The per-unit shock ``Y|do(X = x_i + shift) - Y_factual``, from each unit's baseline.

        A continuous treatment and outcome only. The baseline world ``do(X = x_i)`` pins
        the treatment to its own factual value, which (on invertible mechanisms)
        reproduces the factual outcome exactly, so the subtrahend is just the factual
        column. The shifted world hands the engine a per-unit do-vector ``x_i + shift``
        -- the one query that reads a non-scalar do.
        """
        if isinstance(self.mechanisms[treatment], CategoricalGumbelMax):
            raise ValueError(
                f"shift= is for a continuous treatment; {treatment!r} is categorical/"
                f"binary, which has no additive shift. Use levels=(low_code, high_code)"
            )
        if isinstance(self.mechanisms[outcome], CategoricalGumbelMax):
            raise ValueError(
                f"shift= needs a continuous outcome; {outcome!r} is categorical/binary. "
                f"Use a levels= contrast (do(X=high) - do(X=low)) instead"
            )
        factual = self._factual_values(units)
        x_factual = np.asarray(factual[treatment], dtype=np.float64)
        shifted = _compose.counterfactual_continuous(
            self._order,
            self.mechanisms,
            self._fits,
            factual,
            {treatment: x_factual + shift},
            outcome,
            use_draws=False,
            seed=self._seed(6),
        )
        baseline = np.asarray(factual[outcome], dtype=np.float64)
        return shifted - baseline

    # -- effects, identification-gated ---------------------------------------

    def ate(
        self,
        treatment: str,
        outcome: str,
        *,
        levels: tuple[float, float] | None = None,
        engine: str = "posterior",
        n: int = 2000,
    ) -> SCMResult | Sensitivity:
        """The average treatment effect, routed through identification first.

        ``identify_effect`` decides the return shape. Identified *by back-door
        adjustment* (or the no-path ``"null"`` case): the SCM posterior plug-in
        contrast ``E[Y | do(X = high)] - E[Y | do(X = low)]`` (a posterior
        :class:`~ergodic.scm.SCMResult`). Not identified: a
        :class:`~ergodic.scm.Sensitivity` (the ``b(c)`` curve over the assumed
        confounding), so a naked point is impossible.

        The treatment may be continuous, binary, or categorical. ``levels`` is the
        ``(low, high)`` contrast; it defaults to ``(0.0, 1.0)`` (the two codes of a
        binary treatment, or a unit step for a continuous one). A multi-category
        treatment has no canonical pair, so it *requires* explicit ``levels`` (the
        two category codes to contrast), mirroring the inference pillar's ``treated=``
        pattern. The outcome must be numeric (a categorical outcome has no scalar
        contrast and raises); confounders and mediators may be any type -- a
        categorical one is sampled in the generative model, so the back-door
        g-formula marginalizes over it correctly.

        Identified by back-door adjustment, the no-path ``"null"`` case, or
        **front-door** computes the latent-aware do-contrast: a confounded district
        fits its explicit latent, so the structural mechanism on the causal path fits
        clean and ``do`` severs the treatment from the district latent, leaving the
        identified value (the front-door mediated effect, the C1 world's 1.2). The
        **instrument** strategy still raises with a pointer to
        :func:`~ergodic.inference.estimate_effect`: IV identification is linear and
        rests on the instrument's exclusion restriction, which the latent-district fit
        does not impose, so its treatment->outcome coefficient stays on the same
        non-identified ridge as the bow. The **general ID functional** likewise
        raises for now (verify-then-widen: its latent-fit do-contrast was not held to
        ``estimate_effect`` to tolerance in this phase). ``engine="dml"`` is the
        cross-fit reuse path (post-MVP; raises with a pointer for now).
        """
        self._require_fit()
        self._check_node(treatment)
        self._check_node(outcome)
        # the outcome must be numeric: a categorical contrast is not a scalar.
        if isinstance(self.mechanisms.get(outcome), CategoricalGumbelMax):
            raise ValueError(
                f"ate needs a numeric outcome; {outcome!r} is categorical/binary, which "
                f"has no scalar contrast. Read its post-intervention distribution with "
                f"interventional({outcome!r}, do={{{treatment!r}: ...}}) instead"
            )
        ident = identify_effect(self.graph, treatment, outcome)
        if not ident.identifiable:
            return self._bounds(treatment, outcome, ident)
        if engine == "dml":
            raise NotImplementedError(
                "engine='dml' (cross-fit reuse of the existing DML estimator) is "
                "post-MVP; the default engine='posterior' returns the SCM plug-in "
                "contrast. Use ergodic.estimate_effect(graph, data, X, Y, "
                "strategy='dml') for a DML estimate"
            )
        if engine != "posterior":
            raise ValueError(f"unknown engine {engine!r}; use 'posterior' or 'dml'")
        # The latent-aware do-contrast computes back-door adjustment, the null case,
        # and front-door (the latent absorbs the X-Y confounding, so do() on the
        # treatment severs it from the district latent and the identified value
        # survives). Instruments and the general ID functional are NOT computed by
        # this contrast (decision 3, verify-then-widen): IV rests on an exclusion
        # restriction the latent fit does not impose, and the general functional was
        # not validated to tolerance here, so both raise with a pointer rather than
        # return a confident wrong number.
        if ident.strategy not in ("adjustment", "null", "frontdoor"):
            raise NotImplementedError(
                f"the effect of {treatment!r} on {outcome!r} is identified by the "
                f"{ident.strategy!r} criterion, which the SCM latent-fit do-contrast "
                f"does not yet compute (front-door, adjustment, and the null case do). "
                f"Use estimate_effect(graph, data, {treatment!r}, {outcome!r}), which "
                f"handles {ident.strategy}, or scm.ate(..., engine='dml')"
            )
        low, high = self._ate_levels(treatment, levels)
        draws = _compose.ate_contrast(
            self._order,
            self.mechanisms,
            self._fits,
            treatment,
            outcome,
            low=low,
            high=high,
            n=n,
            seed=self._seed(7),
            latents=self._latent.latents,
        )
        return SCMResult(
            query="ate",
            target=outcome,
            kind="posterior",
            identified=True,
            diagnostics={"strategy": _strategy_code(ident.strategy)},
            _draws=draws,
        )

    def _ate_levels(
        self, treatment: str, levels: tuple[float, float] | None
    ) -> tuple[float, float]:
        """Resolve the ``(low, high)`` contrast for the treatment.

        A continuous or binary treatment defaults to ``(0.0, 1.0)`` (a unit step, or
        the two binary codes). A multi-category treatment has no canonical pair, so
        explicit ``levels`` are required; the two codes are validated against the
        treatment's category count.
        """
        mech = self.mechanisms[treatment]
        is_categorical = isinstance(mech, CategoricalGumbelMax)
        n_cat = self._fits[treatment].n_categories or getattr(mech, "n_categories", 0)
        multi = is_categorical and n_cat > 2
        if levels is None:
            if multi:
                raise ValueError(
                    f"treatment {treatment!r} has {n_cat} categories, which has no default "
                    f"contrast; pass levels=(low_code, high_code) (the two category codes "
                    f"to contrast), as the inference pillar's treated= picks one level"
                )
            return (0.0, 1.0)
        low, high = levels
        if is_categorical:
            for lvl in (low, high):
                v = float(lvl)
                if v != round(v):
                    raise ValueError(
                        f"level {lvl!r} for categorical/binary treatment {treatment!r} is "
                        f"not an integer category code (a fractional level is not snapped); "
                        f"pass integer codes (its codes are 0..{n_cat - 1})"
                    )
                code = round(v)
                if code < 0 or code >= n_cat:
                    raise ValueError(
                        f"level {lvl!r} is not a category code of {treatment!r}; its codes "
                        f"are 0..{n_cat - 1}"
                    )
        return (float(low), float(high))

    def _categorical_result(
        self,
        query: str,
        target: str,
        codes: npt.NDArray[Any],
        *,
        per_unit: npt.NDArray[Any] | None = None,
    ) -> CategoricalResult:
        """Build a :class:`CategoricalResult` from a categorical target's code draws.

        ``codes`` are the integer category draws (pooled over units and posterior
        draws for an interventional/observational marginal, or per-unit for a
        counterfactual). The marginal is the normalized bincount over the node's
        resolved class count, and the category labels come from the data column when
        named. ``per_unit`` carries genuine per-unit codes (a counterfactual) for the
        result's ``per_unit()`` accessor; an interventional marginal leaves it ``None``.
        Every shipped categorical mechanism is invertible, so ``identified=True``.
        """
        mech = self.mechanisms[target]
        n_cat = self._fits[target].n_categories or getattr(mech, "n_categories", 0)
        marginal = _compose.marginal_from_codes(codes, n_cat)
        cats = self.data.categories(target) or tuple(range(n_cat))
        pu = None if per_unit is None else np.asarray(per_unit, dtype=np.int64)
        return CategoricalResult(
            query=query,
            target=target,
            categories=tuple(cats),
            probabilities=marginal,
            identified=True,
            _per_unit=pu,
        )

    def _bounds(self, treatment: str, outcome: str, ident: Any) -> Sensitivity:
        """The honest-bounds return for a non-identified effect.

        A linear-Gaussian treatment-outcome pair gets the exact analytic
        ``b(c) = (cov - c) / var(X)`` curve. A non-linear / non-Gaussian district (a
        categorical/binary treatment or outcome) gets the *latent-prior sweep*: the
        confounded district refits under a grid of confounding-strength priors and the
        range of the resulting do-contrasts is the bound (decision 2). Either way the
        result is a :class:`Sensitivity`, so a naked point read is impossible.
        """
        obstruction = (
            f"no adjustment, front-door, or instrument identifies the effect "
            f"({ident.estimand or 'not a function of the observed data alone'})"
        )
        if self._is_linear_gaussian_pair(treatment, outcome):
            x = np.asarray(self.data.matrix([treatment]).reshape(-1), dtype=np.float64)
            y = np.asarray(self.data.matrix([outcome]).reshape(-1), dtype=np.float64)
            # both at ddof=1 so cov_xy / var_x is the consistent OLS slope (np.cov is
            # ddof=1 by default; np.var defaults to ddof=0, the bug this fixes).
            cov_xy = float(np.cov(x, y)[0, 1])
            var_x = float(np.var(x, ddof=1))
            return linear_sensitivity_curve(
                treatment=treatment,
                outcome=outcome,
                cov_xy=cov_xy,
                var_x=var_x,
                obstruction=obstruction,
                diagnostics={"observational_slope": cov_xy / var_x},
            )
        return self._latent_prior_sweep_bounds(treatment, outcome, obstruction)

    def _is_linear_gaussian_pair(self, treatment: str, outcome: str) -> bool:
        """Whether the analytic ``b(c)`` curve applies: both endpoints Gaussian-additive.

        The closed-form Gaussian-moment derivation of ``b(c)`` needs a continuous
        treatment and outcome. A categorical/binary endpoint breaks it, so that case
        routes to the latent-prior sweep instead.
        """
        return isinstance(self.mechanisms.get(treatment), GaussianAdditive) and isinstance(
            self.mechanisms.get(outcome), GaussianAdditive
        )

    def _latent_prior_sweep_bounds(
        self, treatment: str, outcome: str, obstruction: str
    ) -> Sensitivity:
        """The non-linear non-identified bound: refit the district over a prior grid.

        Refits the confounded district holding everything fixed but the latent loading
        prior sd, and reads the do-contrast at each. The across-prior spread is the
        partial-identification interval (the pressure-test's prior-dependence: the
        answer is whatever the confounding prior says). The treatment must be a
        district endpoint for this to bite; a level pair defaults the same way
        :meth:`ate` does.
        """
        low, high = self._ate_levels(treatment, None)
        nodes = self._latent.node_district.get(treatment, (treatment,))
        order = tuple(n for n in self._order if n in set(nodes))
        latents = self._district_latents(tuple(nodes))
        observed = self._district_observed(order)
        return latent_prior_sweep(
            order,
            {n: self.mechanisms[n] for n in order},
            latents,
            observed,
            treatment=treatment,
            outcome=outcome,
            low=low,
            high=high,
            draws=self.draws,
            tune=self.tune,
            chains=self.chains,
            seed=self._seed(8),
            obstruction=obstruction,
        )

    # -- helpers -------------------------------------------------------------

    def _check_node(self, node: str) -> None:
        if node not in self.graph.nodes:
            nodes = sorted(self.graph.nodes)
            raise ValueError(f"unknown node {node!r}; the SCM's nodes are {nodes}")

    def _check_observable(self, node: str, role: str) -> None:
        """Reject a latent node where a fitted mechanism is required.

        A latent node has no data column and no fitted mechanism (the MVP fits the
        observed margin), so intervening on it or reading its query would fail with
        an opaque ``KeyError`` from deep in the compose/PyMC path. Reject up front.
        """
        self._check_node(node)
        if node in self.graph.latent:
            raise ValueError(
                f"node {node!r} is latent (no data column, no fitted mechanism), so it "
                f"cannot be a {role}; the SCM fits the observed margin. The SCM's "
                f"observed nodes are {list(self.observed)}"
            )

    def _check_do_value(self, node: str, value: DoValue) -> None:
        """Validate a do-target's value against the node's kind.

        A categorical/binary do-target is pinned by its *integer category code*, so a
        fractional value (``do(C=0.5)``) or an out-of-range code (``do(C=99)`` on a
        4-category node) is a user error, not something to silently snap or treat as
        the reference category. This mirrors :meth:`_ate_levels`' code-range check on
        the do-path: ``interventional``/``do``/``counterfactual`` all route through it,
        so neither silently returns a confident wrong answer. A continuous do-target
        takes any float, so nothing is checked. Runs post-fit (the resolved class
        count lives on the fit).
        """
        mech = self.mechanisms[node]
        if not isinstance(mech, CategoricalGumbelMax):
            return
        v = float(value)
        if v != round(v):
            raise ValueError(
                f"do-value {value!r} for categorical/binary node {node!r} is not an "
                f"integer category code; pass an integer code (its codes are 0.."
                f"{(self._fits[node].n_categories or mech.n_categories) - 1}), not a "
                "fractional value (a fractional do-value is not snapped)"
            )
        code = round(v)
        n_cat = self._fits[node].n_categories or mech.n_categories
        if code < 0 or code >= n_cat:
            raise ValueError(
                f"do-value {value!r} is not a category code of {node!r}; its codes are "
                f"0..{n_cat - 1}"
            )

    def _factual_values(self, units: Dataset | Mapping[str, Any]) -> dict[str, npt.NDArray[Any]]:
        """Per-node factual columns for the queried units, as numpy arrays.

        A categorical node is read as integer codes (the Gumbel-max abduction
        needs codes); a continuous node as floats. Only the *observed* nodes (those
        carrying a fitted mechanism) are read: a latent node has no column and is
        abducted implicitly into its district children's residuals (the district
        counterfactual pins the latent at its factual value, the Track-D semantics),
        so iterating ``self.graph.nodes`` would ``KeyError`` on the latent.
        """
        from ..data import Dataset as _Dataset

        out: dict[str, npt.NDArray[Any]] = {}
        for node in self._order:
            mech = self.mechanisms[node]
            if isinstance(units, _Dataset):
                col = self._column_from_dataset(units, node)
            else:
                col = np.asarray(units[node])
            if isinstance(mech, CategoricalGumbelMax):
                out[node] = np.asarray(col, dtype=np.int64)
            else:
                out[node] = np.asarray(col, dtype=np.float64)
        return out

    def _column_from_dataset(self, units: Dataset, node: str) -> npt.NDArray[Any]:
        mech = self.mechanisms[node]
        if isinstance(mech, CategoricalGumbelMax):
            return np.asarray(units.codes(node), dtype=np.int64)
        return np.asarray(units.matrix([node]).reshape(-1), dtype=np.float64)

    def __repr__(self) -> str:
        state = "fitted" if self.fitted else "unfitted"
        return f"SCM({len(self.graph.nodes)} nodes, {state}, n={self.data.n_rows})"

ate(treatment, outcome, *, levels=None, engine='posterior', n=2000)

The average treatment effect, routed through identification first.

identify_effect decides the return shape. Identified by back-door adjustment (or the no-path "null" case): the SCM posterior plug-in contrast E[Y | do(X = high)] - E[Y | do(X = low)] (a posterior :class:~ergodic.scm.SCMResult). Not identified: a :class:~ergodic.scm.Sensitivity (the b(c) curve over the assumed confounding), so a naked point is impossible.

The treatment may be continuous, binary, or categorical. levels is the (low, high) contrast; it defaults to (0.0, 1.0) (the two codes of a binary treatment, or a unit step for a continuous one). A multi-category treatment has no canonical pair, so it requires explicit levels (the two category codes to contrast), mirroring the inference pillar's treated= pattern. The outcome must be numeric (a categorical outcome has no scalar contrast and raises); confounders and mediators may be any type -- a categorical one is sampled in the generative model, so the back-door g-formula marginalizes over it correctly.

Identified by back-door adjustment, the no-path "null" case, or front-door computes the latent-aware do-contrast: a confounded district fits its explicit latent, so the structural mechanism on the causal path fits clean and do severs the treatment from the district latent, leaving the identified value (the front-door mediated effect, the C1 world's 1.2). The instrument strategy still raises with a pointer to :func:~ergodic.inference.estimate_effect: IV identification is linear and rests on the instrument's exclusion restriction, which the latent-district fit does not impose, so its treatment->outcome coefficient stays on the same non-identified ridge as the bow. The general ID functional likewise raises for now (verify-then-widen: its latent-fit do-contrast was not held to estimate_effect to tolerance in this phase). engine="dml" is the cross-fit reuse path (post-MVP; raises with a pointer for now).

Source code in src/ergodic/scm/_scm.py

def ate(
    self,
    treatment: str,
    outcome: str,
    *,
    levels: tuple[float, float] | None = None,
    engine: str = "posterior",
    n: int = 2000,
) -> SCMResult | Sensitivity:
    """The average treatment effect, routed through identification first.

    ``identify_effect`` decides the return shape. Identified *by back-door
    adjustment* (or the no-path ``"null"`` case): the SCM posterior plug-in
    contrast ``E[Y | do(X = high)] - E[Y | do(X = low)]`` (a posterior
    :class:`~ergodic.scm.SCMResult`). Not identified: a
    :class:`~ergodic.scm.Sensitivity` (the ``b(c)`` curve over the assumed
    confounding), so a naked point is impossible.

    The treatment may be continuous, binary, or categorical. ``levels`` is the
    ``(low, high)`` contrast; it defaults to ``(0.0, 1.0)`` (the two codes of a
    binary treatment, or a unit step for a continuous one). A multi-category
    treatment has no canonical pair, so it *requires* explicit ``levels`` (the
    two category codes to contrast), mirroring the inference pillar's ``treated=``
    pattern. The outcome must be numeric (a categorical outcome has no scalar
    contrast and raises); confounders and mediators may be any type -- a
    categorical one is sampled in the generative model, so the back-door
    g-formula marginalizes over it correctly.

    Identified by back-door adjustment, the no-path ``"null"`` case, or
    **front-door** computes the latent-aware do-contrast: a confounded district
    fits its explicit latent, so the structural mechanism on the causal path fits
    clean and ``do`` severs the treatment from the district latent, leaving the
    identified value (the front-door mediated effect, the C1 world's 1.2). The
    **instrument** strategy still raises with a pointer to
    :func:`~ergodic.inference.estimate_effect`: IV identification is linear and
    rests on the instrument's exclusion restriction, which the latent-district fit
    does not impose, so its treatment->outcome coefficient stays on the same
    non-identified ridge as the bow. The **general ID functional** likewise
    raises for now (verify-then-widen: its latent-fit do-contrast was not held to
    ``estimate_effect`` to tolerance in this phase). ``engine="dml"`` is the
    cross-fit reuse path (post-MVP; raises with a pointer for now).
    """
    self._require_fit()
    self._check_node(treatment)
    self._check_node(outcome)
    # the outcome must be numeric: a categorical contrast is not a scalar.
    if isinstance(self.mechanisms.get(outcome), CategoricalGumbelMax):
        raise ValueError(
            f"ate needs a numeric outcome; {outcome!r} is categorical/binary, which "
            f"has no scalar contrast. Read its post-intervention distribution with "
            f"interventional({outcome!r}, do={{{treatment!r}: ...}}) instead"
        )
    ident = identify_effect(self.graph, treatment, outcome)
    if not ident.identifiable:
        return self._bounds(treatment, outcome, ident)
    if engine == "dml":
        raise NotImplementedError(
            "engine='dml' (cross-fit reuse of the existing DML estimator) is "
            "post-MVP; the default engine='posterior' returns the SCM plug-in "
            "contrast. Use ergodic.estimate_effect(graph, data, X, Y, "
            "strategy='dml') for a DML estimate"
        )
    if engine != "posterior":
        raise ValueError(f"unknown engine {engine!r}; use 'posterior' or 'dml'")
    # The latent-aware do-contrast computes back-door adjustment, the null case,
    # and front-door (the latent absorbs the X-Y confounding, so do() on the
    # treatment severs it from the district latent and the identified value
    # survives). Instruments and the general ID functional are NOT computed by
    # this contrast (decision 3, verify-then-widen): IV rests on an exclusion
    # restriction the latent fit does not impose, and the general functional was
    # not validated to tolerance here, so both raise with a pointer rather than
    # return a confident wrong number.
    if ident.strategy not in ("adjustment", "null", "frontdoor"):
        raise NotImplementedError(
            f"the effect of {treatment!r} on {outcome!r} is identified by the "
            f"{ident.strategy!r} criterion, which the SCM latent-fit do-contrast "
            f"does not yet compute (front-door, adjustment, and the null case do). "
            f"Use estimate_effect(graph, data, {treatment!r}, {outcome!r}), which "
            f"handles {ident.strategy}, or scm.ate(..., engine='dml')"
        )
    low, high = self._ate_levels(treatment, levels)
    draws = _compose.ate_contrast(
        self._order,
        self.mechanisms,
        self._fits,
        treatment,
        outcome,
        low=low,
        high=high,
        n=n,
        seed=self._seed(7),
        latents=self._latent.latents,
    )
    return SCMResult(
        query="ate",
        target=outcome,
        kind="posterior",
        identified=True,
        diagnostics={"strategy": _strategy_code(ident.strategy)},
        _draws=draws,
    )

counterfactual(units, do, target, *, use_draws=True)

The twin-network counterfactual for the given units under do.

Abduct each unit's structural noise on the factual graph, then predict target on the action branch reusing it. units is a dataset (or a mapping of node name to a column of factual values) holding every node's factual value per unit. For invertible mechanisms the counterfactual is exact (use_draws=False collapses to the deterministic point; the default carries parameter uncertainty). A categorical target returns a :class:~ergodic.scm.CategoricalResult (the answer is a distribution).

Source code in src/ergodic/scm/_scm.py

def counterfactual(
    self,
    units: Dataset | Mapping[str, Any],
    do: Mapping[str, float],
    target: str,
    *,
    use_draws: bool = True,
) -> SCMResult | CategoricalResult:
    """The twin-network counterfactual for the given units under ``do``.

    Abduct each unit's structural noise on the factual graph, then predict
    ``target`` on the action branch reusing it. ``units`` is a dataset (or a
    mapping of node name to a column of factual values) holding every node's
    factual value per unit. For invertible mechanisms the counterfactual is
    exact (``use_draws=False`` collapses to the deterministic point; the
    default carries parameter uncertainty). A categorical ``target`` returns a
    :class:`~ergodic.scm.CategoricalResult` (the answer is a distribution).
    """
    self._require_fit()
    self._check_observable(target, "counterfactual target")
    for node in do:
        self._check_observable(node, "do-target")
        self._check_do_value(node, do[node])
    factual = self._factual_values(units)

    mech = self.mechanisms[target]
    if isinstance(mech, CategoricalGumbelMax):
        codes, _ = _compose.counterfactual_categorical(
            self.mechanisms,
            self._fits,
            factual,
            do,
            target,
            parents_of_target=mech.parents,
            seed=self._seed(5),
        )
        # the abducted-noise counterfactual is exact (invertible), so identified.
        return self._categorical_result("counterfactual", target, codes, per_unit=codes)

    if not isinstance(mech, GaussianAdditive):
        raise TypeError(f"no counterfactual for mechanism {type(mech).__name__} on {target!r}")
    # The twin-network counterfactual abducts the exact structural noise, which
    # is only valid for invertible mechanisms. A non-invertible (probabilistic-
    # only) node would need the genuinely wider abduction posterior, which is not
    # in the MVP; refuse rather than reuse the exact residual and mislabel the
    # result. Every shipped mechanism is invertible, so this guards a future
    # mechanism, not today's worlds.
    non_invertible = sorted(n for n, m in self.mechanisms.items() if not _is_invertible(m))
    if non_invertible:
        raise NotImplementedError(
            f"node(s) {non_invertible} use a non-invertible (probabilistic-only) "
            "mechanism; non-invertible-mechanism counterfactuals (the genuinely "
            "wider abduction posterior) are not in the MVP. Every built-in "
            "mechanism is invertible, so this only arises for a custom mechanism"
        )
    draws = _compose.counterfactual_continuous(
        self._order,
        self.mechanisms,
        self._fits,
        factual,
        do,
        target,
        use_draws=use_draws,
        seed=self._seed(6),
    )
    return SCMResult(
        query="counterfactual",
        target=target,
        kind="posterior",
        identified=True,
        _draws=draws,
    )

do(intervention)

A lightweight intervened view of the SCM for chaining.

scm.do({"X": 1.0}).interventional("Y") reads the post-intervention distribution. The view shares the fit; it does not refit. A continuous do-target takes a float; a categorical/binary do-target takes its integer category code.

Source code in src/ergodic/scm/_scm.py

def do(self, intervention: Mapping[str, DoValue]) -> InterventionView:
    """A lightweight intervened view of the SCM for chaining.

    ``scm.do({"X": 1.0}).interventional("Y")`` reads the post-intervention
    distribution. The view shares the fit; it does not refit. A continuous
    do-target takes a float; a categorical/binary do-target takes its integer
    category code.
    """
    for node in intervention:
        self._check_observable(node, "do-target")
        # validate a categorical code early when the SCM is fitted (the resolved
        # class count lives on the fit); an unfitted chainable view still builds,
        # and interventional() re-validates before it samples regardless.
        if self.fitted:
            self._check_do_value(node, intervention[node])
    return InterventionView(self, dict(intervention))

fit(*, parallel=False)

Fit every node's mechanism, per node for singletons, jointly per district.

A Markovian singleton district fits per node (cheap, exact, the MVP path); a confounded district (nodes sharing a <->) does not factorize -- the latent couples it -- so it fits as one joint NUTS block with the explicit latents, the gauge, a non-centered parameterization, and a high target_accept (the Track-D recipe). The composed posterior merges the two the same way it merges per-node fits.

Sequential by default (robust, CI-safe). parallel=True spreads the singleton nodes over a process pool with the BLAS/OMP threads pinned in the workers; a confounded district's joint fit stays in the main process (its chains already parallelize and its model is one block). Returns self.

Source code in src/ergodic/scm/_scm.py

def fit(self, *, parallel: bool = False) -> SCM:
    """Fit every node's mechanism, per node for singletons, jointly per district.

    A Markovian singleton district fits per node (cheap, exact, the MVP path); a
    confounded district (nodes sharing a ``<->``) does not factorize -- the latent
    couples it -- so it fits as one joint NUTS block with the explicit latents,
    the gauge, a non-centered parameterization, and a high ``target_accept`` (the
    Track-D recipe). The composed posterior merges the two the same way it merges
    per-node fits.

    Sequential by default (robust, CI-safe). ``parallel=True`` spreads the
    *singleton* nodes over a process pool with the BLAS/OMP threads pinned in the
    workers; a confounded district's joint fit stays in the main process (its
    chains already parallelize and its model is one block). Returns ``self``.
    """
    seed = 0 if self.random_state is None else self.random_state
    singletons = set(self._latent.singletons)
    single_mechs = {n: m for n, m in self.mechanisms.items() if n in singletons}
    if parallel:
        fits = self._fit_parallel(single_mechs, seed)
    else:
        fits = {
            node: mech.fit(
                self.data, draws=self.draws, tune=self.tune, chains=self.chains, seed=seed
            )
            for node, mech in single_mechs.items()
        }
    for offset, nodes in enumerate(self._latent.confounded_districts):
        fits.update(self._fit_district(nodes, seed + offset))
    self._fits = fits
    return self

interventional(target, do=None, *, n=2000)

E[Y | do(X = x)] and its posterior, through pm.do.

do maps each intervened node to its value (broadcast to a length-n vector internally, the verified shape gotcha; an integer code for a categorical/binary do-target, a float for a continuous one). target must not be an intervened node. A categorical target returns a :class:CategoricalResult (the post-intervention per-category marginal); a continuous target an :class:SCMResult. Carries parameter-posterior uncertainty via the merged idata.

Source code in src/ergodic/scm/_scm.py

def interventional(
    self, target: str, do: Mapping[str, DoValue] | None = None, *, n: int = 2000
) -> SCMResult | CategoricalResult:
    """``E[Y | do(X = x)]`` and its posterior, through ``pm.do``.

    ``do`` maps each intervened node to its value (broadcast to a length-n
    vector internally, the verified shape gotcha; an integer code for a
    categorical/binary do-target, a float for a continuous one). ``target`` must
    not be an intervened node. A categorical target returns a
    :class:`CategoricalResult` (the post-intervention per-category marginal); a
    continuous target an :class:`SCMResult`. Carries parameter-posterior
    uncertainty via the merged ``idata``.
    """
    self._require_fit()
    self._check_observable(target, "query target")
    intervention = dict(do or {})
    for node in intervention:
        self._check_observable(node, "do-target")
        self._check_do_value(node, intervention[node])
    draws = _compose.interventional_draws(
        self._order,
        self.mechanisms,
        self._fits,
        intervention,
        target,
        n=n,
        seed=self._seed(3),
        latents=self._latent.latents,
    )
    if isinstance(self.mechanisms[target], CategoricalGumbelMax):
        return self._categorical_result("interventional", target, draws)
    return SCMResult(
        query="interventional", target=target, kind="posterior", identified=True, _draws=draws
    )

ite(units, treatment, outcome, *, levels=None, shift=None)

Per-unit individual treatment effects, through the twin network.

For each unit, the exact counterfactual contrast of outcome under two treatment worlds that share that unit's abducted noise (the same-person, two-worlds twin network). Returns one effect per unit, ordered as units -- the structural-model sibling of a CATE estimator's :meth:~ergodic.inference.OutcomeRegression.cate, which returns a conditional average over the noise rather than the individual's own counterfactual.

Two framings of the contrast, and you pick one:

• an absolute level contrast (the default): Y | do(X=high) - Y | do(X=low), both worlds pinning the treatment to a fixed level for every unit, the very contrast :meth:ate averages. Pass levels=(low, high); it defaults to (0.0, 1.0) (a unit step, or the two codes of a binary treatment). Works for a continuous, binary, or (with explicit codes) categorical treatment, and for a continuous or binary outcome -- a binary outcome's per-unit effect is its counterfactual flip, in {-1, 0, 1}.
• a shift (a "shock"): Y | do(X = x_i + shift) - Y_factual, perturbing each unit's treatment from its own factual value, against the world where nothing was done. Pass shift=delta (a continuous treatment and outcome).

The two coincide for a linear-additive mechanism: the shift contrast is the constant coef * delta for every unit, equal to the matching level step. They part only once a mechanism is nonlinear, where the shift is the unit's local response at its own baseline. So with the built-in linear-Gaussian mechanisms the continuous ITE is constant across units (the coefficient), and genuine per-unit spread comes from a binary or categorical outcome, whose logistic response varies with each unit's baseline.

Each effect is the exact point (invertible mechanisms). For an individual's posterior interval, call :meth:counterfactual on that one unit per arm with use_draws=True and difference the draws.

Source code in src/ergodic/scm/_scm.py

def ite(
    self,
    units: Dataset | Mapping[str, Any],
    treatment: str,
    outcome: str,
    *,
    levels: tuple[float, float] | None = None,
    shift: float | None = None,
) -> ITEResult:
    """Per-unit individual treatment effects, through the twin network.

    For each unit, the exact counterfactual contrast of ``outcome`` under two
    treatment worlds that share *that unit's* abducted noise (the same-person,
    two-worlds twin network). Returns one effect per unit, ordered as ``units`` --
    the structural-model sibling of a CATE estimator's
    :meth:`~ergodic.inference.OutcomeRegression.cate`, which returns a *conditional
    average* over the noise rather than the individual's own counterfactual.

    Two framings of the contrast, and you pick one:

    * an absolute **level** contrast (the default): ``Y | do(X=high) -
      Y | do(X=low)``, both worlds pinning the treatment to a fixed level for every
      unit, the very contrast :meth:`ate` averages. Pass ``levels=(low, high)``; it
      defaults to ``(0.0, 1.0)`` (a unit step, or the two codes of a binary
      treatment). Works for a continuous, binary, or (with explicit codes)
      categorical treatment, and for a continuous or binary outcome -- a binary
      outcome's per-unit effect is its counterfactual flip, in ``{-1, 0, 1}``.
    * a **shift** (a "shock"): ``Y | do(X = x_i + shift) - Y_factual``, perturbing
      each unit's treatment from its *own* factual value, against the world where
      nothing was done. Pass ``shift=delta`` (a continuous treatment and outcome).

    The two coincide for a linear-additive mechanism: the shift contrast is the
    constant ``coef * delta`` for every unit, equal to the matching level step. They
    part only once a mechanism is nonlinear, where the shift is the unit's local
    response at its own baseline. So with the built-in linear-Gaussian mechanisms
    the continuous ITE is constant across units (the coefficient), and genuine
    per-unit spread comes from a **binary or categorical outcome**, whose logistic
    response varies with each unit's baseline.

    Each effect is the exact point (invertible mechanisms). For an individual's
    posterior *interval*, call :meth:`counterfactual` on that one unit per arm with
    ``use_draws=True`` and difference the draws.
    """
    from ..inference import ITEResult

    self._require_fit()
    self._check_observable(treatment, "treatment")
    self._check_observable(outcome, "outcome")
    if levels is not None and shift is not None:
        raise ValueError(
            "pass either levels= or shift=, not both (they are two framings of the "
            "same per-unit contrast)"
        )
    out_mech = self.mechanisms[outcome]
    if isinstance(out_mech, CategoricalGumbelMax):
        n_cat = self._fits[outcome].n_categories or out_mech.n_categories
        if n_cat > 2:
            raise ValueError(
                f"ite needs a numeric outcome; {outcome!r} has {n_cat} categories, "
                f"which has no scalar individual effect. Read per-unit counterfactual "
                f"distributions with counterfactual(units, do, {outcome!r}) instead"
            )
    if shift is not None:
        values = self._ite_shift(units, treatment, outcome, shift)
        framing = f"shift={shift}"
        note = (
            "structural individual treatment effects: per-unit "
            f"Y|do({treatment}=x_i+{shift}) - Y_factual through the twin network "
            "(abduct the unit's noise, predict on the shifted branch)"
        )
    else:
        low, high = self._ate_levels(treatment, levels)
        values = self._ite_levels(units, treatment, outcome, low, high)
        framing = f"levels=({low}, {high})"
        note = (
            "structural individual treatment effects: per-unit "
            f"Y|do({treatment}={high}) - Y|do({treatment}={low}) on shared abducted "
            "noise (the same-person, two-worlds twin network)"
        )
    factual = self._factual_values(units)
    out_col = np.asarray(factual.get(outcome), dtype=np.float64)
    n = int(values.shape[0])
    return ITEResult(
        treatment=treatment,
        outcome=outcome,
        covariates=(),
        n=n,
        kind="point",
        source="scm",
        framing=framing,
        note=note,
        _values=np.asarray(values, dtype=np.float64),
        _index=np.arange(n, dtype=np.int64),
        _factual=out_col,
    )

marginal(target, n=1000)

Observational posterior-predictive draws of one node.

A continuous node returns an :class:SCMResult; a categorical/binary node a :class:CategoricalResult (the per-category marginal), exactly as :meth:counterfactual does for a categorical target.

Source code in src/ergodic/scm/_scm.py

def marginal(self, target: str, n: int = 1000) -> SCMResult | CategoricalResult:
    """Observational posterior-predictive draws of one node.

    A continuous node returns an :class:`SCMResult`; a categorical/binary node a
    :class:`CategoricalResult` (the per-category marginal), exactly as
    :meth:`counterfactual` does for a categorical target.
    """
    self._require_fit()
    self._check_observable(target, "query target")
    return self._observational(target, n=n, seed=self._seed(2))

sample(n=1000)

Observational draws of every node (read the target with :meth:marginal).

Returns the posterior predictive of the generative model over n units; the result carries the draws of the last node in topological order (a sink) as a summary. A categorical sink returns a :class:CategoricalResult (the per-category marginal). Use :meth:marginal to read a specific node, or :meth:interventional / :meth:counterfactual for a query on one.

Source code in src/ergodic/scm/_scm.py

def sample(self, n: int = 1000) -> SCMResult | CategoricalResult:
    """Observational draws of every node (read the target with :meth:`marginal`).

    Returns the posterior predictive of the generative model over ``n`` units;
    the result carries the draws of the last node in topological order (a sink)
    as a summary. A categorical sink returns a :class:`CategoricalResult` (the
    per-category marginal). Use :meth:`marginal` to read a specific node, or
    :meth:`interventional` / :meth:`counterfactual` for a query on one.
    """
    self._require_fit()
    target = self._order[-1]
    return self._observational(target, n=n, seed=self._seed(1))

ergodic.scm.SCMResult dataclass

A posterior over a query's answer.

query names what was asked ("sample", "interventional", "counterfactual", "ate"). target names the node read. _draws holds the posterior draws (flattened); the accessors read them. identified records whether the answer is a function of the data alone.

Source code in src/ergodic/scm/_result.py

@dataclass(frozen=True)
class SCMResult:
    """A posterior over a query's answer.

    ``query`` names what was asked (``"sample"``, ``"interventional"``,
    ``"counterfactual"``, ``"ate"``). ``target`` names the node read.
    ``_draws`` holds the posterior draws (flattened); the accessors read them.
    ``identified`` records whether the answer is a function of the data alone.
    """

    query: str
    target: str
    kind: str  # "posterior"
    identified: bool
    diagnostics: dict[str, float] = field(default_factory=dict)
    _draws: npt.NDArray[np.float64] = field(repr=False, default_factory=lambda: np.empty(0))

    @property
    def value(self) -> float:
        """The mean of the draws (the point summary)."""
        if self._draws.size == 0:
            raise ValueError("result carries no draws")
        return float(np.mean(self._draws))

    def mean(self) -> float:
        """The posterior mean of the answer."""
        return self.value

    @property
    def std_error(self) -> float | None:
        """The posterior standard deviation of the answer."""
        if self._draws.size < 2:
            return None
        return float(np.std(self._draws, ddof=1))

    def draws(self) -> npt.NDArray[np.float64]:
        """The posterior draws of the answer."""
        return np.array(self._draws, copy=True)

    def interval(self, level: float = 0.95) -> tuple[float, float]:
        """A credible interval at ``level``, from the draw quantiles."""
        if not 0.0 < level < 1.0:
            raise ValueError(f"level must be in (0, 1), got {level}")
        if self._draws.size == 0:
            raise ValueError("result carries no draws to form an interval from")
        lo = float(np.quantile(self._draws, (1.0 - level) / 2.0))
        hi = float(np.quantile(self._draws, (1.0 + level) / 2.0))
        return (lo, hi)

    def __repr__(self) -> str:
        flag = "identified" if self.identified else "not identified"
        se = self.std_error
        body = f"{self.target}={self.value:.4g}"
        if se is not None:
            lo, hi = self.interval()
            body += f", sd={se:.3g}, 95% CrI [{lo:.3g}, {hi:.3g}]"
        return f"SCMResult({self.query}, {flag}: {body})"

std_error property

The posterior standard deviation of the answer.

value property

The mean of the draws (the point summary).

draws()

The posterior draws of the answer.

Source code in src/ergodic/scm/_result.py

def draws(self) -> npt.NDArray[np.float64]:
    """The posterior draws of the answer."""
    return np.array(self._draws, copy=True)

interval(level=0.95)

A credible interval at level, from the draw quantiles.

Source code in src/ergodic/scm/_result.py

def interval(self, level: float = 0.95) -> tuple[float, float]:
    """A credible interval at ``level``, from the draw quantiles."""
    if not 0.0 < level < 1.0:
        raise ValueError(f"level must be in (0, 1), got {level}")
    if self._draws.size == 0:
        raise ValueError("result carries no draws to form an interval from")
    lo = float(np.quantile(self._draws, (1.0 - level) / 2.0))
    hi = float(np.quantile(self._draws, (1.0 + level) / 2.0))
    return (lo, hi)

mean()

The posterior mean of the answer.

Source code in src/ergodic/scm/_result.py

def mean(self) -> float:
    """The posterior mean of the answer."""
    return self.value

ergodic.scm.CategoricalResult dataclass

A posterior over a categorical query's answer.

A categorical node's interventional or counterfactual answer is a distribution over its categories, not a single number, so this carries the per-category probabilities (probabilities, summing to one over categories) and the per-unit draws when a counterfactual was asked of specific units. identified carries the same honesty flag as :class:SCMResult.

Source code in src/ergodic/scm/_result.py

@dataclass(frozen=True)
class CategoricalResult:
    """A posterior over a categorical query's answer.

    A categorical node's interventional or counterfactual answer is a
    distribution over its categories, not a single number, so this carries the
    per-category probabilities (``probabilities``, summing to one over
    ``categories``) and the per-unit draws when a counterfactual was asked of
    specific units. ``identified`` carries the same honesty flag as
    :class:`SCMResult`.
    """

    query: str
    target: str
    categories: tuple[object, ...]
    probabilities: npt.NDArray[np.float64]
    identified: bool
    diagnostics: dict[str, float] = field(default_factory=dict)
    _per_unit: npt.NDArray[np.int64] | None = field(repr=False, default=None)

    def mode(self) -> object:
        """The most probable category."""
        return self.categories[int(np.argmax(self.probabilities))]

    def per_unit(self) -> npt.NDArray[np.int64] | None:
        """Per-unit counterfactual category codes, when units were queried."""
        if self._per_unit is None:
            return None
        return np.array(self._per_unit, copy=True)

    def __repr__(self) -> str:
        flag = "identified" if self.identified else "not identified"
        probs = ", ".join(
            f"{cat}={p:.3g}" for cat, p in zip(self.categories, self.probabilities, strict=False)
        )
        return f"CategoricalResult({self.query}, {flag}: {self.target} [{probs}])"

mode()

The most probable category.

Source code in src/ergodic/scm/_result.py

def mode(self) -> object:
    """The most probable category."""
    return self.categories[int(np.argmax(self.probabilities))]

per_unit()

Per-unit counterfactual category codes, when units were queried.

Source code in src/ergodic/scm/_result.py

def per_unit(self) -> npt.NDArray[np.int64] | None:
    """Per-unit counterfactual category codes, when units were queried."""
    if self._per_unit is None:
        return None
    return np.array(self._per_unit, copy=True)

ergodic.scm.Sensitivity dataclass

The effect as a function of the assumed confounding strength.

Returned in place of a point when an effect is not identified. confounding is the grid of assumed confounding-channel values c (a covariance on the treatment-outcome pair), effect the implied causal effect b(c) at each. identified_at_zero is b(0), the answer under the no-confounding assumption (the observational slope), which is where the curve meets the identified region. obstruction is the human-readable reason (the hedge).

There is no scalar accessor on purpose: a non-identified effect is a range, not a number. Read :meth:range, :meth:at, or the arrays.

Source code in src/ergodic/scm/_bounds.py

@dataclass(frozen=True)
class Sensitivity:
    """The effect as a function of the assumed confounding strength.

    Returned in place of a point when an effect is not identified. ``confounding``
    is the grid of assumed confounding-channel values ``c`` (a covariance on the
    treatment-outcome pair), ``effect`` the implied causal effect ``b(c)`` at each.
    ``identified_at_zero`` is ``b(0)``, the answer under the no-confounding
    assumption (the observational slope), which is where the curve meets the
    identified region. ``obstruction`` is the human-readable reason (the hedge).

    There is no scalar accessor on purpose: a non-identified effect is a range, not
    a number. Read :meth:`range`, :meth:`at`, or the arrays.
    """

    treatment: str
    outcome: str
    confounding: npt.NDArray[np.float64]
    effect: npt.NDArray[np.float64]
    identified_at_zero: float
    obstruction: str
    diagnostics: dict[str, float] = field(default_factory=dict)

    @property
    def identified(self) -> bool:
        """Always ``False``: a :class:`Sensitivity` exists because the effect is not."""
        return False

    def range(self) -> tuple[float, float]:
        """The (min, max) implied effect over the swept confounding grid."""
        if self.effect.size == 0:
            raise ValueError("sensitivity curve is empty")
        return (float(np.min(self.effect)), float(np.max(self.effect)))

    def at(self, c: float) -> float:
        """The implied causal effect if the confounding channel strength is ``c``.

        Linear interpolation on the swept grid; ``c`` outside the grid extrapolates
        along the line, which is exact for the linear-Gaussian case.
        """
        return float(np.interp(c, self.confounding, self.effect))

    def __repr__(self) -> str:
        lo, hi = self.range()
        return (
            f"Sensitivity({self.treatment} -> {self.outcome}: effect in "
            f"[{lo:.3g}, {hi:.3g}] over assumed confounding, "
            f"b(0)={self.identified_at_zero:.3g}; not identified: {self.obstruction})"
        )

identified property

Always False: a :class:Sensitivity exists because the effect is not.

at(c)

The implied causal effect if the confounding channel strength is c.

Linear interpolation on the swept grid; c outside the grid extrapolates along the line, which is exact for the linear-Gaussian case.

Source code in src/ergodic/scm/_bounds.py

def at(self, c: float) -> float:
    """The implied causal effect if the confounding channel strength is ``c``.

    Linear interpolation on the swept grid; ``c`` outside the grid extrapolates
    along the line, which is exact for the linear-Gaussian case.
    """
    return float(np.interp(c, self.confounding, self.effect))

range()

The (min, max) implied effect over the swept confounding grid.

Source code in src/ergodic/scm/_bounds.py

def range(self) -> tuple[float, float]:
    """The (min, max) implied effect over the swept confounding grid."""
    if self.effect.size == 0:
        raise ValueError("sensitivity curve is empty")
    return (float(np.min(self.effect)), float(np.max(self.effect)))

ergodic.scm.GaussianAdditive dataclass

Additive-noise Gaussian mechanism, X = b0 + w . PA + U, U ~ N(0, sigma).

The default continuous mechanism: a Bayesian-linear mean over the parents with a half-normal prior on the noise scale. Invertible, so its counterfactual is the exact residual U = X - (b0 + w . PA) reused under the action. A :class:~ergodic.inference.Learner-as-mean override (carrying no parameter posterior, abducting via the residual against the fitted mean) is planned post-MVP; the MVP fits the Bayesian-linear mean.

In a confounded district the node may also carry incident latents: one per-unit L ~ Normal(0, 1) per bidirected edge the node touches, entering the mean as + loading_L * L. latents lists those latent RV names (in a deterministic order); gauged_latents is the subset whose loading is pinned positive (a HalfNormal) to kill the L -> -L reflection. A Markovian node leaves both empty and reduces to the MVP exactly.

Source code in src/ergodic/scm/_mechanism.py

@dataclass(frozen=True)
class GaussianAdditive:
    """Additive-noise Gaussian mechanism, ``X = b0 + w . PA + U``, ``U ~ N(0, sigma)``.

    The default continuous mechanism: a Bayesian-linear mean over the parents with
    a half-normal prior on the noise scale. Invertible, so its counterfactual is
    the exact residual ``U = X - (b0 + w . PA)`` reused under the action. A
    :class:`~ergodic.inference.Learner`-as-mean override (carrying no parameter
    posterior, abducting via the residual against the fitted mean) is planned
    post-MVP; the MVP fits the Bayesian-linear mean.

    In a confounded district the node may also carry incident *latents*: one
    per-unit ``L ~ Normal(0, 1)`` per bidirected edge the node touches, entering
    the mean as ``+ loading_L * L``. ``latents`` lists those latent RV names (in a
    deterministic order); ``gauged_latents`` is the subset whose loading is pinned
    positive (a ``HalfNormal``) to kill the ``L -> -L`` reflection. A Markovian node
    leaves both empty and reduces to the MVP exactly.
    """

    name: str
    parents: tuple[str, ...]
    parent_cards: tuple[int | None, ...] = ()
    coef_prior_sd: float = 5.0
    sigma_prior_sd: float = 2.0
    latents: tuple[str, ...] = ()
    gauged_latents: tuple[str, ...] = ()
    loading_prior_sd: float = 2.0

    @property
    def capability(self) -> Capability:
        return Capability.INVERTIBLE

    @property
    def coef_name(self) -> str:
        return f"{self.name}__coef"

    @property
    def intercept_name(self) -> str:
        return f"{self.name}__intercept"

    @property
    def sigma_name(self) -> str:
        return f"{self.name}__sigma"

    def loading_name(self, latent: str) -> str:
        """The parameter name for this node's loading on latent ``L`` (district fit)."""
        return f"{self.name}__loading__{latent}"

    @property
    def _design_width(self) -> int:
        """The encoded coefficient count, ``len(parents)`` when every parent is
        continuous (no ``parent_cards``), else the one-hot-expanded width."""
        if not self.parents:
            return 0
        if not self.parent_cards:
            return len(self.parents)
        return encoded_width(self.parent_cards)

    def build(self, x: Array, y: Array) -> Any:
        """The per-node PyMC model: ``y ~ Normal(b0 + x_encoded . w, sigma)``.

        ``x`` is the raw parent matrix (codes for discrete parents); it is encoded
        to the one-hot design here, the same encoding the generative model uses.
        """
        import pymc as pm

        x_enc = encode_parent_design(x, self.parent_cards) if self.parents else x
        _, p = x_enc.shape
        with pm.Model() as model:
            x_data = pm.Data("__parents", x_enc) if p else None
            intercept = pm.Normal(self.intercept_name, 0.0, self.coef_prior_sd)
            mean = intercept
            if p:
                coef = pm.Normal(self.coef_name, 0.0, self.coef_prior_sd, shape=p)
                mean = mean + pm.math.dot(x_data, coef)
            sigma = pm.HalfNormal(self.sigma_name, self.sigma_prior_sd)
            pm.Normal(self.name, mu=mean, sigma=sigma, observed=y)
        return model

    def generative_rv(self, nodes: Mapping[str, Any], n: int) -> Any:
        """Add this node's RV to the active generative model: ``Normal(b0 + PA . w, sigma)``.

        ``nodes`` maps each (already-built) parent name to its RV. Parents are
        encoded by the shared pytensor encoder, so the design matches the fit view
        column for column; the RV names match the merged posterior exactly, so
        ``sample_posterior_predictive`` reads the fitted draws. An incident latent
        ``L`` (a per-unit ``Normal(0, 1)`` the compose helper builds first, found in
        ``nodes``) adds ``loading_L * L`` to the mean, so ``pm.do`` on this node
        severs it from ``L`` while a district sibling still reads ``L`` -- the
        confounding cancels in an interventional contrast.
        """
        import pymc as pm
        import pytensor.tensor as pt

        intercept = pm.Normal(self.intercept_name, 0.0, self.coef_prior_sd)
        mean: Any = intercept * pt.ones((n,))
        if self.parents:
            design = encode_parents_pt([nodes[p] for p in self.parents], self.parent_cards)
            coef = pm.Normal(self.coef_name, 0.0, self.coef_prior_sd, shape=self._design_width)
            mean = mean + pm.math.dot(design, coef)
        for lat in self.latents:
            loading = self._loading_rv(lat)
            mean = mean + loading * nodes[lat]
        sigma = pm.HalfNormal(self.sigma_name, self.sigma_prior_sd)
        return pm.Normal(self.name, mu=mean, sigma=sigma, shape=(n,))

    def _loading_rv(self, latent: str) -> Any:
        """The loading RV for ``latent``: ``HalfNormal`` if gauged, else ``Normal``.

        The gauge (one loading per latent pinned positive) lives here so the fit
        view and the generative view declare an identical prior for the same name,
        which is what keeps ``sample_posterior_predictive`` reading the fitted draws
        rather than re-sampling a differently-shaped prior.
        """
        import pymc as pm

        name = self.loading_name(latent)
        if latent in self.gauged_latents:
            return pm.HalfNormal(name, self.loading_prior_sd)
        return pm.Normal(name, 0.0, self.loading_prior_sd)

    def fit(self, data: Dataset, *, draws: int, tune: int, chains: int, seed: int) -> FitResult:
        import pymc as pm

        x = _parent_matrix(data, self.parents)
        y = np.asarray(data.matrix([self.name]).reshape(-1), dtype=np.float64)
        model = self.build(x, y)
        with model:
            idata = pm.sample(
                draws=draws,
                tune=tune,
                chains=chains,
                random_seed=seed,
                progressbar=False,
                compute_convergence_checks=False,
            )
        post = idata.posterior
        params: dict[str, Array] = {self.intercept_name: _flatten(post[self.intercept_name])}
        params[self.sigma_name] = _flatten(post[self.sigma_name])
        if self.parents:
            params[self.coef_name] = _flatten(post[self.coef_name])
        means = {k: np.asarray(v.mean(axis=0), dtype=np.float64) for k, v in params.items()}
        return FitResult(params=params, param_means=means)

    def structural_mean(self, fit: FitResult, parent_values: Array) -> Array:
        """The mean ``b0 + w . encode(PA)`` at the posterior mean.

        ``parent_values`` is the ``(n, p)`` *raw* parent matrix (codes for discrete
        parents); it is one-hot encoded internally, so the coefficients line up with
        the fit. ``(n, p)`` in, ``(n,)`` out.
        """
        intercept = float(np.asarray(fit.param_means[self.intercept_name]))
        pv = np.asarray(parent_values, dtype=np.float64)
        if not self.parents:
            return np.full(pv.shape[0], intercept)
        design = encode_parent_design(pv, self.parent_cards)
        coef = np.asarray(fit.param_means[self.coef_name], dtype=np.float64)
        return intercept + design @ coef

    def mean_per_draw(self, fit: FitResult, parent_values: Array) -> Array:
        """The mean per parameter-posterior draw.

        ``parent_values`` is the ``(S, n, p)`` *raw* per-draw parent matrix (a
        non-intervened parent carries its own counterfactual posterior; codes for a
        discrete parent). It is one-hot encoded per draw, so the encoded design
        matches the ``(S, P)`` coefficient draws. Returns ``(S, n)``. With no parents
        the result is the broadcast intercept draws.
        """
        intercept = np.asarray(fit.params[self.intercept_name], dtype=np.float64)  # (S,)
        pv = np.asarray(parent_values, dtype=np.float64)
        s = intercept.shape[0]
        n = pv.shape[1] if pv.ndim == 3 else pv.shape[0]
        if not self.parents:
            return np.repeat(intercept[:, None], n, axis=1)
        coef = np.asarray(fit.params[self.coef_name], dtype=np.float64)  # (S, P)
        # broadcast an intervened parent's singleton draw axis up to S, then encode
        # each draw's raw (n, p) slice to its (n, P) one-hot design.
        if pv.shape[0] != s:
            pv = np.broadcast_to(pv, (s, n, pv.shape[2]))
        design = np.stack(
            [encode_parent_design(pv[d], self.parent_cards) for d in range(s)], axis=0
        )
        # (S, n, P) * (S, P) summed over P -> (S, n)
        contrib = np.einsum("snp,sp->sn", design, coef)
        return np.asarray(intercept[:, None] + contrib, dtype=np.float64)

    def sigma(self, fit: FitResult, *, use_draws: bool) -> Array:
        src = fit.params if use_draws else fit.param_means
        return np.asarray(src[self.sigma_name], dtype=np.float64)

    def abduct(self, fit: FitResult, value: Array, parent_values: Array) -> Array:
        """The exact residual ``U = X - (b0 + w . PA)`` at the posterior mean."""
        mean = self.structural_mean(fit, parent_values)
        return gaussian_residual(np.asarray(value, dtype=np.float64), mean)

abduct(fit, value, parent_values)

The exact residual U = X - (b0 + w . PA) at the posterior mean.

Source code in src/ergodic/scm/_mechanism.py

def abduct(self, fit: FitResult, value: Array, parent_values: Array) -> Array:
    """The exact residual ``U = X - (b0 + w . PA)`` at the posterior mean."""
    mean = self.structural_mean(fit, parent_values)
    return gaussian_residual(np.asarray(value, dtype=np.float64), mean)

build(x, y)

The per-node PyMC model: y ~ Normal(b0 + x_encoded . w, sigma).

x is the raw parent matrix (codes for discrete parents); it is encoded to the one-hot design here, the same encoding the generative model uses.

Source code in src/ergodic/scm/_mechanism.py

def build(self, x: Array, y: Array) -> Any:
    """The per-node PyMC model: ``y ~ Normal(b0 + x_encoded . w, sigma)``.

    ``x`` is the raw parent matrix (codes for discrete parents); it is encoded
    to the one-hot design here, the same encoding the generative model uses.
    """
    import pymc as pm

    x_enc = encode_parent_design(x, self.parent_cards) if self.parents else x
    _, p = x_enc.shape
    with pm.Model() as model:
        x_data = pm.Data("__parents", x_enc) if p else None
        intercept = pm.Normal(self.intercept_name, 0.0, self.coef_prior_sd)
        mean = intercept
        if p:
            coef = pm.Normal(self.coef_name, 0.0, self.coef_prior_sd, shape=p)
            mean = mean + pm.math.dot(x_data, coef)
        sigma = pm.HalfNormal(self.sigma_name, self.sigma_prior_sd)
        pm.Normal(self.name, mu=mean, sigma=sigma, observed=y)
    return model

generative_rv(nodes, n)

Add this node's RV to the active generative model: Normal(b0 + PA . w, sigma).

nodes maps each (already-built) parent name to its RV. Parents are encoded by the shared pytensor encoder, so the design matches the fit view column for column; the RV names match the merged posterior exactly, so sample_posterior_predictive reads the fitted draws. An incident latent L (a per-unit Normal(0, 1) the compose helper builds first, found in nodes) adds loading_L * L to the mean, so pm.do on this node severs it from L while a district sibling still reads L -- the confounding cancels in an interventional contrast.

Source code in src/ergodic/scm/_mechanism.py

def generative_rv(self, nodes: Mapping[str, Any], n: int) -> Any:
    """Add this node's RV to the active generative model: ``Normal(b0 + PA . w, sigma)``.

    ``nodes`` maps each (already-built) parent name to its RV. Parents are
    encoded by the shared pytensor encoder, so the design matches the fit view
    column for column; the RV names match the merged posterior exactly, so
    ``sample_posterior_predictive`` reads the fitted draws. An incident latent
    ``L`` (a per-unit ``Normal(0, 1)`` the compose helper builds first, found in
    ``nodes``) adds ``loading_L * L`` to the mean, so ``pm.do`` on this node
    severs it from ``L`` while a district sibling still reads ``L`` -- the
    confounding cancels in an interventional contrast.
    """
    import pymc as pm
    import pytensor.tensor as pt

    intercept = pm.Normal(self.intercept_name, 0.0, self.coef_prior_sd)
    mean: Any = intercept * pt.ones((n,))
    if self.parents:
        design = encode_parents_pt([nodes[p] for p in self.parents], self.parent_cards)
        coef = pm.Normal(self.coef_name, 0.0, self.coef_prior_sd, shape=self._design_width)
        mean = mean + pm.math.dot(design, coef)
    for lat in self.latents:
        loading = self._loading_rv(lat)
        mean = mean + loading * nodes[lat]
    sigma = pm.HalfNormal(self.sigma_name, self.sigma_prior_sd)
    return pm.Normal(self.name, mu=mean, sigma=sigma, shape=(n,))

loading_name(latent)

The parameter name for this node's loading on latent L (district fit).

Source code in src/ergodic/scm/_mechanism.py

def loading_name(self, latent: str) -> str:
    """The parameter name for this node's loading on latent ``L`` (district fit)."""
    return f"{self.name}__loading__{latent}"

mean_per_draw(fit, parent_values)

The mean per parameter-posterior draw.

parent_values is the (S, n, p) raw per-draw parent matrix (a non-intervened parent carries its own counterfactual posterior; codes for a discrete parent). It is one-hot encoded per draw, so the encoded design matches the (S, P) coefficient draws. Returns (S, n). With no parents the result is the broadcast intercept draws.

Source code in src/ergodic/scm/_mechanism.py

def mean_per_draw(self, fit: FitResult, parent_values: Array) -> Array:
    """The mean per parameter-posterior draw.

    ``parent_values`` is the ``(S, n, p)`` *raw* per-draw parent matrix (a
    non-intervened parent carries its own counterfactual posterior; codes for a
    discrete parent). It is one-hot encoded per draw, so the encoded design
    matches the ``(S, P)`` coefficient draws. Returns ``(S, n)``. With no parents
    the result is the broadcast intercept draws.
    """
    intercept = np.asarray(fit.params[self.intercept_name], dtype=np.float64)  # (S,)
    pv = np.asarray(parent_values, dtype=np.float64)
    s = intercept.shape[0]
    n = pv.shape[1] if pv.ndim == 3 else pv.shape[0]
    if not self.parents:
        return np.repeat(intercept[:, None], n, axis=1)
    coef = np.asarray(fit.params[self.coef_name], dtype=np.float64)  # (S, P)
    # broadcast an intervened parent's singleton draw axis up to S, then encode
    # each draw's raw (n, p) slice to its (n, P) one-hot design.
    if pv.shape[0] != s:
        pv = np.broadcast_to(pv, (s, n, pv.shape[2]))
    design = np.stack(
        [encode_parent_design(pv[d], self.parent_cards) for d in range(s)], axis=0
    )
    # (S, n, P) * (S, P) summed over P -> (S, n)
    contrib = np.einsum("snp,sp->sn", design, coef)
    return np.asarray(intercept[:, None] + contrib, dtype=np.float64)

structural_mean(fit, parent_values)

The mean b0 + w . encode(PA) at the posterior mean.

parent_values is the (n, p) raw parent matrix (codes for discrete parents); it is one-hot encoded internally, so the coefficients line up with the fit. (n, p) in, (n,) out.

Source code in src/ergodic/scm/_mechanism.py

def structural_mean(self, fit: FitResult, parent_values: Array) -> Array:
    """The mean ``b0 + w . encode(PA)`` at the posterior mean.

    ``parent_values`` is the ``(n, p)`` *raw* parent matrix (codes for discrete
    parents); it is one-hot encoded internally, so the coefficients line up with
    the fit. ``(n, p)`` in, ``(n,)`` out.
    """
    intercept = float(np.asarray(fit.param_means[self.intercept_name]))
    pv = np.asarray(parent_values, dtype=np.float64)
    if not self.parents:
        return np.full(pv.shape[0], intercept)
    design = encode_parent_design(pv, self.parent_cards)
    coef = np.asarray(fit.param_means[self.coef_name], dtype=np.float64)
    return intercept + design @ coef

ergodic.scm.BinaryGumbelMax dataclass

Bases: CategoricalGumbelMax

A binary node as a two-category Gumbel-max (the logistic mechanism).

Specializes :class:CategoricalGumbelMax to two categories; the marginal is the Bernoulli logistic, and the Gumbel-max counterfactual is the binary common-random-numbers flip.

Source code in src/ergodic/scm/_mechanism.py

@dataclass(frozen=True)
class BinaryGumbelMax(CategoricalGumbelMax):
    """A binary node as a two-category Gumbel-max (the logistic mechanism).

    Specializes :class:`CategoricalGumbelMax` to two categories; the marginal is
    the Bernoulli logistic, and the Gumbel-max counterfactual is the binary
    common-random-numbers flip.
    """

    n_categories: int = 2

ergodic.scm.CategoricalGumbelMax dataclass

Categorical mechanism as Gumbel-max over softmax logits.

The logits are linear in the parents with a per-class slope and intercept, fitted by the equivalent :class:Categorical likelihood (the Gumbel-max marginal). The forward mechanism is category = argmax(logits + g) with g an iid Gumbel vector, and the counterfactual reuses the abducted g under shifted logits, so it is invertible. Identified up to the per-row logit gauge, which the argmax washes out.

Source code in src/ergodic/scm/_mechanism.py

@dataclass(frozen=True)
class CategoricalGumbelMax:
    """Categorical mechanism as Gumbel-max over softmax logits.

    The logits are linear in the parents with a per-class slope and intercept,
    fitted by the equivalent :class:`Categorical` likelihood (the Gumbel-max
    marginal). The forward mechanism is ``category = argmax(logits + g)`` with
    ``g`` an iid Gumbel vector, and the counterfactual reuses the abducted ``g``
    under shifted logits, so it is invertible. Identified up to the per-row logit
    gauge, which the argmax washes out.
    """

    name: str
    parents: tuple[str, ...]
    parent_cards: tuple[int | None, ...] = ()
    n_categories: int = 0
    coef_prior_sd: float = 3.0
    latents: tuple[str, ...] = ()
    gauged_latents: tuple[str, ...] = ()
    loading_prior_sd: float = 2.0

    @property
    def capability(self) -> Capability:
        return Capability.INVERTIBLE

    @property
    def weight_name(self) -> str:
        return f"{self.name}__logit_w"

    @property
    def bias_name(self) -> str:
        return f"{self.name}__logit_b"

    def loading_name(self, latent: str) -> str:
        """The per-class logit loading parameter for this node on latent ``L``."""
        return f"{self.name}__loading__{latent}"

    @property
    def _design_width(self) -> int:
        """The encoded weight rows: ``len(parents)`` for all-continuous parents (no
        ``parent_cards``), else the one-hot-expanded width."""
        if not self.parents:
            return 0
        if not self.parent_cards:
            return len(self.parents)
        return encoded_width(self.parent_cards)

    def build(self, x: Array, y_codes: IntArray, n_cat: int) -> Any:
        import pymc as pm

        x_enc = encode_parent_design(x, self.parent_cards) if self.parents else x
        n, p = x_enc.shape
        with pm.Model() as model:
            bias = pm.Normal(self.bias_name, 0.0, self.coef_prior_sd, shape=n_cat)
            logits = bias[None, :] * np.ones((n, 1))
            if p:
                weight = pm.Normal(self.weight_name, 0.0, self.coef_prior_sd, shape=(p, n_cat))
                logits = logits + pm.math.dot(x_enc, weight)
            pm.Categorical(self.name, logit_p=logits, observed=y_codes)
        return model

    def generative_rv(self, nodes: Mapping[str, Any], n: int) -> Any:
        """Add this node's categorical RV to the active generative model.

        ``logits = bias + encode(PA) . weight`` (same per-class linear form and same
        one-hot parent encoding the fit uses), then ``pm.Categorical`` samples the
        node forward (the Gumbel-max marginal). The RV names match the merged
        posterior, so ``sample_posterior_predictive`` reads the fitted logits. The
        node's resolved class count must be passed in (it lives on the fit, not the
        spec). An incident latent ``L`` adds a per-class logit term
        ``L * loading_L`` (``loading_L`` shape ``(K,)``), so the confounding shifts
        the relative class logits, and ``pm.do`` severs the node from ``L``.
        """
        import pymc as pm
        import pytensor.tensor as pt

        if self.n_categories <= 0:
            raise ValueError(
                f"categorical node {self.name!r} has an unresolved class count; the "
                "generative model needs n_categories (read it off the fit)"
            )
        bias = pm.Normal(self.bias_name, 0.0, self.coef_prior_sd, shape=self.n_categories)
        logits: Any = bias[None, :] * pt.ones((n, 1))
        if self.parents:
            design = encode_parents_pt([nodes[p] for p in self.parents], self.parent_cards)
            weight = pm.Normal(
                self.weight_name,
                0.0,
                self.coef_prior_sd,
                shape=(self._design_width, self.n_categories),
            )
            logits = logits + pm.math.dot(design, weight)
        for lat in self.latents:
            loading = self._loading_rv(lat)  # (K,)
            logits = logits + nodes[lat][:, None] * loading[None, :]
        return pm.Categorical(self.name, logit_p=logits, shape=(n,))

    def _loading_rv(self, latent: str) -> Any:
        """The per-class logit loading RV for ``latent`` (gauge: first class pinned).

        Shape ``(K,)``. When ``latent`` is gauged, the first class's loading is a
        ``HalfNormal`` (pinned positive to break the ``L -> -L`` reflection) and the
        remaining classes are ``Normal``; otherwise every class is ``Normal``.
        """
        # pm.math re-exports the pytensor ops as the typed-Any surface (pytensor.tensor's
        # own ops trip mypy's no-untyped-call); this only runs inside a generative/fit
        # model build, so pymc is already in play.
        import pymc as pm

        name = self.loading_name(latent)
        if latent in self.gauged_latents and self.n_categories > 1:
            head = pm.HalfNormal(f"{name}__0", self.loading_prior_sd)
            tail = pm.Normal(
                f"{name}__rest", 0.0, self.loading_prior_sd, shape=self.n_categories - 1
            )
            # the (K,) per-class loading: the pinned-positive first class, then the rest.
            return pm.math.concatenate([head[None], tail])
        return pm.Normal(name, 0.0, self.loading_prior_sd, shape=self.n_categories)

    def fit(self, data: Dataset, *, draws: int, tune: int, chains: int, seed: int) -> FitResult:
        import pymc as pm

        x = _parent_matrix(data, self.parents)
        y_codes = np.asarray(data.codes(self.name), dtype=np.int64)
        cats = data.categories(self.name)
        n_cat = self.n_categories or (len(cats) if cats is not None else int(y_codes.max()) + 1)
        model = self.build(x, y_codes, n_cat)
        with model:
            idata = pm.sample(
                draws=draws,
                tune=tune,
                chains=chains,
                target_accept=0.9,
                random_seed=seed,
                progressbar=False,
                compute_convergence_checks=False,
            )
        post = idata.posterior
        params: dict[str, Array] = {self.bias_name: _flatten(post[self.bias_name])}
        if self.parents:
            params[self.weight_name] = _flatten(post[self.weight_name])
        means = {k: np.asarray(v.mean(axis=0), dtype=np.float64) for k, v in params.items()}
        return FitResult(params=params, param_means=means, n_categories=n_cat)

    def logits(self, fit: FitResult, parent_values: Array) -> Array:
        """The ``(n, K)`` class logits at the posterior mean.

        ``parent_values`` is the ``(n, p)`` *raw* parent matrix (codes for discrete
        parents); it is one-hot encoded internally so the weight rows line up.
        """
        bias = np.asarray(fit.param_means[self.bias_name], dtype=np.float64)
        pv = np.asarray(parent_values, dtype=np.float64)
        out = np.repeat(bias[None, :], pv.shape[0], axis=0)
        if self.parents:
            design = encode_parent_design(pv, self.parent_cards)
            weight = np.asarray(fit.param_means[self.weight_name], dtype=np.float64)
            out = out + design @ weight
        return np.asarray(out, dtype=np.float64)

    def logits_per_draw(self, fit: FitResult, parent_values: Array) -> Array:
        """The ``(S, n, K)`` class logits per parameter-posterior draw.

        The per-draw twin of :meth:`logits`: ``bias`` is ``(S, K)`` and ``weight``
        ``(S, P, K)``, so each draw's logits read ``bias[d] + encode(PA)[d] . w[d]``.
        ``parent_values`` is the ``(S, n, p)`` *raw* per-draw parent matrix (an
        intervened/categorical parent's singleton draw axis broadcasts up to ``S``,
        and is one-hot encoded per draw so the design matches the ``(S, P, K)`` weight
        draws). With no parents the result is the broadcast bias draws. Used by the
        counterfactual to spread a categorical ancestor's decision across draws, so
        its logit-parameter posterior enters the counterfactual interval (the
        posterior-mean collapse it replaces understated the width).
        """
        bias = np.asarray(fit.params[self.bias_name], dtype=np.float64)  # (S, K)
        s, k = bias.shape
        pv = np.asarray(parent_values, dtype=np.float64)
        n = pv.shape[1] if pv.ndim == 3 else pv.shape[0]
        if not self.parents:
            return np.broadcast_to(bias[:, None, :], (s, n, k)).copy()
        weight = np.asarray(fit.params[self.weight_name], dtype=np.float64)  # (S, P, K)
        if pv.ndim == 2:
            pv = pv[None, :, :]
        if pv.shape[0] != s:
            pv = np.broadcast_to(pv, (s, n, pv.shape[2]))
        design = np.stack(
            [encode_parent_design(pv[d], self.parent_cards) for d in range(s)], axis=0
        )  # (S, n, P)
        # (S, n, P) @ (S, P, K) batched over the draw axis -> (S, n, K)
        contrib = np.einsum("snp,spk->snk", design, weight)
        return np.asarray(bias[:, None, :] + contrib, dtype=np.float64)

    def abduct(
        self, fit: FitResult, observed_codes: IntArray, parent_values: Array, *, seed: int
    ) -> Array:
        """Abduct the per-unit Gumbel noise consistent with the observed category."""
        rng = np.random.default_rng(seed)
        return gumbel_max_abduct(self.logits(fit, parent_values), observed_codes, rng)

    def decide(self, logits: Array, gumbel: Array) -> IntArray:
        return gumbel_max_decide(logits, gumbel)

abduct(fit, observed_codes, parent_values, *, seed)

Abduct the per-unit Gumbel noise consistent with the observed category.

Source code in src/ergodic/scm/_mechanism.py

def abduct(
    self, fit: FitResult, observed_codes: IntArray, parent_values: Array, *, seed: int
) -> Array:
    """Abduct the per-unit Gumbel noise consistent with the observed category."""
    rng = np.random.default_rng(seed)
    return gumbel_max_abduct(self.logits(fit, parent_values), observed_codes, rng)

generative_rv(nodes, n)

Add this node's categorical RV to the active generative model.

logits = bias + encode(PA) . weight (same per-class linear form and same one-hot parent encoding the fit uses), then pm.Categorical samples the node forward (the Gumbel-max marginal). The RV names match the merged posterior, so sample_posterior_predictive reads the fitted logits. The node's resolved class count must be passed in (it lives on the fit, not the spec). An incident latent L adds a per-class logit term L * loading_L (loading_L shape (K,)), so the confounding shifts the relative class logits, and pm.do severs the node from L.

Source code in src/ergodic/scm/_mechanism.py

def generative_rv(self, nodes: Mapping[str, Any], n: int) -> Any:
    """Add this node's categorical RV to the active generative model.

    ``logits = bias + encode(PA) . weight`` (same per-class linear form and same
    one-hot parent encoding the fit uses), then ``pm.Categorical`` samples the
    node forward (the Gumbel-max marginal). The RV names match the merged
    posterior, so ``sample_posterior_predictive`` reads the fitted logits. The
    node's resolved class count must be passed in (it lives on the fit, not the
    spec). An incident latent ``L`` adds a per-class logit term
    ``L * loading_L`` (``loading_L`` shape ``(K,)``), so the confounding shifts
    the relative class logits, and ``pm.do`` severs the node from ``L``.
    """
    import pymc as pm
    import pytensor.tensor as pt

    if self.n_categories <= 0:
        raise ValueError(
            f"categorical node {self.name!r} has an unresolved class count; the "
            "generative model needs n_categories (read it off the fit)"
        )
    bias = pm.Normal(self.bias_name, 0.0, self.coef_prior_sd, shape=self.n_categories)
    logits: Any = bias[None, :] * pt.ones((n, 1))
    if self.parents:
        design = encode_parents_pt([nodes[p] for p in self.parents], self.parent_cards)
        weight = pm.Normal(
            self.weight_name,
            0.0,
            self.coef_prior_sd,
            shape=(self._design_width, self.n_categories),
        )
        logits = logits + pm.math.dot(design, weight)
    for lat in self.latents:
        loading = self._loading_rv(lat)  # (K,)
        logits = logits + nodes[lat][:, None] * loading[None, :]
    return pm.Categorical(self.name, logit_p=logits, shape=(n,))

loading_name(latent)

The per-class logit loading parameter for this node on latent L.

Source code in src/ergodic/scm/_mechanism.py

def loading_name(self, latent: str) -> str:
    """The per-class logit loading parameter for this node on latent ``L``."""
    return f"{self.name}__loading__{latent}"

logits(fit, parent_values)

The (n, K) class logits at the posterior mean.

parent_values is the (n, p) raw parent matrix (codes for discrete parents); it is one-hot encoded internally so the weight rows line up.

Source code in src/ergodic/scm/_mechanism.py

def logits(self, fit: FitResult, parent_values: Array) -> Array:
    """The ``(n, K)`` class logits at the posterior mean.

    ``parent_values`` is the ``(n, p)`` *raw* parent matrix (codes for discrete
    parents); it is one-hot encoded internally so the weight rows line up.
    """
    bias = np.asarray(fit.param_means[self.bias_name], dtype=np.float64)
    pv = np.asarray(parent_values, dtype=np.float64)
    out = np.repeat(bias[None, :], pv.shape[0], axis=0)
    if self.parents:
        design = encode_parent_design(pv, self.parent_cards)
        weight = np.asarray(fit.param_means[self.weight_name], dtype=np.float64)
        out = out + design @ weight
    return np.asarray(out, dtype=np.float64)

logits_per_draw(fit, parent_values)

The (S, n, K) class logits per parameter-posterior draw.

The per-draw twin of :meth:logits: bias is (S, K) and weight (S, P, K), so each draw's logits read bias[d] + encode(PA)[d] . w[d]. parent_values is the (S, n, p) raw per-draw parent matrix (an intervened/categorical parent's singleton draw axis broadcasts up to S, and is one-hot encoded per draw so the design matches the (S, P, K) weight draws). With no parents the result is the broadcast bias draws. Used by the counterfactual to spread a categorical ancestor's decision across draws, so its logit-parameter posterior enters the counterfactual interval (the posterior-mean collapse it replaces understated the width).

Source code in src/ergodic/scm/_mechanism.py

def logits_per_draw(self, fit: FitResult, parent_values: Array) -> Array:
    """The ``(S, n, K)`` class logits per parameter-posterior draw.

    The per-draw twin of :meth:`logits`: ``bias`` is ``(S, K)`` and ``weight``
    ``(S, P, K)``, so each draw's logits read ``bias[d] + encode(PA)[d] . w[d]``.
    ``parent_values`` is the ``(S, n, p)`` *raw* per-draw parent matrix (an
    intervened/categorical parent's singleton draw axis broadcasts up to ``S``,
    and is one-hot encoded per draw so the design matches the ``(S, P, K)`` weight
    draws). With no parents the result is the broadcast bias draws. Used by the
    counterfactual to spread a categorical ancestor's decision across draws, so
    its logit-parameter posterior enters the counterfactual interval (the
    posterior-mean collapse it replaces understated the width).
    """
    bias = np.asarray(fit.params[self.bias_name], dtype=np.float64)  # (S, K)
    s, k = bias.shape
    pv = np.asarray(parent_values, dtype=np.float64)
    n = pv.shape[1] if pv.ndim == 3 else pv.shape[0]
    if not self.parents:
        return np.broadcast_to(bias[:, None, :], (s, n, k)).copy()
    weight = np.asarray(fit.params[self.weight_name], dtype=np.float64)  # (S, P, K)
    if pv.ndim == 2:
        pv = pv[None, :, :]
    if pv.shape[0] != s:
        pv = np.broadcast_to(pv, (s, n, pv.shape[2]))
    design = np.stack(
        [encode_parent_design(pv[d], self.parent_cards) for d in range(s)], axis=0
    )  # (S, n, P)
    # (S, n, P) @ (S, P, K) batched over the draw axis -> (S, n, K)
    contrib = np.einsum("snp,spk->snk", design, weight)
    return np.asarray(bias[:, None, :] + contrib, dtype=np.float64)

ergodic.scm.Mechanism

Bases: Protocol

The per-node structural equation.

Implementations carry the node name and its parent names and expose the four jobs below. The SCM holds one mechanism per node and a :class:FitResult per mechanism after fitting. name and parents are read-only: mechanisms are frozen value objects, so filling a blank node name returns a new instance (:func:resolve_mechanism) rather than mutating in place.

Source code in src/ergodic/scm/_mechanism.py

@runtime_checkable
class Mechanism(Protocol):
    """The per-node structural equation.

    Implementations carry the node name and its parent names and expose the four
    jobs below. The SCM holds one mechanism per node and a :class:`FitResult` per
    mechanism after fitting. ``name`` and ``parents`` are read-only: mechanisms are
    frozen value objects, so filling a blank node name returns a new instance
    (:func:`resolve_mechanism`) rather than mutating in place.
    """

    @property
    def name(self) -> str:
        """The node this mechanism is the structural equation for."""
        ...

    @property
    def parents(self) -> tuple[str, ...]:
        """The node's observed parents, in the mechanism's design order."""
        ...

    @property
    def latents(self) -> tuple[str, ...]:
        """Incident latent RV names from confounded-district loadings; () for a Markovian node."""
        ...

    @property
    def gauged_latents(self) -> tuple[str, ...]:
        """The subset of ``latents`` whose loading is pinned positive for the sign gauge."""
        ...

    @property
    def capability(self) -> Capability:
        """``INVERTIBLE`` if counterfactuals are exact, else ``PROBABILISTIC``."""
        ...

    def fit(self, data: Dataset, *, draws: int, tune: int, chains: int, seed: int) -> FitResult:
        """Fit the node on its observed parents; return the posterior."""
        ...

    def generative_rv(self, nodes: Mapping[str, Any], n: int) -> Any:
        """Add this node's RV to the active generative PyMC model, return the RV.

        ``nodes`` maps each already-built parent name to its RV; the structural
        equation reads the parents from it. Names every parameter RV with the same
        scheme :meth:`fit` used, so the merged posterior matches by name. Called
        inside a ``pm.Model()`` context by the compose helper.
        """
        ...

capability property

INVERTIBLE if counterfactuals are exact, else PROBABILISTIC.

gauged_latents property

The subset of latents whose loading is pinned positive for the sign gauge.

latents property

Incident latent RV names from confounded-district loadings; () for a Markovian node.

name property

The node this mechanism is the structural equation for.

parents property

The node's observed parents, in the mechanism's design order.

fit(data, *, draws, tune, chains, seed)

Fit the node on its observed parents; return the posterior.

Source code in src/ergodic/scm/_mechanism.py

def fit(self, data: Dataset, *, draws: int, tune: int, chains: int, seed: int) -> FitResult:
    """Fit the node on its observed parents; return the posterior."""
    ...

generative_rv(nodes, n)

Add this node's RV to the active generative PyMC model, return the RV.

nodes maps each already-built parent name to its RV; the structural equation reads the parents from it. Names every parameter RV with the same scheme :meth:fit used, so the merged posterior matches by name. Called inside a pm.Model() context by the compose helper.

Source code in src/ergodic/scm/_mechanism.py

def generative_rv(self, nodes: Mapping[str, Any], n: int) -> Any:
    """Add this node's RV to the active generative PyMC model, return the RV.

    ``nodes`` maps each already-built parent name to its RV; the structural
    equation reads the parents from it. Names every parameter RV with the same
    scheme :meth:`fit` used, so the merged posterior matches by name. Called
    inside a ``pm.Model()`` context by the compose helper.
    """
    ...