1import json, math
  2from pathlib import Path
  3import numpy as np
  4from scipy.integrate import quad
  5from scipy.special import gammaln
  6
  7EPS = 1.0
  8DELTA_PRIV = 1e-5
  9
 10
 11def log_norm(p, b):
 12    return math.log(p) - math.log(2.0) - math.log(b) - gammaln(1.0 / p)
 13
 14
 15def log_density(x, p, b):
 16    return log_norm(p, b) - (abs(x) / b) ** p
 17
 18
 19def hockey_divergence(p, b, sensitivity, epsilon=EPS):
 20    # Directly integrates [f(x)-exp(eps) f(x-Delta)]+.  Infinite limits
 21    # are important for p>1, whose likelihood ratio has an unbounded tail.
 22    d = float(sensitivity)
 23    ln = log_norm(p, b)
 24    ee = float(epsilon)
 25
 26    def integrand(x):
 27        a = -(abs(x) / b) ** p
 28        c = ee - (abs(x - d) / b) ** p
 29        # common normalization factored out; avoids overflow in exp(eps)
 30        z = math.exp(ln + a) - math.exp(ln + c) if max(ln+a, ln+c) < 700 else 0.0
 31        return max(0.0, z)
 32
 33    val, err = quad(integrand, -np.inf, np.inf, epsabs=2e-11,
 34                    epsrel=2e-8, limit=250, points=None)
 35    return float(val)
 36
 37
 38def calibrate_b(p, sensitivity=1.0, epsilon=EPS, delta=DELTA_PRIV):
 39    # Divergence decreases monotonically with b. Find a conservative bracket.
 40    lo, hi = 1e-8 * sensitivity, 1.0 * sensitivity
 41    while hockey_divergence(p, hi, sensitivity, epsilon) > delta:
 42        hi *= 2.0
 43        if hi > 1e7 * max(1.0, sensitivity):
 44            raise RuntimeError("failed to bracket privacy scale")
 45    for _ in range(55):
 46        mid = (lo + hi) / 2.0
 47        if hockey_divergence(p, mid, sensitivity, epsilon) <= delta:
 48            hi = mid
 49        else:
 50            lo = mid
 51    return hi
 52
 53
 54def moment(p, b, m=2):
 55    return math.exp(m * math.log(b) + gammaln((m + 1.0) / p) - gammaln(1.0 / p))
 56
 57
 58def select_shape(sensitivity=1.0, epsilon=EPS, delta=DELTA_PRIV,
 59                 grid=None, m=2):
 60    if grid is None:
 61        grid = [1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 4.0, 6.0, 8.0]
 62    rows = []
 63    for p in grid:
 64        b = calibrate_b(p, sensitivity, epsilon, delta)
 65        rows.append({"p": p, "b": b, "divergence": hockey_divergence(p, b, sensitivity, epsilon),
 66                     "moment": moment(p, b, m)})
 67    best = min(rows, key=lambda r: r["moment"])
 68    return best, rows
 69
 70
 71def sample_noise(rng, p, b, size):
 72    # If T=(|Z|/b)^p, then T ~ Gamma(1/p, 1).
 73    # The exponential inverse-CDF formula is the special case p=1 only.
 74    radius = rng.gamma(shape=1.0 / p, scale=1.0, size=size) ** (1.0 / p)
 75    return b * rng.choice(np.array([-1.0, 1.0]), size=size) * radius
 76
 77
 78def run_mean_experiment(p_configs, seed=123, n_trials=300, steps=80,
 79                        n=64, lr=0.25, clip=1.0, true_mu=1.0):
 80    # Scalar clipped per-example mean-gradient query, with sensitivity 2C/n.
 81    # Each trial has a fixed dataset; noise is freshly sampled per update.
 82    out = {name: [] for name in p_configs}
 83    rng = np.random.default_rng(seed)
 84    for _ in range(n_trials):
 85        data = rng.normal(true_mu, 0.7, n)
 86        for name, cfg in p_configs.items():
 87            theta = 0.0
 88            p, b = cfg
 89            for _step in range(steps):
 90                per = np.clip(theta - data, -clip, clip)
 91                g = float(np.mean(per))
 92                # b was calibrated for scalar sum-query sensitivity convention;
 93                # this experiment uses the same specified sensitivity in all arms.
 94                theta -= lr * (g + float(sample_noise(rng, p, b, 1)[0]))
 95            out[name].append((theta - true_mu) ** 2)
 96    return {k: {"mse": float(np.mean(v)), "se": float(np.std(v, ddof=1)/math.sqrt(len(v)))}
 97            for k, v in out.items()}
 98
 99
100def main():
101    grid = [1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 4.0, 6.0, 8.0]
102    best, rows = select_shape(grid=grid)
103    # Scaling check: normalized geometry should preserve p and scale b.
104    scaling = []
105    for factor in (0.5, 1.0, 2.0):
106        b, rs = select_shape(sensitivity=factor, grid=grid)
107        scaling.append({"factor": factor, "best_p": b["p"], "best_b": b["b"],
108                        "b_over_factor": b["b"] / factor})
109    lap = next(r for r in rows if r["p"] == 1.0)
110    gau = next(r for r in rows if r["p"] == 2.0)
111    configs = {"laplace_p1": (1.0, lap["b"]), "gaussian_p2": (2.0, gau["b"]),
112               "optimized": (best["p"], best["b"])}
113    task = run_mean_experiment(configs)
114    result = {"settings": {"epsilon": EPS, "delta": DELTA_PRIV, "sensitivity": 1.0,
115                            "grid": grid}, "selected": best, "rows": rows,
116              "scaling": scaling, "task": task}
117    Path("results.json").write_text(json.dumps(result, indent=2))
118    print(json.dumps(result, indent=2))
119
120if __name__ == "__main__":
121    main()