From 92e37566c4f466a7d08acb77da1c7a8d268fc5cc Mon Sep 17 00:00:00 2001 From: Jeff Swigert Date: Mon, 31 Aug 2026 15:19:59 -0600 Subject: [PATCH] Add --per-block-alpha: shrink omega per connected component Omega off-diagonal is block-diagonal over the connected components of the relatedness graph, since a pair only contributes to the block containing both of its members. The eigenvalues of a block-diagonal matrix are the union of its blocks, so positive-definiteness is a PER-BLOCK property and there is no mathematical requirement for alpha to be a single global scalar. With one global alpha the single worst component sets the shrinkage for every individual. On All of Us educational attainment at 3rd degree we measure 13 of 43,903 components breaking PD, holding 1.45% of the sample, forcing alpha to 0.4284 and so a 57% shrinkage on the other 98.5%. One 1,351-person component accounts for almost all of it. Shrinking each component only as far as it requires preserves PD identically: every block independently satisfies pheno_variance + alpha_c*lambda_min_c > 0, so the union of eigenvalues is positive. Opt-in. Default behaviour is unchanged and tested to be so. Implementation notes: - get_block_alphas returns per-component alphas plus the component labels. calculate_omega then scales each entry of the off-diagonal by its own component alpha, which needs no block reassembly since both members of a pair share a component. - A Gershgorin screen skips the eigendecomposition for any block whose largest absolute row sum is already below pheno_variance, since such a block cannot break PD. That covers the overwhelming majority of ordinary families and keeps the cost near the single global call it replaces. A test verifies the screen never skips a block that needed shrinking, by checking every alpha against a direct eigendecomposition. Five tests added, and the existing suite passes unchanged (202 total). One pre-existing issue noticed while testing, NOT addressed here: get_lambda_min calls ARPACK without a fixed starting vector, so repeated calls on the same matrix differ in the last ULP or two (-1.0 vs -0.99999999999999989 observed). Results are therefore not bitwise reproducible on the global path either. Happy to fix separately if wanted. --- grma.py | 15 ++++- grma_lib.py | 106 +++++++++++++++++++++++++++++-- test/grma_lib_test.py | 144 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 256 insertions(+), 9 deletions(-) diff --git a/grma.py b/grma.py index e95d712..f8f6d0a 100755 --- a/grma.py +++ b/grma.py @@ -80,6 +80,7 @@ SNPS_PER_BLOCK = "SNPs Per Block" ID_LIST = "ID List" SNP_LIST = "SNP List" +PER_BLOCK_ALPHA = "Per Block Alpha" # Type declaration @@ -243,6 +244,14 @@ def get_grma_parser(progname: str) -> argp.ArgumentParser: f"Default is {DEFAULT_REL_DEG}") + a_opt.add_argument("--per-block-alpha", action="store_true", + help="Shrink omega's off-diagonal separately for each connected " + "component of the relatedness graph, rather than applying a " + "single global factor. Positive-definiteness is a per-block " + "property, so this preserves it identically while preventing " + "one pedigree from setting the shrinkage for the whole " + "sample. Default is the single global factor.") + infilt_opt = parser.add_argument_group(title="Input Filtering Options") infilt_opt.add_argument("--id-list", metavar="FILE", type=input_file, help="Optional input to specify a whitespace-delimited sample ID file " @@ -446,7 +455,8 @@ def validate_inputs(pargs: argp.Namespace, user_args: Dict[str, Any]): REL_DEG : pargs.rel_thresh, SNPS_PER_BLOCK : pargs.snps_per_block, ID_LIST : pargs.id_list, - SNP_LIST : pargs.snp_list + SNP_LIST : pargs.snp_list, + PER_BLOCK_ALPHA : pargs.per_block_alpha } @@ -507,7 +517,8 @@ def main_func(argv: List[str]): covar_file=iargs[COVAR_FILE], snps_per_block=iargs[SNPS_PER_BLOCK], id_list=iargs[ID_LIST], - snp_list=iargs[SNP_LIST] + snp_list=iargs[SNP_LIST], + per_block_alpha=iargs[PER_BLOCK_ALPHA] ) # Write out the results to disk per chromosome diff --git a/grma_lib.py b/grma_lib.py index 00f7767..9b6e7b9 100755 --- a/grma_lib.py +++ b/grma_lib.py @@ -13,6 +13,7 @@ import pandas as pd import scipy.sparse as sp import scipy.sparse.linalg as spla +from scipy.sparse.csgraph import connected_components from scipy.stats import norm from bedbimfam import ( @@ -354,8 +355,84 @@ def get_rel_covariances(rel_df: pd.DataFrame, return rel_to_cov +def get_block_alphas(off_diag: sp.csr_array, pheno_variance: float, + epsilon=DEFAULT_OMEGA_EPSILON) -> Tuple[np.ndarray, np.ndarray]: + """Computes a shrinkage factor for each connected component of off_diag. + + Omega's off-diagonal is block-diagonal over the connected components of the + relatedness graph: a pair only ever contributes to the block containing both + of its members. The eigenvalues of a block-diagonal matrix are the union of + its blocks', so positive-definiteness is a PER-BLOCK property and there is no + mathematical requirement for alpha to be a single global scalar. + + With one global alpha, the single worst component sets the shrinkage for + every individual in the sample. Shrinking each component only as far as that + component requires preserves positive-definiteness identically -- every block + independently satisfies pheno_variance + alpha_c * lambda_min_c > 0, so the + union of the eigenvalues is positive -- while leaving unaffected blocks + untouched. + + Args: + off_diag: The unshrunk off-diagonal, as built by calculate_omega + pheno_variance: Variance of the unresidualized phenotypes + epsilon: Small positive constant keeping the result strictly definite + + Returns: + A tuple (alphas, labels) where alphas[c] is the shrinkage factor for + component c and labels[i] is the component index of sample i. + """ + + n_comp, labels = connected_components(off_diag, directed=False) + alphas = np.ones(n_comp, dtype=float) + sizes = np.bincount(labels, minlength=n_comp) + + # Gershgorin's circle theorem bounds every eigenvalue of a symmetric matrix + # by its largest absolute row sum, so a block whose bound is already below + # pheno_variance cannot break positive-definiteness and needs no + # eigendecomposition. This skips the overwhelming majority of blocks -- + # ordinary nuclear families -- and keeps the cost close to the single global + # call it replaces. + row_abs_sum = np.abs(off_diag).sum(axis=1) + row_abs_sum = np.asarray(row_abs_sum).ravel() + block_bound = np.zeros(n_comp, dtype=float) + np.maximum.at(block_bound, labels, row_abs_sum) + candidates = np.nonzero((sizes > 1) & (block_bound >= pheno_variance))[0] + logging.debug(f"\t{n_comp} components; {len(candidates)} need an " + f"eigendecomposition after the Gershgorin screen") + + if len(candidates) == 0: + return alphas, labels + + order = np.argsort(labels, kind="stable") + starts = np.searchsorted(labels[order], np.arange(n_comp)) + ends = np.searchsorted(labels[order], np.arange(n_comp), side="right") + + for c in candidates: + idx = order[starts[c]:ends[c]] + block = off_diag[idx][:, idx] + # eigsh needs k < n-1, and tiny blocks are cheaper and exact as dense + if block.shape[0] <= 3: + block_lambda_min = float(np.linalg.eigvalsh(block.toarray()).min()) + else: + block_lambda_min = float(get_lambda_min(block)) + if block_lambda_min < 0.0: + alphas[c] = min((epsilon - pheno_variance) / block_lambda_min, 1.0) + + shrunk = alphas < 1.0 + if np.any(shrunk): + logging.info(f"\tPer-block shrinkage: {int(shrunk.sum())} of {n_comp} " + f"components shrunk; smallest alpha {alphas.min():.6f}, " + f"{int(sizes[shrunk].sum())} of {off_diag.shape[0]} samples " + f"affected") + else: + logging.info("\tPer-block shrinkage: no component required shrinking") + + return alphas, labels + + def calculate_omega(rel_df: pd.DataFrame, N: int, unresidualized_phenotypes: np.ndarray, - epsilon=DEFAULT_OMEGA_EPSILON) -> sp.csr_array: + epsilon=DEFAULT_OMEGA_EPSILON, + per_block_alpha: bool = False) -> sp.csr_array: rel_to_cov = get_rel_covariances(rel_df=rel_df, unresidualized_phenotypes=unresidualized_phenotypes) @@ -371,11 +448,21 @@ def calculate_omega(rel_df: pd.DataFrame, N: int, unresidualized_phenotypes: np. off_diag = sp.coo_array((data, (rows, cols)), shape=(N, N)).tocsr() + pheno_variance = np.var(unresidualized_phenotypes) + logging.debug(f"\t{pheno_variance=}") + + if per_block_alpha: + # Scale each entry by the alpha of the component it belongs to. Both + # members of a pair are in the same component by construction, so + # labels[rows] is the component of every entry in data. + alphas, labels = get_block_alphas(off_diag, pheno_variance, epsilon) + scaled = data * alphas[labels[rows]] + off_diag = sp.coo_array((scaled, (rows, cols)), shape=(N, N)).tocsr() + return pheno_variance * sp.eye(N) + off_diag + lambda_min = get_lambda_min(off_diag) logging.debug(f"\t{lambda_min=}") - pheno_variance = np.var(unresidualized_phenotypes) - logging.debug(f"\t{pheno_variance=}") alpha = min((epsilon - pheno_variance) / lambda_min, 1.0) if lambda_min < 0.0 else 1.0 logging.debug(f"\t{alpha=}\n") @@ -410,7 +497,8 @@ def process_relatedness( rel_file: Union[str, pd.DataFrame], fam_df: pd.DataFrame, rel_degree: Union[str, int], - unresidualized_phenotypes: np.ndarray + unresidualized_phenotypes: np.ndarray, + per_block_alpha: bool = False ) -> Tuple[sp.csr_array, sp.csr_array, np.ndarray]: """Performs processing of the KING-formatted relatedness / pedigree file @@ -482,7 +570,9 @@ def process_relatedness( # Create omega matrix - omega = calculate_omega(rel_df=rel_df, N=N, unresidualized_phenotypes=unresidualized_phenotypes) + omega = calculate_omega(rel_df=rel_df, N=N, + unresidualized_phenotypes=unresidualized_phenotypes, + per_block_alpha=per_block_alpha) # Create R matrix rel_df = rel_df[rel_df[KING_REL_COL] <= rel_degree] @@ -825,7 +915,8 @@ def grma( covar_file: Union[str, pd.DataFrame] = None, snps_per_block: int = DEFAULT_SNPS_PER_BLOCK, id_list: Union[str, pd.DataFrame] = None, - snp_list: Union[str, pd.DataFrame] = None + snp_list: Union[str, pd.DataFrame] = None, + per_block_alpha: bool = False ) -> pd.DataFrame: logging.debug(f"GRMA called with: {locals()}\n\n") @@ -855,7 +946,8 @@ def grma( rel_file=rel_file, fam_df=fam_df, rel_degree=rel_degree, - unresidualized_phenotypes=unresidualized_phenotypes + unresidualized_phenotypes=unresidualized_phenotypes, + per_block_alpha=per_block_alpha ) del unresidualized_phenotypes logging.info(f"Processing relatedness info took {time.time() - rel_time} seconds\n") diff --git a/test/grma_lib_test.py b/test/grma_lib_test.py index 8bddbd7..5419526 100644 --- a/test/grma_lib_test.py +++ b/test/grma_lib_test.py @@ -429,6 +429,150 @@ def test__not_PSD__shrink_as_expected(self): ) assert np.allclose(actual_omega.toarray(), expected_omega.toarray(), atol=1e-3) + def test__per_block_alpha__single_component__matches_global(self): + # With only one multi-person component there is nothing to separate, so + # per-block shrinkage must reproduce the global result exactly. + N = 5 + pheno = np.array([1.0, -1.0, 0.1, 0.1, 0.1]) + rel_df = pd.DataFrame({ + sut.FAM_INDEX_1: [0], + sut.FAM_INDEX_2: [1], + sut.KING_REL_COL: [sut.DEG_PARENT_OFFSPRING], + }) + + global_omega = sut.calculate_omega(rel_df=rel_df, N=N, + unresidualized_phenotypes=pheno) + block_omega = sut.calculate_omega(rel_df=rel_df, N=N, + unresidualized_phenotypes=pheno, + per_block_alpha=True) + + assert np.allclose(global_omega.toarray(), block_omega.toarray()) + + def test__per_block_alpha__default_is_unchanged(self): + # The flag must be opt-in: omitting it has to give the same output as + # the released behaviour. + # + # allclose rather than array_equal: get_lambda_min calls ARPACK, which + # seeds itself randomly, so repeated calls on the SAME matrix differ in + # the last ULP or two (-1.0 vs -0.99999999999999989 observed). That + # non-determinism predates this change and affects the global path too. + N = 6 + pheno = np.array([1.0, -1.0, 0.9, -0.8, 0.1, 0.1]) + rel_df = pd.DataFrame({ + sut.FAM_INDEX_1: [0, 2], + sut.FAM_INDEX_2: [1, 3], + sut.KING_REL_COL: [sut.DEG_PARENT_OFFSPRING, sut.DEG_3RD], + }) + + default_omega = sut.calculate_omega(rel_df=rel_df, N=N, + unresidualized_phenotypes=pheno) + explicit_omega = sut.calculate_omega(rel_df=rel_df, N=N, + unresidualized_phenotypes=pheno, + per_block_alpha=False) + + assert np.allclose(default_omega.toarray(), explicit_omega.toarray(), + rtol=0, atol=1e-12) + + def test__per_block_alpha__innocent_block_not_shrunk(self): + # Two disjoint components carrying the SAME relatedness degree, so both + # get the same covariance rho, and any difference in their shrinkage is + # due to STRUCTURE alone: + # + # block A, a star 0-1 / 0-2 / 0-3 lambda_min = -rho*sqrt(3) + # block B, a single pair 4-5 lambda_min = -rho + # + # The phenotypes below put var(y) between the two, so the star breaks + # positive-definiteness and the pair does not. Under a global alpha the + # innocent pair is shrunk anyway; under per-block alpha it is untouched. + N = 10 + pheno = np.array([0.680193, -0.845163, -0.010330, 0.699210, -1.368729, + -1.756903, 0.215909, 0.036958, -0.204671, 0.032086]) + rel_df = pd.DataFrame({ + sut.FAM_INDEX_1: [0, 0, 0, 4], + sut.FAM_INDEX_2: [1, 2, 3, 5], + sut.KING_REL_COL: [sut.DEG_PARENT_OFFSPRING] * 4, + }) + + rel_to_cov = sut.get_rel_covariances(rel_df=rel_df, + unresidualized_phenotypes=pheno) + rho = rel_to_cov[sut.DEG_PARENT_OFFSPRING] + pheno_var = float(np.var(pheno)) + # the premise of the test + assert rho < pheno_var < rho * np.sqrt(3.0) + + global_omega = sut.calculate_omega(rel_df=rel_df, N=N, + unresidualized_phenotypes=pheno).toarray() + block_omega = sut.calculate_omega(rel_df=rel_df, N=N, + unresidualized_phenotypes=pheno, + per_block_alpha=True).toarray() + + # Global: the star's shrinkage is imposed on the innocent pair too. + assert np.isclose(global_omega[0, 1], global_omega[4, 5]) + assert abs(global_omega[4, 5]) < abs(rho) + + # Per block: the pair keeps its unshrunk covariance, the star does not. + assert np.isclose(block_omega[4, 5], rho) + assert abs(block_omega[0, 1]) < abs(rho) + assert not np.isclose(block_omega[0, 1], block_omega[4, 5]) + + # And the whole matrix is still positive semi-definite. + assert np.all(np.linalg.eigvalsh(block_omega) >= -1e-8) + + def test__per_block_alpha__result_is_psd(self): + # The guarantee that matters: shrinking per block must still leave the + # whole matrix positive semi-definite. + rng = np.random.default_rng(20260831) + N = 40 + pheno = rng.normal(size=N) + pairs = [(i, i + 1) for i in range(0, N - 1, 2)] + [(0, 2), (4, 6), (8, 10)] + rel_df = pd.DataFrame({ + sut.FAM_INDEX_1: [p[0] for p in pairs], + sut.FAM_INDEX_2: [p[1] for p in pairs], + sut.KING_REL_COL: [sut.DEG_PARENT_OFFSPRING] * len(pairs), + }) + + omega = sut.calculate_omega(rel_df=rel_df, N=N, + unresidualized_phenotypes=pheno, + per_block_alpha=True) + + assert np.all(np.linalg.eigvalsh(omega.toarray()) >= -1e-8) + + def test__get_block_alphas__gershgorin_screen_is_safe(self): + # The Gershgorin screen skips blocks it can prove are definite. Verify + # it never skips one that needed shrinking, by checking every returned + # alpha against a direct per-block eigendecomposition. + rng = np.random.default_rng(1234) + N = 60 + pheno = rng.normal(size=N) + pairs = [(i, i + 1) for i in range(0, N - 1, 3)] + [(0, 3), (6, 9)] + rel_df = pd.DataFrame({ + sut.FAM_INDEX_1: [p[0] for p in pairs], + sut.FAM_INDEX_2: [p[1] for p in pairs], + sut.KING_REL_COL: [sut.DEG_PARENT_OFFSPRING] * len(pairs), + }) + rel_to_cov = sut.get_rel_covariances(rel_df=rel_df, + unresidualized_phenotypes=pheno) + i1 = rel_df[sut.FAM_INDEX_1].to_numpy(np.int64) + i2 = rel_df[sut.FAM_INDEX_2].to_numpy(np.int64) + data = np.tile(rel_df[sut.KING_REL_COL].map(rel_to_cov).to_numpy(), 2) + off_diag = sp.coo_array((data, (np.concatenate([i1, i2]), + np.concatenate([i2, i1]))), + shape=(N, N)).tocsr() + pheno_var = float(np.var(pheno)) + + alphas, labels = sut.get_block_alphas(off_diag, pheno_var) + + n_comp = labels.max() + 1 + for c in range(n_comp): + idx = np.flatnonzero(labels == c) + if len(idx) < 2: + continue + block = off_diag[idx][:, idx].toarray() + lmin = float(np.linalg.eigvalsh(block).min()) + expected = (min((sut.DEFAULT_OMEGA_EPSILON - pheno_var) / lmin, 1.0) + if lmin < 0.0 else 1.0) + assert np.isclose(alphas[c], expected, atol=1e-6) +