1"""Accounting of rounds, messages and transmitted numbers/bytes for the protocols (star topology)."""2from __future__ import annotations345def protocol_accounting(K: int, p: int, r: int = 1, q: int = 1, hessian_mode: str = "exact", bytes_per_scalar: int = 8,6 send_local_hessian: bool = False, header_bytes: int = 0):7 """Numbers and bytes moved for: local pilot (phase 1) + q derivative phases.89 Phase 1 uplink per worker: p (corrected local estimate) [+ p(p+1)/2 if the local Hessian is shipped: stale/AEE].10 Each derivative phase: downlink broadcast of p to every worker, uplink of p (score) [+ p(p+1)/2 Hessian if exact].11 hessian_mode: 'exact' | 'stale' | 'single' (no Hessian uplink; machine 0 solves) | 'none' (q = 0).12 Returns dict with rounds, messages, uplink/downlink numbers and bytes, local fits per worker.13 """14 h = p * (p + 1) // 215 up1 = p + (h if (send_local_hessian or hessian_mode == "stale") else 0)16 if q == 0:17 up_d, down_d = 0, 018 elif hessian_mode == "exact" or hessian_mode == "quantized":19 up_d, down_d = q * (p + h), q * p20 elif hessian_mode in ("stale", "single"):21 up_d, down_d = q * p, q * p22 else:23 raise ValueError(hessian_mode)24 uplink_numbers = K * (up1 + up_d)25 downlink_numbers = K * down_d26 rounds = 1 + q # synchronized exchange phases (uplink of estimates, then q derivative phases)27 messages = K * (1 + 2 * q) # K uploads + per phase: K broadcasts + K uploads28 local_fits = 2 ** r - 1 # splits (1,2,4,...): 1 + 2 + 4 + ...29 local_data_passes = r # each level touches every local observation once30 return {"rounds": rounds, "messages": messages, "uplink_numbers": uplink_numbers, "downlink_numbers": downlink_numbers,31 "uplink_bytes": uplink_numbers * bytes_per_scalar + messages * header_bytes,32 "downlink_bytes": downlink_numbers * bytes_per_scalar,33 "local_fits_per_worker": local_fits, "local_data_passes": local_data_passes}343536def sufficient_exponent(r: int, q: int) -> float:37 """alpha(r,q) = 1 - 1/(2^(q+1) r): K = o(N^alpha) is the sufficient regime of Theorem 2 (exact Newton)."""38 return 1.0 - 1.0 / (2 ** (q + 1) * r)394041def sufficient_exponent_stale(r: int, q: int) -> float:42 """Sufficient exponent with a fixed (stale) averaged local Hessian: effective depth r + q (Theorem 4, case (b))."""43 if q == 0:44 return sufficient_exponent(r, 0)45 return 1.0 - 1.0 / (2.0 * (r + q))464748def sufficient_exponent_single(r: int, q: int) -> float:49 """Sufficient exponent with a single machine's Hessian (CSL-type curvature): effective depth r + q/2."""50 if q == 0:51 return sufficient_exponent(r, 0)52 return 1.0 - 1.0 / (2.0 * r + q)53