#!/usr/bin/env python3 """Reference analysis — privacy-preserving allele-frequency aggregation via additive homomorphic encryption (Paillier). Question shape: "Can a fast homomorphic-encryption scheme compute a pooled genomic statistic (allele frequency) across cohorts WITHOUT any party revealing its per-cohort counts — and at what cost?" Textbook use: secure/federated genomic analytics; the privacy-preserving- computation family (homomorphic encryption, secure aggregation, SMPC) applied to population genetics. Design: the five 1000 Genomes super-populations (AFR/AMR/EAS/EUR/SAS) stand in as five mutually distrusting data holders. For the most differentiated variant in the fetched slice, each party holds its private (alt allele count, total allele count). Each encrypts BOTH counts under an additive Paillier scheme; an untrusted aggregator homomorphically SUMS the ciphertexts (E(a)*E(b) = E(a+b) mod n^2) without ever decrypting a per-party value, then only the pooled ciphertexts are decrypted to yield the global allele frequency. What is measured (all REAL, all from public zero-auth 1000G data): * correctness — the decrypted pooled counts must EXACTLY equal the plaintext pooled counts (max_abs_error == 0). This is the scientific claim: additive HE is lossless for this aggregation. It is VERIFIED, not assumed — an undersized modulus would wrap and is honestly reported as a null_result. * overhead — ciphertext expansion (bytes) and keygen/encrypt/add/decrypt latency at the chosen security level. Input: --kg CSV from fetch_1000g.py (chrom,pos,id,ref,alt,af,eas_af,amr_af,afr_af,eur_af,sas_af) Output: figure.png — per-cohort private alt-allele contributions + the verified pooled allele frequency and measured HE overhead stats.json — pooled_allele_freq, plaintext_allele_freq, max_abs_error, exact_match, ciphertext_expansion, latencies, key_bits Deterministic: --seed seeds prime generation and blinding, so the whole run (including the key) reproduces bit-for-bit; the pooled statistic is a pure function of the fetched slice regardless of the key, so reproduction matches. Usage: python3 he_secure_aggregation.py --kg work/kg_apoe.csv \ --label "APOE region" --seed 1234 --key-bits 1024 --outdir work Exit: 0 computed, 1 could not compute, 2 usage. """ from __future__ import annotations import argparse import json import random import sys import time from functools import reduce import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import pandas as pd # 1000G phase-3 super-population sample sizes (2N haplotypes) — fixed public # constants of the panel, used to turn a published AF back into the integer # allele counts each "party" would hold. Identical constants to # population_allele_freq so the two templates agree on the panel. SUPERPOP_2N = {"afr": 1322, "amr": 694, "eas": 1008, "eur": 1006, "sas": 978} POPS = ["afr", "amr", "eas", "eur", "sas"] # --------------------------------------------------------------------------- # # Minimal, dependency-free Paillier cryptosystem (additively homomorphic). # Pure Python big-int modular arithmetic — no crypto wheels needed. # --------------------------------------------------------------------------- # def _miller_rabin(n: int, rng: random.Random, rounds: int = 40) -> bool: if n < 2: return False for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): if n % p == 0: return n == p d, r = n - 1, 0 while d % 2 == 0: d //= 2 r += 1 for _ in range(rounds): a = rng.randrange(2, n - 1) x = pow(a, d, n) if x in (1, n - 1): continue for _ in range(r - 1): x = pow(x, 2, n) if x == n - 1: break else: return False return True def _gen_prime(bits: int, rng: random.Random) -> int: while True: candidate = rng.getrandbits(bits) | (1 << (bits - 1)) | 1 # top+bottom bit set if _miller_rabin(candidate, rng): return candidate class Paillier: def __init__(self, key_bits: int, rng: random.Random): half = key_bits // 2 p = _gen_prime(half, rng) q = _gen_prime(half, rng) while q == p: q = _gen_prime(half, rng) self.n = p * q self.n2 = self.n * self.n # g = n + 1 lets pow(g, m, n^2) collapse to (1 + m*n) mod n^2. self.lam = (p - 1) * (q - 1) // _gcd(p - 1, q - 1) self.mu = pow(self.lam, -1, self.n) self._rng = rng def encrypt(self, m: int) -> int: # Blinding factor r coprime to n (any r in Z*_n; primes p,q >> n ensures it). r = self._rng.randrange(1, self.n) return ((1 + m * self.n) * pow(r, self.n, self.n2)) % self.n2 def add(self, c1: int, c2: int) -> int: return (c1 * c2) % self.n2 def decrypt(self, c: int) -> int: u = pow(c, self.lam, self.n2) return (((u - 1) // self.n) * self.mu) % self.n def _gcd(a: int, b: int) -> int: while b: a, b = b, a % b return a def _ms(fn): start = time.perf_counter() result = fn() return result, (time.perf_counter() - start) * 1000.0 def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--kg", required=True) ap.add_argument("--label", default="region") ap.add_argument("--seed", type=int, default=1234) ap.add_argument("--key-bits", type=int, default=1024, help="Paillier modulus size; 1024 for a fast demo, 2048+ in production") ap.add_argument("--outdir", default=".") args = ap.parse_args() if args.key_bits < 256: _fail(args.outdir, args.label, "key-bits must be at least 256 for a meaningful scheme") return 2 df = pd.read_csv(args.kg, dtype=str).fillna("") for p in POPS: df[f"{p}_af"] = pd.to_numeric(df.get(f"{p}_af"), errors="coerce") df = df.dropna(subset=[f"{p}_af" for p in POPS]) if df.empty: _fail(args.outdir, args.label, "no variants with all-population AF in the slice") return 1 # Lead variant = largest spread across populations (the most interesting one # to pool privately); same selection rule as population_allele_freq. afmat = df[[f"{p}_af" for p in POPS]].to_numpy(dtype=float) spread = afmat.max(axis=1) - afmat.min(axis=1) lead = df.iloc[int(np.argmax(spread))] lead_af = {p: float(lead[f"{p}_af"]) for p in POPS} # Each party's PRIVATE integer counts for the lead variant. alt_counts = {p: round(lead_af[p] * SUPERPOP_2N[p]) for p in POPS} tot_counts = {p: SUPERPOP_2N[p] for p in POPS} plain_alt = sum(alt_counts.values()) plain_tot = sum(tot_counts.values()) if plain_tot == 0: _fail(args.outdir, args.label, "degenerate zero-total allele panel for the lead variant") return 1 rng = random.Random(args.seed) # ---- keygen -> encrypt -> homomorphic aggregate -> decrypt, all timed ---- scheme, keygen_ms = _ms(lambda: Paillier(args.key_bits, rng)) def _encrypt_all(): return ( [scheme.encrypt(alt_counts[p]) for p in POPS], [scheme.encrypt(tot_counts[p]) for p in POPS], ) (enc_alt, enc_tot), encrypt_ms = _ms(_encrypt_all) (agg_alt, agg_tot), add_ms = _ms(lambda: ( reduce(scheme.add, enc_alt), reduce(scheme.add, enc_tot), )) (dec_alt, dec_tot), decrypt_ms = _ms(lambda: ( scheme.decrypt(agg_alt), scheme.decrypt(agg_tot), )) # ---- correctness: the HE aggregate must EXACTLY equal the plaintext one ---- exact_counts = (dec_alt == plain_alt) and (dec_tot == plain_tot) plain_af = plain_alt / plain_tot he_af = (dec_alt / dec_tot) if dec_tot else float("nan") max_abs_error = 0.0 if exact_counts else abs(he_af - plain_af) exact_match = bool(exact_counts and max_abs_error == 0.0) # ---- overhead ---- ciphertext_bytes = (scheme.n2.bit_length() + 7) // 8 plaintext_bytes = max(1, (max(plain_alt, plain_tot).bit_length() + 7) // 8) ciphertext_expansion = round(ciphertext_bytes / plaintext_bytes, 1) outcome = "success" if exact_match else "null_result" result = { "label": args.label, "seed": args.seed, "reference_template": "he_secure_aggregation", "scheme": "paillier_additive_homomorphic", "key_bits": int(args.key_bits), "n_parties": len(POPS), "n_variants_in_slice": int(len(df)), "lead_variant": lead.get("id") or f"{lead['chrom']}:{lead['pos']}", "private_alt_counts_by_pop": alt_counts, "pooled_alt_allele_count_he": int(dec_alt), "pooled_alt_allele_count_plaintext": int(plain_alt), "pooled_total_alleles": int(plain_tot), "pooled_allele_freq": round(he_af, 6), "plaintext_allele_freq": round(plain_af, 6), "max_abs_error": max_abs_error, "exact_match": exact_match, "outcome": outcome, "ciphertext_bytes": int(ciphertext_bytes), "plaintext_bytes": int(plaintext_bytes), "ciphertext_expansion": ciphertext_expansion, "keygen_ms": round(keygen_ms, 2), "encrypt_ms": round(encrypt_ms, 2), "homomorphic_add_ms": round(add_ms, 2), "decrypt_ms": round(decrypt_ms, 2), "headline_statistic": ( f"pooled alt AF = {he_af:.4f} over {len(POPS)} private cohorts, " f"exact match to plaintext (err {max_abs_error:g}); " f"{ciphertext_expansion:g}x ciphertext, " f"{(encrypt_ms + add_ms + decrypt_ms):.0f} ms secure aggregate at {args.key_bits}-bit" ), } with open(f"{args.outdir}/stats.json", "w") as fh: json.dump(result, fh, indent=2) _render(args.outdir, alt_counts, result) print(json.dumps(result, indent=2)) return 0 def _fail(outdir, label, reason): result = {"label": label, "reference_template": "he_secure_aggregation", "outcome": "failed", "reason": reason} with open(f"{outdir}/stats.json", "w") as fh: json.dump(result, fh, indent=2) print(json.dumps(result, indent=2)) def _render(outdir, alt_counts, result): fig, ax = plt.subplots(figsize=(6.4, 4.6)) colors = ["#8e44ad", "#e67e22", "#16a085", "#2980b9", "#c0392b"] ax.bar([p.upper() for p in POPS], [alt_counts[p] for p in POPS], color=colors, alpha=0.85) ax.set_ylabel("Private alt-allele count (never revealed)") ax.set_title( f"{result['label']}: {result['lead_variant']} — privacy-preserving " f"pooled allele frequency\n{result['headline_statistic']}", fontsize=9) ax.grid(axis="y", alpha=0.3) verdict = "EXACT" if result["exact_match"] else f"error {result['max_abs_error']:g}" ax.text(0.98, 0.95, f"Paillier {result['key_bits']}-bit\n" f"HE pooled AF = {result['pooled_allele_freq']:.4f}\n" f"plaintext AF = {result['plaintext_allele_freq']:.4f} ({verdict})\n" f"{result['ciphertext_expansion']:g}x ciphertext expansion", transform=ax.transAxes, ha="right", va="top", fontsize=8, bbox=dict(boxstyle="round", facecolor="#f4f4f4", edgecolor="#999")) fig.tight_layout() fig.savefig(f"{outdir}/figure.png", dpi=130) plt.close(fig) if __name__ == "__main__": sys.exit(main())