Add --per-block-alpha: shrink omega per connected component

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

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

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