track-certify 0.1.0__py3-none-any.whl

track_certify/__init__.py

@@ -0,0 +1,49 @@
+"""track-certify: anytime joint certification of causal graphs and unknown
+intervention targets.
+
+Given a declared finite class of candidate causal graphs over a shared
+covariance, and K interventional environments whose targets and amplitudes
+are unknown, this package answers three questions with guarantees:
+
+* **Where to sample next, and when to stop.** :class:`Certifier` streams:
+  it emits the next environment to sample, consumes one fresh observation at
+  a time, and stops with a jointly certified (graph, target-vector) answer
+  whose error probability is at most ``delta`` under optional stopping -- or
+  with a refusal, never a guess.
+* **How hard the problem is, before sampling.** :func:`characteristic_time`
+  returns the instance's first-order sample requirement ``T*`` and the
+  optimal budget split across environments; :func:`coupling_tax` quantifies
+  how much harder joint certification is than its staged relaxations.
+* **What survives an estimated covariance.** :func:`certified_envelope` and
+  :func:`robust_certifier` implement the split-sample workflow with fully
+  explicit constants, refusing before starting when the envelope is
+  infeasible.
+
+Ready-made benchmark instances with theorem-known difficulty are in
+:func:`solvable_instance` and :func:`scalable_instance`; the audited
+simulation driver behind the paper's experiments is exposed as
+:func:`simulate`.
+
+The statistical core is vendored byte-identically from the audited research
+archive; ``tests/test_equivalence.py`` proves trajectory equivalence between
+the streaming interface and that driver.
+"""
+
+from .api import (Certifier, Envelope, Refusal, StepOutcome, build_model,
+                  certified_envelope)
+from .design import (CouplingReport, DesignReport, characteristic_time,
+                     coupling_tax)
+from .instances import Instance, scalable_instance, solvable_instance
+from .robust import RobustSetup, estimate_covariance, robust_certifier
+from . import _vendor  # noqa: F401
+from track_and_certify_general import run_general as simulate  # noqa: E402
+
+__all__ = [
+    "Certifier", "Envelope", "Refusal", "StepOutcome", "build_model",
+    "certified_envelope",
+    "DesignReport", "CouplingReport", "characteristic_time", "coupling_tax",
+    "Instance", "solvable_instance", "scalable_instance",
+    "RobustSetup", "estimate_covariance", "robust_certifier",
+    "simulate",
+]
+__version__ = "0.1.0"

track_certify/_vendor/__init__.py

@@ -0,0 +1,58 @@
+"""Vendored audited modules, byte-identical to the paper's evidence archive.
+
+The four modules import one another by their original top-level names
+(``assignment_oracle``, ``best_alternative``, ``finite_ray_model``,
+``track_and_certify_general``).  To keep them byte-identical, this package
+loads them from files in dependency order and registers them in
+``sys.modules`` under those names -- **without** touching ``sys.path``, so
+nothing else on the import path is shadowed.
+
+If one of these names is already imported from somewhere else (a user module
+with the same name), loading raises ``ImportError`` immediately with an
+explanation instead of silently mixing implementations.  Each module is also
+registered under ``track_certify._vendor.<name>`` for unambiguous access.
+"""
+
+import importlib.util
+import os
+import sys
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_ORDER = (
+    "assignment_oracle",
+    "best_alternative",
+    "finite_ray_model",
+    "track_and_certify_general",
+)
+
+
+def _load(name):
+    path = os.path.join(_HERE, name + ".py")
+    existing = sys.modules.get(name)
+    if existing is not None:
+        existing_file = getattr(existing, "__file__", None)
+        if existing_file is None or \
+                os.path.dirname(os.path.abspath(existing_file)) != _HERE:
+            raise ImportError(
+                f"track-certify vendors an audited module named {name!r}, "
+                f"but a different module with that name is already imported "
+                f"from {existing_file!r}. Import track_certify before the "
+                f"conflicting module, or rename it.")
+        return existing
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    try:
+        spec.loader.exec_module(module)
+    except Exception:
+        del sys.modules[name]
+        raise
+    return module
+
+
+for _name in _ORDER:
+    _mod = _load(_name)
+    sys.modules[__name__ + "." + _name] = _mod
+    globals()[_name] = _mod
+
+del _name, _mod

track_certify/_vendor/assignment_oracle.py

@@ -0,0 +1,100 @@
+"""Exact target-assignment oracle for joint graph--target certification.
+
+For a fixed graph, each row is an intervention environment and each column is
+a candidate target.  The certification loss is additive across rows and the
+main model requires distinct targets, so the best target vector is a
+rectangular linear assignment.  No hypothesis is pruned.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, Hashable, Iterable, Optional, Sequence, Tuple
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+
+@dataclass(frozen=True)
+class AssignmentSolution:
+    cost: float
+    targets: Tuple[int, ...]
+
+
+def _checked_cost(cost: np.ndarray) -> np.ndarray:
+    arr = np.asarray(cost, dtype=float)
+    if arr.ndim != 2:
+        raise ValueError("cost must be a K-by-d matrix")
+    k, d = arr.shape
+    if k == 0 or k > d:
+        raise ValueError("distinct-target assignment requires 1 <= K <= d")
+    if np.isnan(arr).any():
+        raise ValueError("cost contains NaN")
+    return arr
+
+
+def best_assignment(
+    cost: np.ndarray,
+    forbidden_pairs: Iterable[Tuple[int, int]] = (),
+) -> Optional[AssignmentSolution]:
+    """Return the minimum-cost injective target vector, or None if infeasible."""
+
+    arr = _checked_cost(cost).copy()
+    for e, target in forbidden_pairs:
+        if not (0 <= e < arr.shape[0] and 0 <= target < arr.shape[1]):
+            raise ValueError("forbidden pair is outside the cost matrix")
+        arr[e, target] = np.inf
+    try:
+        rows, cols = linear_sum_assignment(arr)
+    except ValueError:
+        return None
+    if len(rows) != arr.shape[0] or not np.isfinite(arr[rows, cols]).all():
+        return None
+    targets = np.empty(arr.shape[0], dtype=int)
+    targets[rows] = cols
+    return AssignmentSolution(float(arr[rows, cols].sum()), tuple(targets.tolist()))
+
+
+def best_distinct_assignment(
+    cost: np.ndarray,
+    incumbent_targets: Sequence[int],
+) -> Optional[AssignmentSolution]:
+    """Return the best injective assignment different from the incumbent.
+
+    Every different assignment omits at least one incumbent row--target pair.
+    Solving K assignments, each forbidding one such pair, is therefore exact.
+    """
+
+    arr = _checked_cost(cost)
+    incumbent = tuple(int(t) for t in incumbent_targets)
+    if len(incumbent) != arr.shape[0] or len(set(incumbent)) != len(incumbent):
+        raise ValueError("incumbent_targets must be an injective K-vector")
+    if any(t < 0 or t >= arr.shape[1] for t in incumbent):
+        raise ValueError("incumbent target is outside the cost matrix")
+
+    best: Optional[AssignmentSolution] = None
+    for e, target in enumerate(incumbent):
+        candidate = best_assignment(arr, forbidden_pairs=((e, target),))
+        if candidate is not None and (best is None or candidate.cost < best.cost):
+            best = candidate
+    return best
+
+
+def best_wrong_hypothesis(
+    graph_costs: Dict[Hashable, np.ndarray],
+    incumbent_graph: Hashable,
+    incumbent_targets: Sequence[int],
+) -> Tuple[Hashable, AssignmentSolution]:
+    """Return the exact lowest-loss graph--target answer excluding incumbent."""
+
+    if incumbent_graph not in graph_costs:
+        raise ValueError("incumbent graph is absent from graph_costs")
+    answer = None
+    for graph, cost in graph_costs.items():
+        if graph == incumbent_graph:
+            solution = best_distinct_assignment(cost, incumbent_targets)
+        else:
+            solution = best_assignment(cost)
+        if solution is not None and (answer is None or solution.cost < answer[1].cost):
+            answer = (graph, solution)
+    if answer is None:
+        raise ValueError("the declared class contains no alternative hypothesis")
+    return answer

track_certify/_vendor/best_alternative.py

@@ -0,0 +1,245 @@
+"""Unified best-hypothesis and best-alternative interface.
+
+The sequential algorithm has two consumers of the same optimization problem:
+
+* the stopping statistic selects the profiled incumbent and its best wrong
+  graph--target answer;
+* the allocation game asks nature for the best wrong answer under the current
+  environment weights.
+
+Both consumers call the functions in this module.  ``backend="enumeration"``
+is an exact small-instance reference implementation.  ``backend="assignment"``
+uses the rectangular assignment oracle and is the production implementation.
+The returned ``row_costs`` are always *unweighted*, which lets the allocation
+learner update with the full environment-wise reward vector.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import itertools
+from typing import Dict, Hashable, Mapping, Optional, Sequence, Tuple
+
+import numpy as np
+
+from assignment_oracle import best_assignment, best_wrong_hypothesis
+
+
+Hypothesis = Tuple[Hashable, Tuple[Hashable, ...]]
+
+
+@dataclass(frozen=True)
+class HypothesisSolution:
+    """Exact solution returned by either backend."""
+
+    graph: Hashable
+    targets: Tuple[Hashable, ...]
+    cost: float
+    row_costs: Tuple[float, ...]
+
+    @property
+    def hypothesis(self) -> Hypothesis:
+        return self.graph, self.targets
+
+
+def _validated_problem(
+    graph_costs: Mapping[Hashable, np.ndarray],
+    target_labels: Optional[Sequence[Hashable]],
+    environment_weights: Optional[Sequence[float]],
+) -> Tuple[Dict[Hashable, np.ndarray], Tuple[Hashable, ...], np.ndarray]:
+    if not graph_costs:
+        raise ValueError("graph_costs must contain at least one graph")
+
+    arrays: Dict[Hashable, np.ndarray] = {}
+    shape = None
+    for graph, cost in graph_costs.items():
+        arr = np.asarray(cost, dtype=float)
+        if arr.ndim != 2:
+            raise ValueError("every graph cost must be a K-by-d matrix")
+        if not np.isfinite(arr).all():
+            raise ValueError("graph costs must be finite")
+        if shape is None:
+            shape = arr.shape
+        elif arr.shape != shape:
+            raise ValueError("all graph cost matrices must have the same shape")
+        arrays[graph] = arr
+
+    assert shape is not None
+    k, d = shape
+    if k == 0 or k > d:
+        raise ValueError("injective targets require 1 <= K <= d")
+
+    labels = tuple(range(d)) if target_labels is None else tuple(target_labels)
+    if len(labels) != d or len(set(labels)) != d:
+        raise ValueError("target_labels must contain d distinct labels")
+
+    if environment_weights is None:
+        weights = np.ones(k, dtype=float)
+    else:
+        weights = np.asarray(environment_weights, dtype=float)
+        if weights.shape != (k,):
+            raise ValueError("environment_weights must have length K")
+        if not np.isfinite(weights).all() or np.any(weights < 0):
+            raise ValueError("environment_weights must be finite and nonnegative")
+        if weights.sum() <= 0:
+            raise ValueError("environment_weights must have positive total mass")
+    return arrays, labels, weights
+
+
+def _solution_from_columns(
+    graph: Hashable,
+    columns: Sequence[int],
+    arrays: Mapping[Hashable, np.ndarray],
+    labels: Sequence[Hashable],
+    weights: np.ndarray,
+) -> HypothesisSolution:
+    cols = np.asarray(columns, dtype=int)
+    row = arrays[graph][np.arange(len(cols)), cols]
+    return HypothesisSolution(
+        graph=graph,
+        targets=tuple(labels[j] for j in cols),
+        cost=float(weights @ row),
+        row_costs=tuple(float(x) for x in row),
+    )
+
+
+def _lex_key(solution: HypothesisSolution, graph_order: Mapping[Hashable, int],
+             label_order: Mapping[Hashable, int]) -> Tuple[int, Tuple[int, ...]]:
+    return graph_order[solution.graph], tuple(label_order[t] for t in solution.targets)
+
+
+def _prefer(
+    candidate: HypothesisSolution,
+    incumbent: Optional[HypothesisSolution],
+    graph_order: Mapping[Hashable, int],
+    label_order: Mapping[Hashable, int],
+    tie_atol: float,
+) -> bool:
+    if incumbent is None:
+        return True
+    scale = 1.0 + max(abs(candidate.cost), abs(incumbent.cost))
+    if candidate.cost < incumbent.cost - tie_atol * scale:
+        return True
+    if abs(candidate.cost - incumbent.cost) <= tie_atol * scale:
+        return _lex_key(candidate, graph_order, label_order) < _lex_key(
+            incumbent, graph_order, label_order
+        )
+    return False
+
+
+def _enumeration_solve(
+    arrays: Mapping[Hashable, np.ndarray],
+    labels: Sequence[Hashable],
+    weights: np.ndarray,
+    excluded: Optional[Hypothesis],
+    tie_atol: float,
+) -> HypothesisSolution:
+    graph_order = {graph: i for i, graph in enumerate(arrays)}
+    label_order = {label: i for i, label in enumerate(labels)}
+    best = None
+    for graph, arr in arrays.items():
+        for cols in itertools.permutations(range(arr.shape[1]), arr.shape[0]):
+            targets = tuple(labels[j] for j in cols)
+            if excluded is not None and (graph, targets) == excluded:
+                continue
+            candidate = _solution_from_columns(graph, cols, arrays, labels, weights)
+            if np.isfinite(candidate.cost) and _prefer(
+                candidate, best, graph_order, label_order, tie_atol
+            ):
+                best = candidate
+    if best is None:
+        raise ValueError("the declared class contains no feasible hypothesis")
+    return best
+
+
+def _assignment_incumbent(
+    arrays: Mapping[Hashable, np.ndarray],
+    labels: Sequence[Hashable],
+    weights: np.ndarray,
+    tie_atol: float,
+) -> HypothesisSolution:
+    graph_order = {graph: i for i, graph in enumerate(arrays)}
+    label_order = {label: i for i, label in enumerate(labels)}
+    best = None
+    for graph, arr in arrays.items():
+        assignment = best_assignment(arr * weights[:, None])
+        if assignment is None:
+            continue
+        candidate = _solution_from_columns(
+            graph, assignment.targets, arrays, labels, weights
+        )
+        if _prefer(candidate, best, graph_order, label_order, tie_atol):
+            best = candidate
+    if best is None:
+        raise ValueError("the declared class contains no feasible hypothesis")
+    return best
+
+
+def _assignment_alternative(
+    arrays: Mapping[Hashable, np.ndarray],
+    labels: Sequence[Hashable],
+    weights: np.ndarray,
+    incumbent: Hypothesis,
+) -> HypothesisSolution:
+    incumbent_graph, incumbent_targets = incumbent
+    if incumbent_graph not in arrays:
+        raise ValueError("incumbent graph is absent from graph_costs")
+    label_to_col = {label: i for i, label in enumerate(labels)}
+    try:
+        incumbent_columns = tuple(label_to_col[t] for t in incumbent_targets)
+    except KeyError as error:
+        raise ValueError("incumbent target is absent from target_labels") from error
+    if len(incumbent_columns) != next(iter(arrays.values())).shape[0]:
+        raise ValueError("incumbent target vector must have length K")
+    if len(set(incumbent_columns)) != len(incumbent_columns):
+        raise ValueError("incumbent target vector must be injective")
+
+    weighted = {graph: arr * weights[:, None] for graph, arr in arrays.items()}
+    graph, assignment = best_wrong_hypothesis(
+        weighted, incumbent_graph, incumbent_columns
+    )
+    return _solution_from_columns(graph, assignment.targets, arrays, labels, weights)
+
+
+def best_hypothesis(
+    graph_costs: Mapping[Hashable, np.ndarray],
+    *,
+    environment_weights: Optional[Sequence[float]] = None,
+    target_labels: Optional[Sequence[Hashable]] = None,
+    backend: str = "assignment",
+    tie_atol: float = 1e-12,
+) -> HypothesisSolution:
+    """Return the exact minimum-cost graph--target hypothesis."""
+
+    arrays, labels, weights = _validated_problem(
+        graph_costs, target_labels, environment_weights
+    )
+    if backend == "assignment":
+        return _assignment_incumbent(arrays, labels, weights, tie_atol)
+    if backend == "enumeration":
+        return _enumeration_solve(arrays, labels, weights, None, tie_atol)
+    raise ValueError("backend must be 'assignment' or 'enumeration'")
+
+
+def best_alternative(
+    graph_costs: Mapping[Hashable, np.ndarray],
+    incumbent: Hypothesis,
+    *,
+    environment_weights: Optional[Sequence[float]] = None,
+    target_labels: Optional[Sequence[Hashable]] = None,
+    backend: str = "assignment",
+    tie_atol: float = 1e-12,
+) -> HypothesisSolution:
+    """Return the exact best graph--target answer different from ``incumbent``."""
+
+    arrays, labels, weights = _validated_problem(
+        graph_costs, target_labels, environment_weights
+    )
+    normalized_incumbent = (incumbent[0], tuple(incumbent[1]))
+    if backend == "assignment":
+        return _assignment_alternative(arrays, labels, weights, normalized_incumbent)
+    if backend == "enumeration":
+        return _enumeration_solve(
+            arrays, labels, weights, normalized_incumbent, tie_atol
+        )
+    raise ValueError("backend must be 'assignment' or 'enumeration'")

track_certify/_vendor/finite_ray_model.py

@@ -0,0 +1,198 @@
+"""Finite Gaussian response-ray models for staged Track-and-Certify smoke tests.
+
+Each candidate graph is a complete DAG induced by a declared topological order.
+All candidates therefore belong to the same complete-skeleton MEC and share one
+known covariance matrix.  A Cholesky column is a structural propagation vector
+times a target-specific residual standard deviation; because amplitudes are
+profiled, it represents exactly the same response ray.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import itertools
+from typing import Dict, Iterable, Mapping, Sequence, Tuple
+
+import numpy as np
+from scipy.optimize import linprog
+
+from best_alternative import Hypothesis
+
+
+@dataclass(frozen=True)
+class FiniteRayModel:
+    sigma: np.ndarray
+    graph_orders: Mapping[str, Tuple[int, ...]]
+    directions: Mapping[str, np.ndarray]
+
+    @property
+    def dimension(self) -> int:
+        return int(self.sigma.shape[0])
+
+    @property
+    def graph_labels(self) -> Tuple[str, ...]:
+        return tuple(self.graph_orders)
+
+    @property
+    def target_labels(self) -> Tuple[int, ...]:
+        return tuple(range(self.dimension))
+
+
+def correlation_matrix(dimension: int, rng: np.random.Generator,
+                       ridge: float) -> np.ndarray:
+    factor = rng.normal(size=(dimension, dimension))
+    covariance = factor @ factor.T + ridge * np.eye(dimension)
+    scale = 1.0 / np.sqrt(np.diag(covariance))
+    return covariance * np.outer(scale, scale)
+
+
+def rays_for_order(sigma: np.ndarray, order: Sequence[int]) -> np.ndarray:
+    dimension = sigma.shape[0]
+    order = tuple(int(node) for node in order)
+    if sorted(order) != list(range(dimension)):
+        raise ValueError("order must be a permutation of all nodes")
+    cholesky = np.linalg.cholesky(sigma[np.ix_(order, order)])
+    directions = np.zeros((dimension, dimension))
+    for position, node in enumerate(order):
+        column = np.zeros(dimension)
+        column[list(order)] = cholesky[:, position]
+        directions[:, node] = column
+    return directions
+
+
+def model_from_orders(sigma: np.ndarray,
+                      graph_orders: Mapping[str, Sequence[int]]) -> FiniteRayModel:
+    sigma = np.asarray(sigma, dtype=float)
+    if sigma.ndim != 2 or sigma.shape[0] != sigma.shape[1]:
+        raise ValueError("sigma must be square")
+    if not np.allclose(sigma, sigma.T, atol=1e-12, rtol=1e-12):
+        raise ValueError("sigma must be symmetric")
+    if np.linalg.eigvalsh(sigma).min() <= 0:
+        raise ValueError("sigma must be positive definite")
+    if not graph_orders:
+        raise ValueError("at least one graph is required")
+    normalized = {
+        str(graph): tuple(int(node) for node in order)
+        for graph, order in graph_orders.items()
+    }
+    directions = {
+        graph: rays_for_order(sigma, order)
+        for graph, order in normalized.items()
+    }
+    return FiniteRayModel(sigma=sigma, graph_orders=normalized, directions=directions)
+
+
+def profiled_divergence(mu: np.ndarray, ray: np.ndarray,
+                        sigma_inv: np.ndarray) -> float:
+    a = mu @ sigma_inv @ mu
+    b = mu @ sigma_inv @ ray
+    c = ray @ sigma_inv @ ray
+    return float(0.5 * (a - b * b / c))
+
+
+def true_means(model: FiniteRayModel, true_hypothesis: Hypothesis,
+               amplitudes: Sequence[float]) -> Tuple[np.ndarray, ...]:
+    graph, targets = true_hypothesis
+    amplitudes = np.asarray(amplitudes, dtype=float)
+    if graph not in model.directions:
+        raise ValueError("true graph is absent from the declared class")
+    if amplitudes.shape != (len(targets),):
+        raise ValueError("one amplitude is required per environment")
+    if len(set(targets)) != len(targets):
+        raise ValueError("targets must be injective")
+    return tuple(
+        amplitudes[e] * model.directions[graph][:, target]
+        for e, target in enumerate(targets)
+    )
+
+
+def alternative_kind(candidate: Hypothesis, truth: Hypothesis) -> str:
+    graph_changed = candidate[0] != truth[0]
+    targets_changed = tuple(candidate[1]) != tuple(truth[1])
+    if graph_changed and targets_changed:
+        return "coupled"
+    if graph_changed:
+        return "graph"
+    if targets_changed:
+        return "target"
+    return "truth"
+
+
+def alternative_divergences(model: FiniteRayModel, true_hypothesis: Hypothesis,
+                            amplitudes: Sequence[float]):
+    means = true_means(model, true_hypothesis, amplitudes)
+    sigma_inv = np.linalg.inv(model.sigma)
+    n_environments = len(true_hypothesis[1])
+    rows = []
+    kinds = []
+    hypotheses = []
+    for graph in model.graph_labels:
+        for targets in itertools.permutations(model.target_labels, n_environments):
+            candidate = (graph, tuple(targets))
+            if candidate == true_hypothesis:
+                continue
+            rows.append(
+                [
+                    profiled_divergence(
+                        means[e], model.directions[graph][:, targets[e]], sigma_inv
+                    )
+                    for e in range(n_environments)
+                ]
+            )
+            kinds.append(alternative_kind(candidate, true_hypothesis))
+            hypotheses.append(candidate)
+    return np.asarray(rows), tuple(kinds), tuple(hypotheses)
+
+
+def characteristic_time(divergences: np.ndarray):
+    divergences = np.asarray(divergences, dtype=float)
+    if divergences.ndim != 2 or divergences.shape[0] == 0:
+        return np.inf, None
+    n_alternatives, n_environments = divergences.shape
+    objective = np.zeros(n_environments + 1)
+    objective[-1] = -1.0
+    result = linprog(
+        objective,
+        A_ub=np.hstack([-divergences, np.ones((n_alternatives, 1))]),
+        b_ub=np.zeros(n_alternatives),
+        A_eq=np.array([[*([1.0] * n_environments), 0.0]]),
+        b_eq=[1.0],
+        bounds=[(0.0, None)] * n_environments + [(None, None)],
+        method="highs",
+    )
+    if not result.success or result.x[-1] <= 1e-12:
+        return np.inf, None
+    return float(1.0 / result.x[-1]), tuple(float(x) for x in result.x[:-1])
+
+
+def diagnose_characteristic_times(model: FiniteRayModel,
+                                  true_hypothesis: Hypothesis,
+                                  amplitudes: Sequence[float]):
+    divergences, kinds, _ = alternative_divergences(
+        model, true_hypothesis, amplitudes
+    )
+    output = {}
+    for kind in ("graph", "target", "coupled"):
+        mask = np.asarray([label == kind for label in kinds])
+        time, weights = characteristic_time(divergences[mask])
+        output[kind] = {"tstar": time, "weights": weights}
+    time, weights = characteristic_time(divergences)
+    output["all"] = {"tstar": time, "weights": weights}
+    output["gamma_uniform"] = float(
+        np.min(divergences @ (np.ones(divergences.shape[1]) / divergences.shape[1]))
+    )
+    return output
+
+
+def graph_cost_matrices(model: FiniteRayModel, means: Sequence[np.ndarray]):
+    sigma_inv = np.linalg.inv(model.sigma)
+    costs: Dict[str, np.ndarray] = {}
+    for graph in model.graph_labels:
+        matrix = np.empty((len(means), model.dimension))
+        for e, mu in enumerate(means):
+            for target in model.target_labels:
+                matrix[e, target] = profiled_divergence(
+                    mu, model.directions[graph][:, target], sigma_inv
+                )
+        costs[graph] = matrix
+    return costs