mecfs_bio.build_system.task.ppp_ldsc.batched_ldsc_h2

Batched univariate LD-score-regression heritability for the UKB-PPP database.

Because every protein's chi-square is defined on the SAME shared SNP set (see ppp_ldsc_context), we process K proteins at once: stack chi-square into an (S, K) matrix and run the weighted LD-score regression plus block jackknife for all K in vectorized numpy. Per-protein exclusions (missing variants, the chi-square filter, and cis variants) are encoded as ZERO regression weights rather than by dropping rows, so all proteins keep the same row set and the same shared contiguous jackknife blocks.

This is validated (experiments/claude/ppp_ldsc/batched_vs_exact_h2_probe.py) to agree with the repo's exact per-protein GenomicSEM port (genomic_sem_ldsc.estimate_h2): the h2 point estimate is machine-identical and the jackknife SE agrees to <=~1.5% (the only difference is that the exact method cuts the KEPT SNPs into blocks whereas the shared-block method cuts the FULL set, shifting block membership by a few SNPs).

Weighting note (the one easy bug): the effective WLS weight is omega = het * oc (the PRODUCT), because estimate_h2 multiplies BOTH the design and the response by initial_w = sqrt(het * oc). The sum-to-1 normalization is a common scalar that cancels in the 2x2 slope solve, so it is omitted here.

Classes:

• BatchedH2Result –

  Per-protein heritability outputs; all arrays are length K (protein order preserved).

• ExactH2Result –

Functions:

• batched_h2 –

  Estimate LDSC heritability for K proteins sharing the SNP set.

• exact_h2_single –

  Reference single-protein heritability via the repo's exact GenomicSEM port. Drops

Attributes:

• DEFAULT_N_BLOCKS –

DEFAULT_N_BLOCKS module-attribute

DEFAULT_N_BLOCKS = 200

BatchedH2Result

Per-protein heritability outputs; all arrays are length K (protein order preserved).

Attributes:

• h2 (ndarray) –
• h2_se (ndarray) –
• intercept (ndarray) –
• lambda_gc (ndarray) –
• mean_chi2 (ndarray) –
• n_snps (ndarray) –

h2 instance-attribute

h2: ndarray

h2_se instance-attribute

h2_se: ndarray

intercept instance-attribute

intercept: ndarray

lambda_gc instance-attribute

lambda_gc: ndarray

mean_chi2 instance-attribute

mean_chi2: ndarray

n_snps instance-attribute

n_snps: ndarray

ExactH2Result

Attributes:

• h2 (float) –
• h2_se (float) –
• intercept (float) –

h2 instance-attribute

h2: float

h2_se instance-attribute

h2_se: float

intercept instance-attribute

intercept: float

batched_h2

batched_h2(
    chi2: ndarray,
    ld: ndarray,
    n: ndarray,
    m: float,
    *,
    n_blocks: int = DEFAULT_N_BLOCKS,
    exclude: ndarray | None = None,
) -> BatchedH2Result

Estimate LDSC heritability for K proteins sharing the SNP set.

chi2: (S, K) chi-square (= (BETA/SE)^2), NaN where a protein lacks the variant. ld: (S,) LD score (== weight LD score), genome-sorted for contiguous blocks. n: (K,) per-protein sample size (constant across a protein's SNPs). m: total reference-SNP count. exclude: optional (S, K) boolean; True drops that variant for that protein (e.g. cis).

NOTE: - To understand the weighting, see the docstring for _het_oc_initial_weight in genomic_sem_ldsc.

Source code in mecfs_bio/build_system/task/ppp_ldsc/batched_ldsc_h2.py

def batched_h2(
    chi2: np.ndarray,
    ld: np.ndarray,
    n: np.ndarray,
    m: float,
    *,
    n_blocks: int = DEFAULT_N_BLOCKS,
    exclude: np.ndarray | None = None,
) -> BatchedH2Result:
    """Estimate LDSC heritability for K proteins sharing the SNP set.

    chi2: (S, K) chi-square (= (BETA/SE)^2), NaN where a protein lacks the variant.
    ld: (S,) LD score (== weight LD score), genome-sorted for contiguous blocks.
    n: (K,) per-protein sample size (constant across a protein's SNPs).
    m: total reference-SNP count.
    exclude: optional (S, K) boolean; True drops that variant for that protein (e.g. cis).


    NOTE:
        - To understand the weighting, see the docstring for _het_oc_initial_weight in genomic_sem_ldsc.
    """
    s, k = chi2.shape
    keep = np.isfinite(chi2) & (chi2 <= _chisq_threshold(n)[None, :])
    if exclude is not None:
        keep &= ~exclude
    chi = np.where(keep, chi2, 0.0)

    # Aggregate-h2 per protein over KEPT SNPs (GenomicSEM tot.agg), for the weights.
    cnt = keep.sum(0).astype(float)  # (K,)
    ld_col = ld[:, None]
    mean_chi = chi.sum(0) / cnt
    mean_ldn = (
        np.where(keep, ld_col, 0.0).sum(0) * n
    ) / cnt  # mean(ld)*N == mean(ld*N)
    tot_agg = np.clip(m * (mean_chi - 1.0) / mean_ldn, 0.0, 1.0)  # (K,)

    # Effective WLS weight omega = het * oc (see module docstring), zeroed off the kept set.
    ldm = np.maximum(ld, 1.0)[:, None]  # wLD == LD here
    c = (tot_agg * n / m)[None, :]
    het = 1.0 / (2.0 * (1.0 + c * ldm) ** 2)
    oc = 1.0 / ldm
    w = het * oc * keep  # (S, K)

    # Per-block weighted sums for the 2x2 normal equations of design [ld, 1]:
    #   a = sum w*ld^2, b = sum w*ld, d = sum w ; e = sum w*ld*chi, f = sum w*chi.
    ld2 = (ld * ld)[:, None]
    blk = np.empty((n_blocks, k, 5))
    for i, (lo, hi) in enumerate(block_bounds(s, n_blocks)):
        wb = w[lo:hi]
        cb = chi[lo:hi]
        blk[i, :, 0] = (wb * ld2[lo:hi]).sum(0)
        blk[i, :, 1] = (wb * ld_col[lo:hi]).sum(0)
        blk[i, :, 2] = wb.sum(0)
        blk[i, :, 3] = (wb * ld_col[lo:hi] * cb).sum(0)
        blk[i, :, 4] = (wb * cb).sum(0)
    tot = blk.sum(0)  # (K, 5)

    def _solve(sums: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        a, b, d, e, f = (sums[..., j] for j in range(5))
        det = a * d - b * b
        slope = (d * e - b * f) / det
        intercept = (a * f - b * e) / det
        return slope, intercept

    slope_full, intercept_full = _solve(tot)

    # Leave-one-block-out slopes -> pseudo-values of the raw slope -> jackknife SE.
    slope_loo, _ = _solve(tot[None, :, :] - blk)  # (n_blocks, K)
    pseudo = n_blocks * slope_full[None, :] - (n_blocks - 1) * slope_loo
    denom = (n * np.sqrt(n_blocks) / m) ** 2
    se = np.sqrt(np.var(pseudo, axis=0, ddof=1) / denom)

    # Genomic control uses the median chi-square over kept SNPs.
    chi_masked = np.where(keep, chi2, np.nan)
    lambda_gc = np.nanmedian(chi_masked, axis=0) / _CHI2_1DF_MEDIAN

    return BatchedH2Result(
        h2=slope_full / n * m,
        h2_se=se,
        intercept=intercept_full,
        mean_chi2=mean_chi,
        lambda_gc=lambda_gc,
        n_snps=cnt.astype(np.int64),
    )

exact_h2_single

exact_h2_single(
    chi2: ndarray,
    ld: ndarray,
    n: float,
    m: float,
    *,
    n_blocks: int = DEFAULT_N_BLOCKS,
    exclude: ndarray | None = None,
) -> ExactH2Result

Reference single-protein heritability via the repo's exact GenomicSEM port. Drops (rather than zero-weights) filtered/cis SNPs, then blocks the KEPT SNPs -- this is what the batched kernel is validated against. chi2/ld are genome-sorted (S,) arrays.

Source code in mecfs_bio/build_system/task/ppp_ldsc/batched_ldsc_h2.py

def exact_h2_single(
    chi2: np.ndarray,
    ld: np.ndarray,
    n: float,
    m: float,
    *,
    n_blocks: int = DEFAULT_N_BLOCKS,
    exclude: np.ndarray | None = None,
) -> ExactH2Result:
    """Reference single-protein heritability via the repo's exact GenomicSEM port. Drops
    (rather than zero-weights) filtered/cis SNPs, then blocks the KEPT SNPs -- this is what
    the batched kernel is validated against. chi2/ld are genome-sorted (S,) arrays."""
    keep = np.isfinite(chi2) & (chi2 <= max(0.001 * n, 80.0))
    if exclude is not None:
        keep &= ~exclude
    chi = chi2[keep]
    ld_kept = ld[keep]
    est = estimate_h2(
        chi=chi,
        ld_raw=ld_kept,
        wld_raw=ld_kept,
        n=np.full(chi.shape, n),
        m=m,
        n_blocks=n_blocks,
    )
    # run_ldsc: V_h2 = var(pseudo, ddof=1) / (n_bar * sqrt(n_blocks) / m)^2 ; se = sqrt(V).
    denom = (est.n_bar * np.sqrt(n_blocks) / m) ** 2
    se = float(np.sqrt(np.var(est.pseudo_coef, ddof=1) / denom))
    return ExactH2Result(h2=est.reg_tot, h2_se=se, intercept=est.intercept)