repro-bundle-v0/src/graphs.py

"""Time-varying directed graphs and compatible row-/column-stochastic mixing matrices.

Follows the paper's Section 2 notation and Assumption 3.2: at each iteration k we
draw a directed ring (guarantees strong connectivity, Assumption 3.1) plus up to
`extra_edges` additional random directed edges (an Erdos-Renyi-style augmentation,
matching the paper's "ER model + directed rings" dynamic-topology construction in
Section 4). Degrees are capped so the resulting mixing weights admit a *uniform*
lower bound a, b > 0 across all k, as Assumption 3.2 requires.
"""
from __future__ import annotations

import numpy as np


class TimeVaryingDigraphs:
    """Generates a fresh strongly-connected directed graph each call to `step`,
    and the row-stochastic (A) / column-stochastic (B) matrices compatible with it.
    """

    def __init__(self, n: int, extra_edges: int, seed: int = 0):
        self.n = n
        self.extra_edges = extra_edges
        self.rng = np.random.default_rng(seed)
        # Uniform lower bounds guaranteed by construction (ring + <= extra_edges extra
        # in-/out-edges per node => max in/out degree <= 1 + extra_edges, plus self-loop).
        max_degree_plus_self = 2 + extra_edges
        self.a_min = 1.0 / max_degree_plus_self
        self.b_min = 1.0 / max_degree_plus_self

    def step(self):
        n = self.n
        # Directed ring i -> i+1 (mod n) guarantees strong connectivity on its own;
        # each node gets up to `extra_edges` extra random outgoing edges on top.
        out_neighbors = [set() for _ in range(n)]
        for i in range(n):
            out_neighbors[i].add((i + 1) % n)  # ring edge i -> i+1
        for i in range(n):
            possible = [j for j in range(n) if j != i and j not in out_neighbors[i]]
            m = min(self.extra_edges, len(possible))
            if m > 0:
                chosen = self.rng.choice(possible, size=m, replace=False)
                for j in chosen:
                    out_neighbors[i].add(int(j))

        in_neighbors = [set() for _ in range(n)]
        for i in range(n):
            for j in out_neighbors[i]:
                in_neighbors[j].add(i)

        # Row-stochastic A: node i averages over in-neighbors + self (the "pull" step).
        A = np.zeros((n, n))
        for i in range(n):
            support = sorted(in_neighbors[i] | {i})
            w = 1.0 / len(support)
            for j in support:
                A[i, j] = w

        # Column-stochastic B: node i splits its own weight across out-neighbors + self
        # (the "push" step) -> column i sums to 1.
        B = np.zeros((n, n))
        for i in range(n):
            support = sorted(out_neighbors[i] | {i})
            w = 1.0 / len(support)
            for j in support:
                B[j, i] = w

        assert np.allclose(A.sum(axis=1), 1.0)
        assert np.allclose(B.sum(axis=0), 1.0)
        return A, B


def static_ring_matrices(n: int, extra_edges: int, seed: int = 0):
    """A single fixed (non-time-varying) topology, for the Static_FAB ablation."""
    gen = TimeVaryingDigraphs(n, extra_edges, seed)
    return gen.step()