"""Erdos-Renyi graph distributions for binary graph observations.
Data type: a binary graph observation represented as a square adjacency matrix,
a NetworkX-like graph, or a mapping accepted by ``GraphDataEncoder``.
"""
from __future__ import annotations
import math
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
import numpy as np
from numpy.random import RandomState
from mixle.enumeration.algorithms import BufferedStream, ProductEnumerator
from mixle.stats.compute.pdist import (
DistributionEnumerator,
DistributionSampler,
EnumerationError,
ParameterEstimator,
SequenceEncodableProbabilityDistribution,
SequenceEncodableStatisticAccumulator,
StatisticAccumulatorFactory,
)
from mixle.utils.exact import require_exact_bool
if TYPE_CHECKING:
from mixle.data.sources.graph_source import GraphDataEncoder, GraphObservation
# mixle.data.sources.graph_source's own top-level import of mixle.stats.compute.pdist.DataSequenceEncoder
# forces mixle.stats's package __init__ to run first, which eagerly imports this very module -- so a
# module-level `from mixle.data.sources.graph_source import ...` here is circular whenever
# mixle.data.sources.graph_source is the entry point (e.g. `import mixle.data.sources.graph_source`
# directly, before mixle.stats has been warmed by any other path): graph_source would still be
# mid-import (paused inside its own DataSequenceEncoder import, above this module in the same chain),
# so none of GraphDataEncoder/GraphObservation/the coercion helpers would exist as attributes yet.
# Every usage below is either purely a type annotation (deferred to a string by the `from __future__
# import annotations` above) or a call inside a function/method body, so the imports are deferred to
# call time, once every module in the cycle has finished loading normally -- mirroring the
# pilot_ladder.py / mixture.py fixes for the same shape of cycle.
def _optional_num_nodes(value: Any, *, name: str = "num_nodes") -> int | None:
if value is None:
return None
if isinstance(value, (bool, np.bool_)) or not isinstance(value, (int, np.integer)):
raise ValueError(f"{name} must be an exact non-Boolean integer or None.")
result = int(value)
if result < 0:
raise ValueError(f"{name} must be non-negative.")
return result
def _require_graph_size(adjacency: np.ndarray, num_nodes: int | None) -> None:
if num_nodes is not None and adjacency.shape != (num_nodes, num_nodes):
raise ValueError(
f"graph observation has shape {adjacency.shape}; expected ({num_nodes}, {num_nodes}) "
"for this fixed-size distribution."
)
def _finite_nonnegative(value: Any, *, name: str) -> float:
result = float(value)
if not np.isfinite(result) or result < 0.0:
raise ValueError(f"{name} must be finite and non-negative.")
return result
def _validate_statistics(suff_stat: Any) -> tuple[float, float]:
if not isinstance(suff_stat, (tuple, list)) or len(suff_stat) != 2:
raise ValueError("Erdos-Renyi statistics must be (edge_opportunities, edge_count).")
total = _finite_nonnegative(suff_stat[0], name="edge_opportunities")
successes = _finite_nonnegative(suff_stat[1], name="edge_count")
if successes > total:
raise ValueError("edge_count cannot exceed edge_opportunities.")
return total, successes
def _require_matching_configuration(left: ErdosRenyiGraphAccumulator, right: ErdosRenyiGraphAccumulator) -> None:
if left.directed != right.directed or left.self_loops != right.self_loops or left.num_nodes != right.num_nodes:
raise ValueError("cannot share Erdos-Renyi statistics across incompatible graph configurations.")
[docs]
class ErdosRenyiGraphDistribution(SequenceEncodableProbabilityDistribution):
"""Independent Bernoulli distribution over binary graph edges."""
[docs]
@classmethod
def compute_capabilities(cls):
"""Return backend capabilities for Bernoulli graph scoring."""
from mixle.stats.compute.capabilities import DistributionCapabilities
return DistributionCapabilities(engine_ready=("numpy", "torch"), kernel_status="generic_object")
[docs]
@classmethod
def compute_declaration(cls):
"""Return the structured declaration for the Erdos-Renyi graph family."""
from mixle.stats.compute.declarations import DistributionDeclaration, ParameterSpec, StatisticSpec
return DistributionDeclaration(
name="erdos_renyi_graph",
distribution_type=cls,
parameters=(
ParameterSpec("p", constraint="unit_interval"),
ParameterSpec("directed", constraint="fixed", differentiable=False),
ParameterSpec("self_loops", constraint="fixed", differentiable=False),
ParameterSpec("num_nodes", constraint="optional_integer", differentiable=False),
),
statistics=(
StatisticSpec("edge_opportunities"),
StatisticSpec("edge_count"),
),
support="binary_graph",
)
def __init__(
self,
p: float,
directed: bool = False,
self_loops: bool = False,
num_nodes: int | None = None,
name: str | None = None,
keys: str | None = None,
) -> None:
from mixle.data.sources.graph_source import _clip_prob
self.p = _clip_prob(p)
self.log_p = -math.inf if self.p == 0.0 else math.log(self.p)
self.log_1p = -math.inf if self.p == 1.0 else math.log1p(-self.p)
self.directed = require_exact_bool(directed, "directed")
self.self_loops = require_exact_bool(self_loops, "self_loops")
self.num_nodes = _optional_num_nodes(num_nodes)
self.name = name
self.keys = keys
def __str__(self) -> str:
return "ErdosRenyiGraphDistribution(p=%s, directed=%s, self_loops=%s, num_nodes=%s, name=%s, keys=%s)" % (
repr(self.p),
repr(self.directed),
repr(self.self_loops),
repr(self.num_nodes),
repr(self.name),
repr(self.keys),
)
[docs]
@classmethod
def from_model(cls, model: Any) -> ErdosRenyiGraphDistribution:
"""Create a distribution wrapper from an Erdos-Renyi model."""
return cls(
model.p,
directed=model.directed,
self_loops=model.self_loops,
num_nodes=getattr(model, "num_nodes", None),
name=getattr(model, "name", None),
)
[docs]
def to_model(self) -> Any:
"""Convert this distribution to the corresponding random-graph model."""
from mixle.models.random_graph import ErdosRenyiGraphModel
return ErdosRenyiGraphModel(self.p, directed=self.directed, self_loops=self.self_loops, name=self.name)
[docs]
def density(self, x: Any) -> float:
"""Return the probability mass of one graph observation."""
return math.exp(self.log_density(x))
[docs]
def log_density(self, x: Any) -> float:
"""Return the Bernoulli edge log probability of one graph."""
from mixle.data.sources.graph_source import (
_bernoulli_log_likelihood,
_edge_counts,
_extract_observation,
_validate_graph_constraints,
)
obs = _extract_observation(x, directed=self.directed)
_validate_graph_constraints(obs.adjacency, self.directed, self.self_loops)
_require_graph_size(obs.adjacency, self.num_nodes)
total, successes = _edge_counts(obs.adjacency, self.directed, self.self_loops)
return _bernoulli_log_likelihood(successes, total, self.p)
[docs]
def seq_log_density(self, x: Sequence[GraphObservation]) -> np.ndarray:
"""Score a batch of graph observations."""
# Extract (opportunities, successes) per graph once (ragged adjacency forces the per-graph
# extraction), then score the whole batch with one vectorized Bernoulli log-likelihood instead
# of a Python log_density call per graph.
from mixle.data.sources.graph_source import _edge_counts, _extract_observation, _validate_graph_constraints
if len(x) == 0:
return np.zeros(0, dtype=np.float64)
counts = []
for o in x:
adj = _extract_observation(o, directed=self.directed).adjacency
_validate_graph_constraints(adj, self.directed, self.self_loops)
_require_graph_size(adj, self.num_nodes)
counts.append(_edge_counts(adj, self.directed, self.self_loops))
from mixle.data.sources.graph_source import _bernoulli_log_likelihood
return np.asarray(
[_bernoulli_log_likelihood(successes, total, self.p) for total, successes in counts],
dtype=np.float64,
)
[docs]
def backend_seq_log_density(self, x: Sequence[GraphObservation], engine: Any) -> Any:
"""Engine-routed Bernoulli edge log-likelihood.
Per-graph edge opportunities/successes are extracted host-side (the graphs are ragged
object data), but the Bernoulli reduction runs on the active engine, so the model's scoring
math is engine-native (and differentiable in ``p`` on torch).
"""
from mixle.data.sources.graph_source import (
_bernoulli_log_likelihood,
_edge_counts,
_extract_observation,
_validate_graph_constraints,
)
counts = []
for o in x:
adj = _extract_observation(o, directed=self.directed).adjacency
_validate_graph_constraints(adj, self.directed, self.self_loops)
_require_graph_size(adj, self.num_nodes)
counts.append(_edge_counts(adj, self.directed, self.self_loops))
return engine.asarray(
np.asarray(
[_bernoulli_log_likelihood(successes, total, self.p) for total, successes in counts],
dtype=np.float64,
)
)
[docs]
def edge_probability(self, i: int | None = None, j: int | None = None, context: Any | None = None) -> float:
"""Return the common edge probability ``p``."""
return self.p
[docs]
def edge_marginals(self, num_nodes: int | None = None) -> np.ndarray:
"""Return the matrix of marginal edge probabilities for ``num_nodes``."""
n = self.num_nodes if num_nodes is None else _optional_num_nodes(num_nodes)
if n is None:
raise ValueError("num_nodes is required when distribution.num_nodes is None.")
if self.num_nodes is not None and n != self.num_nodes:
raise ValueError(f"num_nodes={n} conflicts with this fixed-size distribution ({self.num_nodes}).")
mat = np.full((n, n), self.p, dtype=np.float64)
if not self.self_loops:
np.fill_diagonal(mat, 0.0)
return mat
[docs]
def posterior(self, x: Any) -> dict[str, float]:
"""Return edge opportunities, observed edge count, and fitted probability for ``x``."""
from mixle.data.sources.graph_source import _edge_counts, _extract_observation, _validate_graph_constraints
adj = _extract_observation(x, directed=self.directed).adjacency
_validate_graph_constraints(adj, self.directed, self.self_loops)
_require_graph_size(adj, self.num_nodes)
total, successes = _edge_counts(adj, self.directed, self.self_loops)
return {"edge_opportunities": total, "edge_count": successes, "p": self.p}
[docs]
def sampler(self, seed: int | None = None) -> ErdosRenyiGraphSampler:
"""Return a sampler for Erdos-Renyi graph observations."""
return ErdosRenyiGraphSampler(self, seed)
[docs]
def enumerator(self) -> DistributionEnumerator:
"""Enumerate binary graphs in descending probability order (requires ``num_nodes``).
The edges are independent Bernoulli(p) over the free positions ``_edge_indices(n, directed,
self_loops)``, so the graph distribution is a product of edge factors and enumerates by
best-first over the per-edge supports -- exactly like a composite of Bernoullis, with the
combined value assembled into an adjacency matrix (mirrored for the undirected case). Each
graph carries its exact ``log_density``. ``num_nodes`` must be set so the edge set is finite.
"""
if self.num_nodes is None:
raise EnumerationError(self, reason="num_nodes must be set to enumerate graphs")
return ErdosRenyiGraphEnumerator(self)
[docs]
def estimator(self, pseudo_count: float | None = None) -> ErdosRenyiGraphEstimator:
"""Return the edge-count estimator for this graph family."""
return ErdosRenyiGraphEstimator(
directed=self.directed,
self_loops=self.self_loops,
pseudo_count=pseudo_count,
prior_p=self.p,
num_nodes=self.num_nodes,
name=self.name,
keys=self.keys,
)
[docs]
def dist_to_encoder(self) -> GraphDataEncoder:
"""Return the graph encoder used by vectorized scoring and fitting."""
from mixle.data.sources.graph_source import GraphDataEncoder
return GraphDataEncoder(directed=self.directed)
[docs]
class ErdosRenyiGraphEnumerator(DistributionEnumerator):
"""Enumerator over finite Erdos-Renyi binary graph support."""
def __init__(self, dist: ErdosRenyiGraphDistribution) -> None:
"""Best-first enumeration of binary graphs over independent edge factors.
Args:
dist (ErdosRenyiGraphDistribution): Distribution whose graphs are enumerated (its
``num_nodes`` must be set).
"""
from mixle.data.sources.graph_source import _edge_indices
super().__init__(dist)
n = dist.num_nodes
edges = list(_edge_indices(n, dist.directed, dist.self_loops))
directed = dist.directed
# Each free edge is an independent Bernoulli(p): present (1) at log_p, absent (0) at log_1p,
# ordered by descending probability so the per-edge stream is sorted.
present, absent = (1, 0)
if dist.p == 0.0:
edge_pair = [(absent, 0.0)]
elif dist.p == 1.0:
edge_pair = [(present, 0.0)]
else:
edge_pair = (
[(present, dist.log_p), (absent, dist.log_1p)]
if dist.log_p >= dist.log_1p
else [(absent, dist.log_1p), (present, dist.log_p)]
)
def combine(edge_values: tuple[int, ...]) -> np.ndarray:
adj = np.zeros((n, n), dtype=np.int8)
for (i, j), v in zip(edges, edge_values):
adj[i, j] = v
if not directed:
adj[j, i] = v
return adj
streams = [BufferedStream(iter(list(edge_pair))) for _ in edges]
self._product = ProductEnumerator(streams, combine=combine)
def __next__(self) -> tuple[np.ndarray, float]:
"""Return the next adjacency matrix and its log probability."""
return next(self._product)
[docs]
class ErdosRenyiGraphSampler(DistributionSampler):
"""Sample binary graphs from an Erdos-Renyi distribution."""
def __init__(self, dist: ErdosRenyiGraphDistribution, seed: int | None = None) -> None:
self.dist = dist
self.rng = RandomState(seed)
[docs]
def sample_graph(self, num_nodes: int | None = None) -> np.ndarray:
"""Draw one binary graph adjacency matrix."""
n = self.dist.num_nodes if num_nodes is None else _optional_num_nodes(num_nodes)
if n is None:
raise ValueError("num_nodes is required when distribution.num_nodes is None.")
if self.dist.num_nodes is not None and n != self.dist.num_nodes:
raise ValueError(f"num_nodes={n} conflicts with this fixed-size distribution ({self.dist.num_nodes}).")
mat = (self.rng.rand(n, n) < self.dist.p).astype(np.int8)
if self.dist.directed:
if not self.dist.self_loops:
np.fill_diagonal(mat, 0)
return mat
upper = np.triu(mat, k=0 if self.dist.self_loops else 1)
mat = upper + upper.T
if self.dist.self_loops:
diag = (self.rng.rand(n) < self.dist.p).astype(np.int8)
np.fill_diagonal(mat, diag)
return mat
[docs]
def sample(
self, size: int | None = None, num_nodes: int | None = None, *, batched: bool = True
) -> np.ndarray | list[np.ndarray]:
"""Draw one graph or a list of graphs."""
if size is None:
return self.sample_graph(num_nodes=num_nodes)
sample_size = _optional_num_nodes(size, name="size")
return [self.sample_graph(num_nodes=num_nodes) for _ in range(sample_size)]
[docs]
class ErdosRenyiGraphAccumulator(SequenceEncodableStatisticAccumulator):
"""Accumulate edge counts for Erdos-Renyi graph fitting."""
def __init__(
self,
directed: bool = False,
self_loops: bool = False,
num_nodes: int | None = None,
name: str | None = None,
keys: str | None = None,
) -> None:
self.edge_opportunities = 0.0
self.edge_count = 0.0
self.directed = require_exact_bool(directed, "directed")
self.self_loops = require_exact_bool(self_loops, "self_loops")
self.num_nodes = _optional_num_nodes(num_nodes)
self.name = name
self.keys = keys
[docs]
def update(self, x: Any, weight: float, estimate: ErdosRenyiGraphDistribution | None) -> None:
"""Accumulate weighted edge opportunities and successes from one graph."""
from mixle.data.sources.graph_source import _edge_counts, _extract_observation, _validate_graph_constraints
obs = _extract_observation(x, directed=self.directed)
_validate_graph_constraints(obs.adjacency, self.directed, self.self_loops)
expected_nodes = self.num_nodes if estimate is None else estimate.num_nodes
_require_graph_size(obs.adjacency, expected_nodes)
total, successes = _edge_counts(obs.adjacency, self.directed, self.self_loops)
checked_weight = _finite_nonnegative(weight, name="weight")
self.edge_opportunities += checked_weight * total
self.edge_count += checked_weight * successes
[docs]
def initialize(self, x: Any, weight: float, rng: RandomState | None) -> None:
"""Initialize sufficient statistics from one weighted graph."""
self.update(x, weight, None)
[docs]
def seq_update(
self, x: Sequence[GraphObservation], weights: np.ndarray, estimate: ErdosRenyiGraphDistribution | None
) -> None:
"""Accumulate weighted edge counts from a batch of graphs."""
from mixle.data.sources.graph_source import _edge_counts, _validate_graph_constraints
checked_weights = np.asarray(weights, dtype=np.float64)
if checked_weights.ndim != 1 or len(checked_weights) != len(x):
raise ValueError("weights must be a one-dimensional array aligned with the graph batch.")
if np.any(~np.isfinite(checked_weights)) or np.any(checked_weights < 0.0):
raise ValueError("weights must be finite and non-negative.")
pending: list[tuple[float, float]] = []
for obs, weight in zip(x, checked_weights):
_validate_graph_constraints(obs.adjacency, self.directed, self.self_loops)
expected_nodes = self.num_nodes if estimate is None else estimate.num_nodes
_require_graph_size(obs.adjacency, expected_nodes)
total, successes = _edge_counts(obs.adjacency, self.directed, self.self_loops)
pending.append((float(weight) * total, float(weight) * successes))
for total, successes in pending:
self.edge_opportunities += total
self.edge_count += successes
[docs]
def seq_initialize(self, x: Sequence[GraphObservation], weights: np.ndarray, rng: RandomState | None) -> None:
"""Initialize sufficient statistics from a weighted graph batch."""
self.seq_update(x, weights, None)
[docs]
def combine(self, suff_stat: tuple[float, float]) -> ErdosRenyiGraphAccumulator:
"""Merge serialized edge-count sufficient statistics."""
total, successes = _validate_statistics(suff_stat)
combined = _validate_statistics((self.edge_opportunities + total, self.edge_count + successes))
self.edge_opportunities, self.edge_count = combined
return self
[docs]
def value(self) -> tuple[float, float]:
"""Return serialized edge-count sufficient statistics."""
return self.edge_opportunities, self.edge_count
[docs]
def from_value(self, x: tuple[float, float]) -> ErdosRenyiGraphAccumulator:
"""Restore accumulator state from serialized edge counts."""
self.edge_opportunities, self.edge_count = _validate_statistics(x)
return self
[docs]
def key_merge(self, stats_dict: dict[str, Any]) -> None:
"""Merge this accumulator under its configured sharing key."""
if self.keys is None:
return
if self.keys in stats_dict:
other = stats_dict[self.keys]
if not isinstance(other, ErdosRenyiGraphAccumulator):
raise ValueError("shared Erdos-Renyi key is bound to incompatible statistics.")
_require_matching_configuration(self, other)
other.combine(self.value())
else:
stats_dict[self.keys] = self
[docs]
def key_replace(self, stats_dict: dict[str, Any]) -> None:
"""Replace this accumulator from its configured sharing key."""
if self.keys is not None and self.keys in stats_dict:
other = stats_dict[self.keys]
if not isinstance(other, ErdosRenyiGraphAccumulator):
raise ValueError("shared Erdos-Renyi key is bound to incompatible statistics.")
_require_matching_configuration(self, other)
self.from_value(other.value())
[docs]
def acc_to_encoder(self) -> GraphDataEncoder:
"""Return the encoder associated with this accumulator."""
from mixle.data.sources.graph_source import GraphDataEncoder
return GraphDataEncoder(directed=self.directed)
[docs]
class ErdosRenyiGraphAccumulatorFactory(StatisticAccumulatorFactory):
"""Factory for ErdosRenyiGraphAccumulator."""
def __init__(
self,
directed: bool = False,
self_loops: bool = False,
num_nodes: int | None = None,
name: str | None = None,
keys: str | None = None,
) -> None:
self.directed = require_exact_bool(directed, "directed")
self.self_loops = require_exact_bool(self_loops, "self_loops")
self.num_nodes = _optional_num_nodes(num_nodes)
self.name = name
self.keys = keys
[docs]
def make(self) -> ErdosRenyiGraphAccumulator:
"""Create a fresh Erdos-Renyi graph accumulator."""
return ErdosRenyiGraphAccumulator(
directed=self.directed,
self_loops=self.self_loops,
num_nodes=self.num_nodes,
name=self.name,
keys=self.keys,
)
[docs]
class ErdosRenyiGraphEstimator(ParameterEstimator):
"""Estimate an Erdos-Renyi graph distribution from edge counts."""
def __init__(
self,
directed: bool = False,
self_loops: bool = False,
pseudo_count: float | None = None,
prior_p: float = 0.5,
num_nodes: int | None = None,
name: str | None = None,
keys: str | None = None,
) -> None:
self.directed = require_exact_bool(directed, "directed")
self.self_loops = require_exact_bool(self_loops, "self_loops")
self.pseudo_count = None if pseudo_count is None else _finite_nonnegative(pseudo_count, name="pseudo_count")
from mixle.data.sources.graph_source import _clip_prob
self.prior_p = _clip_prob(prior_p)
self.num_nodes = _optional_num_nodes(num_nodes)
self.name = name
self.keys = keys
[docs]
def accumulator_factory(self) -> ErdosRenyiGraphAccumulatorFactory:
"""Return the accumulator factory used by this estimator."""
return ErdosRenyiGraphAccumulatorFactory(
directed=self.directed,
self_loops=self.self_loops,
num_nodes=self.num_nodes,
name=self.name,
keys=self.keys,
)
[docs]
def estimate(self, nobs: float | None, suff_stat: tuple[float, float]) -> ErdosRenyiGraphDistribution:
"""Estimate the Bernoulli edge probability from edge-count statistics."""
total, successes = _validate_statistics(suff_stat)
if self.pseudo_count is not None:
successes += self.pseudo_count * self.prior_p
total += self.pseudo_count
if total <= 0.0:
raise ValueError("cannot estimate an Erdos-Renyi distribution without data or positive pseudo-count.")
p = successes / total
return ErdosRenyiGraphDistribution(
p,
directed=self.directed,
self_loops=self.self_loops,
num_nodes=self.num_nodes,
name=self.name,
keys=self.keys,
)