1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6SEED = 17
  7TREE = {0: [1, 2], 1: [3], 2: []}
  8PARENT = {1: 0, 2: 0, 3: 1}
  9N = 4
 10
 11
 12def correction_terms(u, s, rho):
 13    """One dual update and the exact parent corrections -sum(delta s)."""
 14    s_new = s.copy()
 15    for child, parent in PARENT.items():
 16        s_new[child] += rho * (u[child] - u[parent])
 17    corr = np.zeros(N)
 18    for parent, children in TREE.items():
 19        corr[parent] = -sum(s_new[ch] - s[ch] for ch in children)
 20    return s_new, corr
 21
 22
 23def tree_step(u, s, targets, rho, eta, penalty, corrected=True):
 24    # Jacobi primal step. Each node reads old neighbors; the root/parents can
 25    # additionally read child dual increments through the correction channel.
 26    s_new, corr = correction_terms(u, s, rho)
 27    grad = u - targets
 28    for child, parent in PARENT.items():
 29        d = u[child] - u[parent]
 30        grad[parent] -= s[child] + penalty * d
 31        grad[child] += s[child] + penalty * d
 32    if corrected:
 33        # At a parent, -delta s is precisely the paper's correction term.
 34        grad += corr
 35    unew = u - eta * grad
 36    return unew, s_new, corr
 37
 38
 39def corrected_tree(targets, rounds=160, rho=.20, eta=.16, penalty=.8, corrected=True):
 40    u = np.zeros(N)
 41    s = np.zeros(N)
 42    losses, residuals, bytes_sent = [], [], 0
 43    for _ in range(rounds):
 44        u, s, _ = tree_step(u, s, targets, rho, eta, penalty, corrected)
 45        loss = .5 * np.sum((u-targets)**2) + .5 * penalty * sum((u[ch]-u[p])**2 for ch,p in PARENT.items())
 46        losses.append(float(loss))
 47        residuals.append(float(np.sqrt(np.mean([(u[ch]-u[p])**2 for ch,p in PARENT.items()]))))
 48        # Each edge sends a scalar state and a scalar correction in this toy.
 49        bytes_sent += 2 * len(PARENT) * 8
 50    return u, losses, residuals, bytes_sent
 51
 52
 53def fedavg(targets, rounds=160, eta=.16, local_steps=1):
 54    u = np.zeros(N)
 55    losses, residuals = [], []
 56    for _ in range(rounds):
 57        proposals = []
 58        for x in u:
 59            z = x
 60            for _ in range(local_steps):
 61                z -= eta * (z - targets[len(proposals)])
 62            proposals.append(z)
 63        mean = np.mean(proposals)
 64        u[:] = mean
 65        loss = .5 * np.sum((u-targets)**2) + .5*.8*sum((u[ch]-u[p])**2 for ch,p in PARENT.items())
 66        losses.append(float(loss))
 67        residuals.append(float(np.sqrt(np.mean([(u[ch]-u[p])**2 for ch,p in PARENT.items()]))))
 68    # four workers upload and root broadcasts a scalar each round
 69    return u, losses, residuals, rounds * 2 * N * 8
 70
 71
 72def math_checks():
 73    # Prediction 1: correction conservation error is numerical zero for every rho.
 74    u = np.array([.2, 1.3, -.7, .4])
 75    s = np.array([0., .8, -.2, .5])
 76    conservation = []
 77    for rho in [.01, .1, .4, .9]:
 78        sn, corr = correction_terms(u, s, rho)
 79        expected = np.array([-(sn[1]-s[1])-(sn[2]-s[2]),
 80                             -(sn[3]-s[3]), 0., 0.])
 81        conservation.append(float(abs(corr[0]-expected[0]) + abs(corr[1]-expected[1])))
 82
 83    # Prediction 2: fixed-state correction norm is exactly linear in rho.
 84    norms = []
 85    for rho in [.05, .10, .20, .40]:
 86        _, c = correction_terms(u, s, rho)
 87        norms.append(float(np.linalg.norm(c)))
 88    ratios = [norms[i]/norms[0] for i in range(4)]
 89    predicted_ratios = [1., 2., 4., 8.]
 90
 91    # Prediction 3: stale parent dual error is the correction magnitude and
 92    # therefore has zero intercept and linear rho scaling.
 93    stale = norms
 94    slope = float(np.dot([.05,.1,.2,.4], stale) / np.dot([.05,.1,.2,.4], [.05,.1,.2,.4]))
 95    intercept = float(np.mean(stale) - slope*np.mean([.05,.1,.2,.4]))
 96    return {
 97        "conservation_abs_errors": conservation,
 98        "correction_norms": norms,
 99        "rho_ratio_observed": ratios,
100        "rho_ratio_predicted": predicted_ratios,
101        "stale_error_linear_fit_slope": slope,
102        "stale_error_linear_fit_intercept": intercept,
103        "max_conservation_error": max(conservation),
104    }
105
106
107def main():
108    np.random.seed(SEED)
109    targets = np.array([-2.0, 1.0, 2.5, -1.0])
110    checks = math_checks()
111    sweep = []
112    for rho in [.0, .05, .10, .20, .40]:
113        uc, lc, rc, bc = corrected_tree(targets, rho=rho, corrected=True)
114        us, ls, rs, bs = corrected_tree(targets, rho=rho, corrected=False)
115        sweep.append({
116            "rho": rho,
117            "corrected_final_loss": lc[-1], "stale_final_loss": ls[-1],
118            "corrected_final_residual": rc[-1], "stale_final_residual": rs[-1],
119            "corrected_loss_40": lc[39], "stale_loss_40": ls[39],
120        })
121    uc, lc, rc, bc = corrected_tree(targets, rho=.20, corrected=True)
122    uf, lf, rf, bf = fedavg(targets)
123    result = {
124        "seed": SEED, "targets": targets.tolist(), "math_checks": checks,
125        "rho_sweep": sweep,
126        "mini_experiment": {
127            "tree_correction": {"final_loss": lc[-1], "final_residual": rc[-1], "bytes": bc},
128            "fedavg": {"final_loss": lf[-1], "final_residual": rf[-1], "bytes": bf},
129            "stale_tree_rho_0.2": {"final_loss": corrected_tree(targets, rho=.2, corrected=False)[1][-1],
130                                    "final_residual": corrected_tree(targets, rho=.2, corrected=False)[2][-1],
131                                    "bytes": bc},
132        },
133        "interpretation": "Checks test algebraic correction conservation and rho scaling; optimization comparison is a small scalar quadratic proxy, not a neural-network benchmark."
134    }
135    Path("results.json").write_text(json.dumps(result, indent=2))
136    print(json.dumps(result, indent=2))
137
138if __name__ == "__main__":
139    main()