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) +