Polygenic Score Aggregate

#!/usr/bin/env python3"""local_project_owner.py — LOCAL stages the PROJECT OWNER (researcher) runs.The project owner holds the key and receives the result. These three functionsrun only on the owner's machine — the secret context never leaves it:  * keygen()  — create the BFV crypto context; return (secret, public). The kit    shim 00_keygen.py writes the two halves to disk; only the public half is ever    published.  * decrypt() — the ONLY use of the secret key: aggregate result ciphertext ->    plaintext integer vector.  * decode()  — plaintext vector -> released result (weighted counts + cohort PGS).    The cross-coordinate reduction (Sigma_j) is done HERE, post-decrypt, never    under encryption — which is what keeps the protocol on minimal additive-tier    params.The data owner's stages (encode, encrypt) live in local_data_owner.py; the blindserver stage (compute, incl. the public-weight multiply) lives in server.py.BFV parameters--------------poly_modulus_degree = 8192  -> 8192 packing slots (>> length+1). FIXED across all                               security levels.plain_modulus       = 1073692673 (a 30-bit batching prime, == 1 mod 16384) ->                               exact integer arithmetic in Z_t. The flagship's                               20-bit t = 1032193 is UNDER-SIZED here: the public                               plaintext-weight multiply inflates each coordinate                               to w_scaled[j] * (sum_i g_ij), whose maximum is                               max_j(w_scaled[j]) * 2N. With S = 1000 and effect                               weights <= ~2.0 (w_scaled <= ~2000), t = 1073692673                               stays exact for N up to ~250k. FIXED — a function of                               the value envelope, not security.Security levels (`--security {128,192,256}`)--------------------------------------------`--security` selects the coeff-modulus chain (`coeff_mod_bit_sizes`) and nothingelse. At FIXED N=8192 the security level is the q-band: SMALLER Σbits ⇒ MOREsecure. This protocol's one extra op — a single ciphertext × PLAINTEXT weightmultiply (degree-preserving) — combined with the 30-bit t needs an EFFECTIVE q≳ 80 bits; TenSEAL reserves the LAST coeff prime as a key-switching special prime,so every level ships a 3-prime chain (a 2-prime chain leaves < 80 effective bitsand FAILS to decrypt). All four additive protocols standardize on these samePGS-safe 3-prime chains so the SECURITY table is byte-identical across bundles.  * 128 -> [60,60,60] (Σ=180, eff 120, lands in the 8192 128-band 153–218)  * 192 -> [50,50,50] (Σ=150, eff 100, lands in the 8192 192-band 119–152)  * 256 -> [45,45,28] (Σ=118, eff  90, lands in the 8192 256-band ≤118)"""from __future__ import annotationsDEFAULT_POLY_MODULUS_DEGREE = 8192# 30-bit NTT batching prime (== 1 mod 16384); exact BFV in Z_t, t > max slot value# = max_j(w_scaled[j]) * 2N after the public-weight multiply. See module docstring.DEFAULT_PLAIN_MODULUS = 1073692673DEFAULT_SECURITY = 128# coeff_mod_bit_sizes chain per requested HE security level, at FIXED N=8192.# Verified by real TenSEAL 0.3.16 measurement: each chain decrypts the public-# weighted aggregate bit-exact and its achieved level (Σbits vs the# HomomorphicEncryption.org caps) == requested. PGS-safe 3-prime chains (2-prime# chains leave < 80 effective bits after the special prime and FAIL for this# protocol's 30-bit t + weight multiply). Shared verbatim across all four# additive protocols.SECURITY: dict[int, list[int]] = {    128: [60, 60, 60],  # Σ=180 -> achieved 128    192: [50, 50, 50],  # Σ=150 -> achieved 192    256: [45, 45, 28],  # Σ=118 -> achieved 256}# Published fixed-point factor S (must match server.WEIGHT_SCALE).WEIGHT_SCALE = 1000def keygen(    poly_modulus_degree: int = DEFAULT_POLY_MODULUS_DEGREE,    plain_modulus: int = DEFAULT_PLAIN_MODULUS,    security: int = DEFAULT_SECURITY,) -> tuple[bytes, bytes]:    """Return ``(secret_context_bytes, public_context_bytes)``.    The secret context carries the secret key; the public context is the same    context with the secret key removed (``make_context_public``). Additive-tier    protocol (add + plaintext-scalar multiply) => we generate no relin/Galois keys.    ``security`` selects the ``coeff_mod_bit_sizes`` chain from ``SECURITY`` — the    only parameter that varies with the requested HE security level.    """    import tenseal as ts    if security not in SECURITY:        raise ValueError(            f"unsupported security level {security!r}; choose one of {sorted(SECURITY)}"        )    context = ts.context(        ts.SCHEME_TYPE.BFV,        poly_modulus_degree=poly_modulus_degree,        plain_modulus=plain_modulus,        coeff_mod_bit_sizes=SECURITY[security],    )    # Serialize the private half (with secret key) first.    secret_bytes = context.serialize(save_secret_key=True)    # Derive the public half from an independent copy so we never mutate the    # secret context in place.    public_context = ts.context_from(secret_bytes)    public_context.make_context_public()    public_bytes = public_context.serialize()    return secret_bytes, public_bytesdef decrypt(secret_context_bytes: bytes, result_bytes: bytes) -> list[int]:    """Decrypt the aggregate ciphertext -> plaintext integer vector (length L+1).    This is the ONLY point the secret key is used, and it runs on the owner's    machine — never on the server. The decrypted tensor has length ``L + 1``: the    first L slots are the per-coordinate weighted counts ``w_scaled[j] *    (sum_i g_ij)`` (integer, fixed-point domain of scale S), the trailing slot is    the append-1 sentinel — weighted by 1 in stage 30, so it still decrypts to    exactly N.    """    import tenseal as ts    context = ts.context_from(secret_context_bytes)    if not context.is_private():        raise ValueError("decrypt stage needs the secret context (has no secret key)")    return [int(value) for value in ts.bfv_vector_from(context, result_bytes).decrypt()]def decode(plain: list[int], length: int, scale: int = WEIGHT_SCALE) -> dict:    """Split sentinel from weighted counts and compute the cohort PGS.    Raises ValueError if the tensor is not exactly ``length + 1`` slots, which    would mean the sentinel is missing or the coordinate length disagrees.        cohort_pgs_scaled = sum_j weighted_counts[j]        # integer, exact        cohort_pgs_sum    = cohort_pgs_scaled / S           # sum_i PGS_i  (real)        mean_pgs          = cohort_pgs_scaled / (S * N)     # mean per-contributor PGS    """    expected = length + 1    if len(plain) != expected:        raise ValueError(            f"expected {expected} slots (L={length} + 1 sentinel), got {len(plain)}"        )    if scale <= 0:        raise ValueError(f"scale must be positive, got {scale}")    weighted_counts = [int(value) for value in plain[:length]]    n_contributors = int(plain[length])    if n_contributors <= 0:        raise ValueError(f"sentinel decoded to N={n_contributors}; expected N > 0")    cohort_pgs_scaled = sum(weighted_counts)    cohort_pgs_sum = cohort_pgs_scaled / scale    mean_pgs = cohort_pgs_scaled / (scale * n_contributors)    return {        "protocol": "polygenic_score_aggregate",        "coordinates_length": length,        "weight_scale": scale,        "n_contributors": n_contributors,        "weighted_counts": weighted_counts,        "cohort_pgs_scaled": cohort_pgs_scaled,        "cohort_pgs_sum": cohort_pgs_sum,        "mean_pgs": mean_pgs,    }