Allele Frequency with Variance

#!/usr/bin/env python3"""local_data_owner.py — LOCAL stages each DATA OWNER (contributor) runs.A data owner turns their own raw genotype data into one ciphertext, on their ownmachine, using the project's PUBLIC context only. Nothing but the ciphertext (andthe already-public context) ever leaves the machine; the secret key is nevertouched here.  * encode()  — raw dosage vector -> validated, zero-padded, length-L vector.  * encrypt() — encoded vector (+ append-1 sentinel) -> serialized ciphertext.The project owner's stages (keygen, decrypt, decode) live inlocal_project_owner.py; the blind server stage (compute, incl. the ct x ctsquare) lives in server.py.`allele_frequency_with_variance` shares the flagship's input schema VERBATIM: theraw input is an alt-allele dosage vector ``g in {0,1,2}^L``. The client sends ONLY``g`` (a SINGLE ciphertext); the server derives ``sum_g2 = sum_i enc(g_ij)^2``under encryption. That is the whole point of the multiplicative version — thesecond moment is server-derived, so the contributor payload stays identical tothe flagship."""from __future__ import annotationsVALUE_DOMAIN = (0, 1, 2)MISSING_ENCODED_AS = 0SENTINEL_VALUE = 1def encode(raw: list[int | None], length: int) -> list[int]:    """Return a validated, zero-padded, length-``length`` dosage vector.    Missing calls (``None``) encode as 0; every present call is validated in    {0,1,2}; the vector is zero-padded to length L (rejected if longer).    Raises ValueError on an out-of-domain call or an over-length input.    """    if length <= 0:        raise ValueError(f"length must be positive, got {length}")    if len(raw) > length:        raise ValueError(            f"raw vector length {len(raw)} exceeds coordinate length {length}"        )    encoded: list[int] = []    for index, call in enumerate(raw):        if call is None:            encoded.append(MISSING_ENCODED_AS)            continue        if call not in VALUE_DOMAIN:            raise ValueError(                f"coordinate {index}: dosage {call!r} not in {VALUE_DOMAIN}"            )        encoded.append(int(call))    # Zero-pad any trailing coordinates the raw vector did not cover.    encoded.extend([MISSING_ENCODED_AS] * (length - len(encoded)))    return encodeddef append_sentinel(encoded: list[int], sentinel: int = SENTINEL_VALUE) -> list[int]:    """Return ``encoded`` with the append-1 sentinel as the final slot.    When the server homomorphically sums N contributions, the sum path's trailing    slot sums to exactly N, and the SQUARE path's trailing slot sums to    ``sum_i 1^2 = N`` too — both decrypt to the exact contributor count. An    integrity/corruption check, NOT a MAC.    """    return list(encoded) + [sentinel]def encrypt(public_context_bytes: bytes, encoded: list[int]) -> bytes:    """BFV-encrypt ``encoded`` (with sentinel appended) -> serialized ciphertext.    Uses the PUBLIC context only. A SINGLE ciphertext per contributor is produced    (the server squares it under encryption); the public context carries relin    keys so that square is legal. The secret key is not touched here.    """    import tenseal as ts    context = ts.context_from(public_context_bytes)    plaintext = append_sentinel(encoded)    vector = ts.bfv_vector(context, plaintext)    return vector.serialize()