sofic 0.1.0__py3-none-any.whl

sofic/inference/bayesian/diversity.py

@@ -0,0 +1,254 @@
+"""Posterior diversity diagnostics for Bayesian epsilon-machine model comparison.
+
+Two complementary notions of posterior spread are tracked:
+
+* **Machine diversity** — Shannon entropy of the topology weights returned by
+  :meth:`~sofic.inference.bayesian.comparison.ModelComparisonEM.model_probabilities`.
+  This measures uncertainty over *presentations* (topologies), not processes.
+
+* **Process diversity** — weighted Jensen–Shannon divergence (JSD) over length-:math:`L`
+  word distributions, one per posterior component. When many high-weight machines
+  generate nearly the same stochastic process, process diversity can be much smaller
+  than machine diversity.
+
+Word length conventions
+-----------------------
+For an :math:`n`-state presentation, classical HMM identification (Paz 1971;
+Finesso 1991) uses the full length-:math:`(2n-1)` word distribution. Some
+minimal-realization algorithms use a conservative window of :math:`2n+1`.
+Upper's rank-growing history/future word lists (see
+:func:`~sofic.generators.process_equivalence.is_equal_process`) provide a
+data-driven alternative. All topologies in a comparison share the same :math:`L`,
+chosen from the largest state count in the posterior.
+
+JSD is computed in bits via ``dit``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from itertools import product
+from typing import TYPE_CHECKING, Any, Literal
+
+import numpy as np
+
+from sofic.generators.base import HiddenMarkovModel
+from sofic.generators.process_equivalence import _HistoryFutureWordList
+from sofic.generators.words import hmm_words_of_length
+from sofic.inference.bayesian.counts import BayesianInferenceError
+from sofic.inference.bayesian.epsilon import EpsilonMachinePosterior
+
+if TYPE_CHECKING:
+    from sofic.inference.bayesian.comparison import ModelComparisonEM
+
+_TOL = 1e-15
+WordLengthConvention = Literal["paz", "conservative", "upper_list"]
+ProcessDiversityMethod = Literal["posterior_mean", "monte_carlo"]
+
+
+@dataclass(frozen=True)
+class PosteriorDiversityResult:
+    """Posterior machine and process diversity diagnostics."""
+
+    process_diversity: float
+    machine_diversity: float
+    word_length: int
+    method: ProcessDiversityMethod
+    n_components: int
+
+
+def _require_dit():
+    try:
+        from dit.divergences.jensen_shannon_divergence import jensen_shannon_divergence_pmf
+    except ImportError as exc:
+        raise ImportError("dit is required for posterior process diversity; install with `pip install dit`") from exc
+    return jensen_shannon_divergence_pmf
+
+
+def _comparison_alphabet(comparison: ModelComparisonEM) -> tuple[Any, ...]:
+    alphabets: set[Any] = set()
+    for posterior in comparison.em_dict.values():
+        alphabets.update(posterior.machine.observation_alphabet)
+    if not alphabets:
+        raise BayesianInferenceError("posterior has no viable machines")
+    return tuple(sorted(alphabets, key=repr))
+
+
+def _max_state_count(comparison: ModelComparisonEM) -> int:
+    if not comparison.em_dict:
+        raise BayesianInferenceError("posterior has no viable machines")
+    return max(len(posterior.dirichlet.nodes) for posterior in comparison.em_dict.values())
+
+
+def _upper_list_word_length(machine: HiddenMarkovModel) -> int:
+    hf = _HistoryFutureWordList.from_hmm(machine)
+    future_words = hf.future_word_list()
+    history_words = hf.history_word_list()
+    candidates = [len(word) for word in (*future_words, *history_words)]
+    return max(candidates) if candidates else 0
+
+
+def process_identification_word_length(
+    comparison: ModelComparisonEM,
+    *,
+    convention: WordLengthConvention | str = "paz",
+    word_length: int | None = None,
+) -> int:
+    """Return the word length used to compare processes in a model comparison."""
+    if word_length is not None:
+        length = int(word_length)
+        if length < 0:
+            raise BayesianInferenceError("word_length must be nonnegative")
+        return length
+
+    n_max = _max_state_count(comparison)
+    if convention == "paz":
+        return max(0, 2 * n_max - 1)
+    if convention == "conservative":
+        return 2 * n_max + 1
+    if convention == "upper_list":
+        lengths = []
+        for posterior in comparison.em_dict.values():
+            start_probs = posterior.start_node_probabilities()
+            if not start_probs:
+                continue
+            start_node = max(start_probs, key=start_probs.get)
+            mean_machine = posterior.posterior_mean_machine(start_node)
+            if mean_machine is not None:
+                lengths.append(_upper_list_word_length(mean_machine))
+        if not lengths:
+            return max(0, 2 * n_max - 1)
+        return max(lengths)
+    raise BayesianInferenceError(f"unknown word-length convention: {convention!r}")
+
+
+def word_distribution_to_pmf(
+    distribution: dict[tuple[Any, ...], float],
+    alphabet: Sequence[Any],
+    length: int,
+) -> np.ndarray:
+    """Align a sparse word distribution to a dense PMF over ``alphabet**length``."""
+    symbols = tuple(alphabet)
+    if length == 0:
+        total = sum(float(prob) for prob in distribution.values())
+        return np.array([total if total > _TOL else 0.0], dtype=float)
+
+    outcomes = list(product(symbols, repeat=length))
+    pmf = np.zeros(len(outcomes), dtype=float)
+    for index, word in enumerate(outcomes):
+        prob = distribution.get(word, 0.0)
+        if prob > _TOL:
+            pmf[index] = float(prob)
+
+    total = float(pmf.sum())
+    if total > _TOL and abs(total - 1.0) > _TOL:
+        pmf /= total
+    return pmf
+
+
+def machine_diversity(comparison: ModelComparisonEM) -> float:
+    """Shannon entropy (bits) of the topology posterior weights."""
+    probs = np.array(list(comparison.model_probabilities().values()), dtype=float)
+    if probs.size == 0:
+        raise BayesianInferenceError("posterior has no viable machines")
+    positive = probs[probs > 0.0]
+    if positive.size <= 1:
+        return 0.0
+    return float(-np.sum(positive * np.log2(positive)))
+
+
+def posterior_mean_word_distribution(
+    posterior: EpsilonMachinePosterior,
+    length: int,
+) -> dict[tuple[Any, ...], float]:
+    """Start-marginalized word distribution from posterior-mean transition probabilities."""
+    if length < 0:
+        raise BayesianInferenceError("length must be nonnegative")
+
+    start_probs = posterior.start_node_probabilities()
+    if not start_probs:
+        return {}
+
+    distribution: dict[tuple[Any, ...], float] = {}
+    for start_node, start_weight in start_probs.items():
+        if start_weight <= _TOL:
+            continue
+        machine = posterior.posterior_mean_machine(start_node)
+        if machine is None:
+            continue
+        words = hmm_words_of_length(machine, length)
+        for word, prob in words.items():
+            distribution[word] = distribution.get(word, 0.0) + start_weight * float(prob)
+    return distribution
+
+
+def _jsd_from_word_distributions(
+    distributions: Sequence[dict[tuple[Any, ...], float]],
+    weights: Sequence[float],
+    alphabet: Sequence[Any],
+    length: int,
+) -> float:
+    jensen_shannon_divergence_pmf = _require_dit()
+    if not distributions:
+        raise BayesianInferenceError("no word distributions supplied")
+    if len(distributions) != len(weights):
+        raise BayesianInferenceError("number of weights must match number of distributions")
+
+    pmfs = np.vstack([word_distribution_to_pmf(dist, alphabet, length) for dist in distributions])
+    weight_array = np.asarray(weights, dtype=float)
+    if weight_array.sum() <= _TOL:
+        raise BayesianInferenceError("posterior weights must sum to a positive value")
+    return float(jensen_shannon_divergence_pmf(pmfs, weight_array))
+
+
+def posterior_process_diversity(
+    comparison: ModelComparisonEM,
+    *,
+    method: ProcessDiversityMethod = "posterior_mean",
+    n_samples: int = 500,
+    rng: np.random.Generator | None = None,
+    convention: WordLengthConvention | str = "paz",
+    word_length: int | None = None,
+) -> PosteriorDiversityResult:
+    """Compute weighted JSD over posterior word distributions."""
+    if not comparison.em_dict:
+        raise BayesianInferenceError("posterior has no viable machines")
+
+    length = process_identification_word_length(comparison, convention=convention, word_length=word_length)
+    machine_div = machine_diversity(comparison)
+    alphabet = _comparison_alphabet(comparison)
+
+    if method == "posterior_mean":
+        model_probs = comparison.model_probabilities()
+        names = list(model_probs)
+        distributions = [posterior_mean_word_distribution(comparison.em_dict[name], length) for name in names]
+        weights = [model_probs[name] for name in names]
+        process_div = _jsd_from_word_distributions(distributions, weights, alphabet, length)
+        return PosteriorDiversityResult(
+            process_diversity=process_div,
+            machine_diversity=machine_div,
+            word_length=length,
+            method="posterior_mean",
+            n_components=len(names),
+        )
+
+    if method == "monte_carlo":
+        if n_samples <= 0:
+            raise BayesianInferenceError("n_samples must be positive")
+        generator = rng if rng is not None else np.random.default_rng()
+        distributions = []
+        for _ in range(n_samples):
+            _start, machine = comparison.generate_sample(rng=generator)
+            distributions.append(hmm_words_of_length(machine, length))
+        weights = [1.0 / n_samples] * n_samples
+        process_div = _jsd_from_word_distributions(distributions, weights, alphabet, length)
+        return PosteriorDiversityResult(
+            process_diversity=process_div,
+            machine_diversity=machine_div,
+            word_length=length,
+            method="monte_carlo",
+            n_components=n_samples,
+        )
+
+    raise BayesianInferenceError(f"unknown process diversity method: {method!r}")

sofic/inference/bayesian/epsilon.py

@@ -0,0 +1,270 @@
+"""Conjugate Bayesian inference for fixed unifilar generator topologies."""
+
+from __future__ import annotations
+
+from collections.abc import Hashable, Mapping, Sequence
+from typing import Any
+
+import numpy as np
+from scipy.special import logsumexp
+
+from sofic.generators.epsilon_machine import EpsilonMachine
+from sofic.generators.mealy import MealyHMM
+from sofic.graph import ATTR_EMISSION, ATTR_PROB
+from sofic.inference.bayesian.counts import (
+    BayesianInferenceError,
+    PathCountEM,
+    dirichlet_multinomial_log_evidence,
+    scan_unifilar_topology,
+)
+
+
+class DirichletDistributionEM:
+    """Product-of-Dirichlets prior/posterior for a fixed unifilar topology."""
+
+    def __init__(self, machine: MealyHMM, data: Sequence[Any] | None = None, state_path: bool = False):
+        self.machine = machine
+        self.alphas: dict[Hashable | tuple[Hashable, Any], float] = {}
+        self.edges: list[tuple[Hashable, Any]] = []
+        self.valid_edges: list[tuple[Hashable, Any]] = []
+        self.nodes: list[Hashable] = list(machine.states())
+        self.trace: dict[tuple[Hashable, Any], Hashable] = {}
+        self.data: PathCountEM | None = None
+        self.deterministic = False
+        self._process_machine_topology()
+        self.valid_startnodes = list(self.nodes)
+        self._generate_uniform_alphas()
+        if data is not None:
+            self.data = PathCountEM(machine, data, state_path=state_path)
+            self.valid_startnodes = self.data.get_possible_start_nodes()
+
+    def _process_machine_topology(self) -> None:
+        self.trace, self.edges = scan_unifilar_topology(self.machine)
+        outgoing: dict[Hashable, list[tuple[Hashable, Any]]] = {}
+        for edge in self.edges:
+            outgoing.setdefault(edge[0], []).append(edge)
+        for _source, edges in outgoing.items():
+            if len(edges) > 1:
+                self.valid_edges.extend(edges)
+            else:
+                edge = edges[0]
+                prob = self._topology_probability(edge)
+                if not np.isclose(prob, 1.0):
+                    self.valid_edges.append(edge)
+        self.valid_edges.sort(key=repr)
+        self.deterministic = not self.valid_edges
+
+    def _topology_probability(self, edge: tuple[Hashable, Any]) -> float:
+        source, symbol = edge
+        for transition in self.machine.graph.out_transitions(source):
+            if transition.data.get(ATTR_EMISSION) == symbol:
+                return float(transition.data.get(ATTR_PROB, 0.0))
+        return 0.0
+
+    def _generate_uniform_alphas(self) -> None:
+        self.alphas.clear()
+        for edge in self.valid_edges:
+            source, _symbol = edge
+            self.alphas[edge] = 1.0
+            self.alphas[source] = self.alphas.get(source, 0.0) + 1.0
+
+    def get_edges(self) -> list[tuple[Hashable, Any]]:
+        return list(self.edges)
+
+    def get_nodes(self) -> list[Hashable]:
+        return list(self.nodes)
+
+    def get_possible_start_nodes(self) -> list[Hashable]:
+        return list(self.valid_startnodes)
+
+    def get_edge_alpha(self, start_node: Hashable, edge: tuple[Hashable, Any]) -> float | None:
+        del start_node
+        return self.alphas.get(edge)
+
+    def get_node_alpha(self, start_node: Hashable, node: Hashable) -> float | None:
+        del start_node
+        return self.alphas.get(node)
+
+    def get_edge_count(self, start_node: Hashable, edge: tuple[Hashable, Any]) -> int | None:
+        return None if self.data is None else self.data.get_edge_count(start_node, edge)
+
+    def get_node_count(self, start_node: Hashable, node: Hashable) -> int | None:
+        return None if self.data is None else self.data.get_node_count(start_node, node)
+
+    def get_last_node(self, start_node: Hashable) -> Hashable | None:
+        return None if self.data is None else self.data.get_last_node(start_node)
+
+    def get_state_path(self, start_node: Hashable) -> tuple[Hashable, ...]:
+        return () if self.data is None else self.data.get_state_path(start_node)
+
+    def log_evidence_start_node(self, start_node: Hashable) -> float:
+        if start_node not in self.valid_startnodes:
+            return -np.inf
+        rows: dict[Hashable, list[tuple[float, float]]] = {}
+        for edge in self.valid_edges:
+            source, _symbol = edge
+            alpha = self.get_edge_alpha(start_node, edge)
+            alpha_row = self.get_node_alpha(start_node, source)
+            if alpha is None or alpha_row is None:
+                raise BayesianInferenceError("missing Dirichlet alpha")
+            count = self.get_edge_count(start_node, edge) or 0
+            rows.setdefault(source, []).append((alpha, count))
+        evidence = 0.0
+        for source, cells in rows.items():
+            alpha_row = self.get_node_alpha(start_node, source)
+            count_row = self.get_node_count(start_node, source) or 0
+            evidence += dirichlet_multinomial_log_evidence(alpha_row, count_row, cells)
+        return float(evidence)
+
+    def mean_edge_probability(self, start_node: Hashable, edge: tuple[Hashable, Any]) -> float | None:
+        if start_node not in self.valid_startnodes:
+            return None
+        if edge in self.edges and edge not in self.valid_edges:
+            return 1.0
+        count = self.get_edge_count(start_node, edge) or 0
+        row_count = self.get_node_count(start_node, edge[0]) or 0
+        alpha = self.get_edge_alpha(start_node, edge)
+        row_alpha = self.get_node_alpha(start_node, edge[0])
+        if alpha is None or row_alpha is None:
+            return None
+        return float((count + alpha) / (row_count + row_alpha))
+
+    def set_edge_alpha(self, edge: tuple[Hashable, Any], value: float) -> None:
+        if edge not in self.valid_edges:
+            raise BayesianInferenceError("cannot set alpha for deterministic edge")
+        self.alphas[edge] = float(value)
+        for node in self.nodes:
+            self.alphas[node] = 0.0
+        for valid_edge in self.valid_edges:
+            self.alphas[valid_edge[0]] += self.alphas[valid_edge]
+
+    def _machine_from_probabilities(
+        self, start_node: Hashable, probabilities: Mapping[tuple[Hashable, Any], float], name: str
+    ) -> MealyHMM:
+        machine = MealyHMM(observation_alphabet=getattr(self.machine, "observation_alphabet", frozenset()))
+        machine.name = name
+        initial_state = self.get_last_node(start_node) if self.data is not None else start_node
+        machine.initial_distribution = {initial_state: 1.0} if initial_state is not None else {}
+        for node in self.nodes:
+            machine.graph.add_state(node)
+        for edge in self.edges:
+            source, symbol = edge
+            prob = probabilities[edge]
+            machine.graph.add_transition(source, self.trace[edge], **{ATTR_EMISSION: symbol, ATTR_PROB: prob})
+        machine.validate()
+        if machine.is_unifilar():
+            eps = EpsilonMachine.from_networkx(
+                machine.to_networkx(),
+                initial_distribution=machine.initial_distribution,
+                observation_alphabet=machine.observation_alphabet,
+            )
+            eps.name = name
+            eps.validate()
+            return eps
+        return machine
+
+    def posterior_mean_machine(self, start_node: Hashable) -> MealyHMM | None:
+        if start_node not in self.valid_startnodes:
+            return None
+        probabilities = {}
+        for edge in self.edges:
+            prob = self.mean_edge_probability(start_node, edge)
+            if prob is None:
+                raise BayesianInferenceError(f"missing probability for edge {edge!r}")
+            probabilities[edge] = prob
+        return self._machine_from_probabilities(
+            start_node, probabilities, f"Posterior Mean Machine, Start Node: {start_node}"
+        )
+
+    def generate_sample(self, start_node: Hashable, rng: np.random.Generator | None = None) -> MealyHMM | None:
+        if start_node not in self.valid_startnodes:
+            return None
+        generator = rng if rng is not None else np.random.default_rng()
+        probabilities: dict[tuple[Hashable, Any], float] = {}
+        grouped: dict[Hashable, list[tuple[Hashable, Any]]] = {}
+        for edge in self.edges:
+            if edge in self.valid_edges:
+                grouped.setdefault(edge[0], []).append(edge)
+            else:
+                probabilities[edge] = 1.0
+        for _source, edges in grouped.items():
+            alpha = []
+            for edge in edges:
+                a = self.get_edge_alpha(start_node, edge)
+                if a is None:
+                    raise BayesianInferenceError("missing Dirichlet alpha")
+                alpha.append(a + (self.get_edge_count(start_node, edge) or 0))
+            sample = generator.dirichlet(np.asarray(alpha, dtype=float))
+            for edge, prob in zip(edges, sample, strict=True):
+                probabilities[edge] = float(prob)
+        return self._machine_from_probabilities(start_node, probabilities, f"Sampled Machine, Start Node: {start_node}")
+
+
+class EpsilonMachinePosterior:
+    """Posterior over parameters and unknown start state for a topology."""
+
+    def __init__(
+        self,
+        machine: MealyHMM,
+        data: Sequence[Any] | None = None,
+        start_dist: Mapping[Hashable, float] | None = None,
+        state_path: bool = False,
+    ):
+        self.machine = machine
+        self.dirichlet = DirichletDistributionEM(machine, data, state_path=state_path)
+        if start_dist is None:
+            nodes = self.dirichlet.get_nodes()
+            self.start_dist = {node: 1.0 / len(nodes) for node in nodes} if nodes else {}
+        else:
+            self.start_dist = dict(start_dist)
+
+    def log_evidence(self) -> float:
+        terms = []
+        for node in self.dirichlet.get_possible_start_nodes():
+            prior = self.start_dist.get(node, 0.0)
+            if prior > 0:
+                terms.append(np.log(prior) + self.dirichlet.log_evidence_start_node(node))
+        return float(logsumexp(terms)) if terms else -np.inf
+
+    def probability_start_node(self, start_node: Hashable) -> float:
+        if start_node not in self.dirichlet.get_possible_start_nodes():
+            return 0.0
+        log_norm = self.log_evidence()
+        prior = self.start_dist.get(start_node, 0.0)
+        if prior <= 0:
+            return 0.0
+        return float(np.exp(np.log(prior) + self.dirichlet.log_evidence_start_node(start_node) - log_norm))
+
+    def start_node_probabilities(self) -> dict[Hashable, float]:
+        return {node: self.probability_start_node(node) for node in self.dirichlet.get_possible_start_nodes()}
+
+    def sample_start_node(self, rng: np.random.Generator | None = None) -> Hashable:
+        generator = rng if rng is not None else np.random.default_rng()
+        probs = self.start_node_probabilities()
+        nodes = tuple(probs)
+        weights = np.array([probs[node] for node in nodes], dtype=float)
+        return nodes[int(generator.choice(len(nodes), p=weights))]
+
+    def generate_sample(self, rng: np.random.Generator | None = None) -> tuple[Hashable, MealyHMM]:
+        generator = rng if rng is not None else np.random.default_rng()
+        start = self.sample_start_node(generator)
+        machine = self.dirichlet.generate_sample(start, rng=generator)
+        if machine is None:
+            raise BayesianInferenceError("sampled impossible start node")
+        return start, machine
+
+    def posterior_mean_machine(self, start_node: Hashable | None = None) -> MealyHMM | None:
+        if start_node is None:
+            probs = self.start_node_probabilities()
+            if not probs:
+                return None
+            start_node = max(probs, key=probs.get)
+        return self.dirichlet.posterior_mean_machine(start_node)
+
+    def as_pymc_model(self, start_node: Hashable | None = None) -> Any:
+        from sofic.inference.bayesian.pymc_backend import epsilon_machine_model
+
+        return epsilon_machine_model(self, start_node=start_node)
+
+
+InferEM = EpsilonMachinePosterior