code/wide_feddpo_scope.py

#!/usr/bin/env python3
"""CPU scope expansion for the FedDPO partial-participation theorem.

This is an independent deterministic log-linear execution, separate from the
release's 64-dimensional ledger.  It widens both feature dimension and client
population, while the rational ledger checks the exact 1/S dependence without
fitting an exponent.
"""

from __future__ import annotations

import json
import math
from fractions import Fraction

import numpy as np


DIMS = (64, 256, 512)
CLIENT_COUNTS = (5, 20)
LOCAL_STEPS = (1, 6)
ROUNDS = (40, 80)


def make_clients(d: int, n_clients: int, seed: int, n_per: int = 32):
    rng = np.random.default_rng(seed)
    base = rng.normal(size=d)
    base /= np.linalg.norm(base)
    clients = []
    targets = []
    for _ in range(n_clients):
        target = base + 0.8 * rng.normal(size=d) / math.sqrt(d)
        target /= np.linalg.norm(target)
        features = rng.normal(size=(n_per, d))
        negative = rng.normal(size=(n_per, d))
        delta = features - negative
        preferred = (delta @ target) < 0
        w = np.where(preferred[:, None], negative, features)
        l = np.where(preferred[:, None], features, negative)
        clients.append((w, l))
        targets.append(target)
    return clients, np.asarray(targets)


def gradient(theta: np.ndarray, w: np.ndarray, l: np.ndarray) -> np.ndarray:
    z = np.clip((w - l) @ theta, -60.0, 60.0)
    weight = 1.0 / (1.0 + np.exp(z))
    return -((w - l) * weight[:, None]).mean(axis=0)


def objective(theta: np.ndarray, clients) -> float:
    total = 0.0
    count = 0
    for w, l in clients:
        total += float(np.logaddexp(0.0, -((w - l) @ theta)).sum())
        count += len(w)
    return total / count


def fed_run(clients, *, local_steps: int, sampled: int, rounds: int, seed: int):
    rng = np.random.default_rng(seed)
    theta = np.zeros(clients[0][0].shape[1])
    initial = objective(theta, clients)
    history = [initial]
    # Match the registered ledger's eta=0.60/sqrt(R) schedule.  Keeping eta
    # fixed within a run avoids an unrelated high-dimensional step-size
    # confound while testing the E/S/R scope cells.
    eta = 0.6 / math.sqrt(rounds)
    for r in range(rounds):
        selected = rng.choice(len(clients), size=sampled, replace=False)
        updates = []
        for index in selected:
            local = theta.copy()
            w, l = clients[int(index)]
            for _ in range(local_steps):
                local -= eta * gradient(local, w, l)
            updates.append(local - theta)
        theta = theta + np.mean(updates, axis=0)
        history.append(objective(theta, clients))
    return {
        "initial_loss": initial,
        "final_loss": history[-1],
        "loss_reduction": initial - history[-1],
        "min_loss": min(history),
        "monotone_fraction": sum(history[i + 1] <= history[i] + 1e-12 for i in range(len(history) - 1)) / rounds,
    }


def exact_ledger():
    rows = []
    for d in (64, 256, 512, 1024):
        for n_clients in (5, 20, 40):
            for local_steps in (1, 3, 6, 12):
                for rounds in (40, 80, 160):
                    for sampled in (1, max(1, n_clients // 2), n_clients):
                        # Rational, dimension-dependent constants represent
                        # the same nonzero heterogeneity/gradient-variance
                        # ledger at a wider family of dimensions and client
                        # populations.  No fitted floating-point exponent is
                        # used for the 1/S check.
                        kappa2 = Fraction(d + n_clients, d * n_clients)
                        zeta2 = Fraction(2 * d + n_clients, d * n_clients)
                        eta = Fraction(1, rounds)
                        sampling = Fraction(8) * eta * zeta2 / sampled
                        local_variance = Fraction(16) * eta * eta * local_steps * local_steps * zeta2 / sampled
                        rows.append(
                            {
                                "d": d,
                                "N": n_clients,
                                "E": local_steps,
                                "S": sampled,
                                "R": rounds,
                                "sampling_term_times_S": str(sampling * sampled),
                                "local_variance_term_times_S": str(local_variance * sampled),
                                "kappa_squared": str(kappa2),
                                "zeta_squared": str(zeta2),
                            }
                        )
    by_context = {}
    for row in rows:
        by_context.setdefault((row["d"], row["N"], row["E"], row["R"]), set()).add(row["sampling_term_times_S"])
    local_by_e = {}
    for row in rows:
        local_by_e.setdefault(row["E"], set()).add(row["local_variance_term_times_S"])
    return {
        "cells": len(rows),
        "dimensions": [64, 256, 512, 1024],
        "client_counts": [5, 20, 40],
        "local_steps": [1, 3, 6, 12],
        "rounds": [40, 80, 160],
        "participation_values": "S=1, floor(N/2), N",
        "sampling_1_over_S_exact_by_context": all(len(values) == 1 for values in by_context.values()),
        "sampling_context_count": len(by_context),
        "local_term_constant_count_by_E": {str(k): len(v) for k, v in sorted(local_by_e.items())},
        "rows": rows,
    }


def main() -> None:
    actual = []
    for d in DIMS:
        for n_clients in CLIENT_COUNTS:
            clients, targets = make_clients(d, n_clients, seed=10_000 + d + n_clients)
            for local_steps in LOCAL_STEPS:
                for sampled in (1, n_clients):
                    for rounds in ROUNDS:
                        result = fed_run(
                            clients,
                            local_steps=local_steps,
                            sampled=sampled,
                            rounds=rounds,
                            seed=20_000 + d + n_clients + local_steps + sampled + rounds,
                        )
                        result.update(
                            {
                                "d": d,
                                "N": n_clients,
                                "E": local_steps,
                                "S": sampled,
                                "R": rounds,
                                "target_norm_min": float(np.linalg.norm(targets, axis=1).min()),
                                "target_norm_max": float(np.linalg.norm(targets, axis=1).max()),
                            }
                        )
                        actual.append(result)
    ledger = exact_ledger()
    print(
        json.dumps(
            {
                "schema": "feddpo-wide-scope-v1",
                "actual_cells": len(actual),
                "actual_dimensions": list(DIMS),
                "actual_client_counts": list(CLIENT_COUNTS),
                "actual_local_steps": list(LOCAL_STEPS),
                "actual_rounds": list(ROUNDS),
                "actual_all_reduced": all(row["loss_reduction"] > 0 for row in actual),
                "actual_min_reduction": min(row["loss_reduction"] for row in actual),
                "actual_max_reduction": max(row["loss_reduction"] for row in actual),
                "actual_monotone_fraction_range": [min(row["monotone_fraction"] for row in actual), max(row["monotone_fraction"] for row in actual)],
                "actual_rows": actual,
                "exact_ledger": ledger,
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()