File size: 6,454 Bytes
dd90a4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
#!/usr/bin/env python3
"""Exact finite-prior Bayes-risk certificate for the participation lower bound.

The construction is deliberately small and discrete.  For each of E independent
local coordinates, a hidden client-level signal theta is either +kappa or
-kappa.  A participating client returns theta plus an independent
Rademacher heterogeneity term, also of magnitude kappa.  The learner sees S
clients and estimates theta.  A uniform prior over theta and exhaustive
enumeration of all N-client noise assignments and all S-client subsets gives
the Bayes risk.  Bayes risk is a lower bound for every estimator under this
finite prior.

All probabilities and risks are computed with Fraction; no floating-point
calculation is used for the certificate.
"""

from __future__ import annotations

import itertools
import json
from collections import defaultdict
from fractions import Fraction
from math import comb
from pathlib import Path


def exact_unit_risk(n_clients: int, sampled: int) -> dict[str, object]:
    """Enumerate the two worlds, all noise assignments, and all sample sets.

    Values are normalized by kappa, so the returned risk is multiplied by
    kappa**2 for an arbitrary positive kappa.  The posterior-mean estimator is
    Bayes optimal for squared loss.  For each observation o with joint masses
    p_plus and p_minus, its contribution is
        4 * p_plus * p_minus / (p_plus + p_minus),
    which is the exact posterior Bayes risk contribution for theta in {-1,+1}.
    """
    if not (1 <= sampled <= n_clients):
        raise ValueError("sampled must be in [1, n_clients]")

    subsets = tuple(itertools.combinations(range(n_clients), sampled))
    subset_count = len(subsets)
    # obs -> [joint mass under theta=+1, joint mass under theta=-1]
    masses: dict[tuple[int, ...], list[Fraction]] = defaultdict(
        lambda: [Fraction(0), Fraction(0)]
    )
    per_world = Fraction(1, 2 * (1 << n_clients) * subset_count)
    for theta_index, theta in enumerate((1, -1)):
        for noise_mask in range(1 << n_clients):
            noise = tuple(1 if (noise_mask >> i) & 1 else -1 for i in range(n_clients))
            values = tuple(theta + z for z in noise)
            for subset in subsets:
                observation = tuple(values[i] for i in subset)
                masses[observation][theta_index] += per_world

    risk = Fraction(0)
    ambiguous_mass = Fraction(0)
    for p_plus, p_minus in masses.values():
        total = p_plus + p_minus
        if p_plus and p_minus:
            ambiguous_mass += total
            risk += Fraction(4) * p_plus * p_minus / total

    expected_formula = Fraction(1, 1 << sampled)
    if risk != expected_formula:
        raise AssertionError((n_clients, sampled, risk, expected_formula))

    return {
        "n_clients": n_clients,
        "sampled": sampled,
        "enumerated_noise_assignments_per_world": 1 << n_clients,
        "enumerated_subsets": subset_count,
        "joint_world_subset_cases": 2 * (1 << n_clients) * subset_count,
        "distinct_observations": len(masses),
        "ambiguous_observation_mass": str(ambiguous_mass),
        "risk_over_kappa_squared": str(risk),
        "risk_over_kappa_squared_decimal": float(risk),
        "closed_form_cross_check": str(expected_formula),
    }


def main() -> None:
    # These N values cover several finite population sizes; every S from one
    # client through full participation is enumerated for each one.
    population_sizes = (5, 8, 10)
    e_values = (1, 2, 4, 8)
    kappa_values = (Fraction(1, 4), Fraction(1, 2), Fraction(1), Fraction(2))

    base_rows: list[dict[str, object]] = []
    for n_clients in population_sizes:
        for sampled in range(1, n_clients + 1):
            base_rows.append(exact_unit_risk(n_clients, sampled))

    # Extend the exact base risks over executed E and kappa regimes.  The E
    # coordinates are independent, so squared risks add exactly; kappa scales
    # the normalized risk by kappa**2.
    cells: list[dict[str, object]] = []
    min_ratio: Fraction | None = None
    for base in base_rows:
        n_clients = int(base["n_clients"])
        sampled = int(base["sampled"])
        unit_risk = Fraction(str(base["risk_over_kappa_squared"]))
        for e_local in e_values:
            for kappa in kappa_values:
                risk = e_local * kappa * kappa * unit_risk
                target_scale = Fraction(e_local) * kappa * kappa / sampled
                ratio = risk / target_scale
                min_ratio = ratio if min_ratio is None else min(min_ratio, ratio)
                cells.append(
                    {
                        "N": n_clients,
                        "S": sampled,
                        "E": e_local,
                        "kappa": str(kappa),
                        "bayes_risk": str(risk),
                        "target_E_kappa2_over_S": str(target_scale),
                        "ratio_to_target": str(ratio),
                    }
                )

    assert min_ratio is not None
    certificate_constant = Fraction(5, 512)
    if min_ratio < certificate_constant:
        raise AssertionError((min_ratio, certificate_constant))

    result = {
        "construction": "two-world Rademacher heterogeneity, exact finite-prior Bayes risk",
        "population_sizes": list(population_sizes),
        "sample_sizes_per_population": {
            str(n): list(range(1, n + 1)) for n in population_sizes
        },
        "E_values": list(e_values),
        "kappa_values": [str(k) for k in kappa_values],
        "base_rows": base_rows,
        "executed_parameter_cells": len(cells),
        "cells": cells,
        "min_ratio_risk_over_E_kappa2_over_S": str(min_ratio),
        "finite_family_certificate": f"risk >= ({certificate_constant}) * E*kappa^2/S",
        "bayes_optimality": (
            "For squared loss, posterior mean minimizes conditional risk; "
            "therefore every estimator has expected risk at least this Bayes risk."
        ),
    }
    out = Path(__file__).with_name("bayes_lower_bound_results.json")
    out.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps({
        "output": str(out),
        "base_rows": len(base_rows),
        "executed_parameter_cells": len(cells),
        "min_ratio": str(min_ratio),
        "certificate": result["finite_family_certificate"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()