graphld.heritability¶

Heritability estimation using graphREML.

This module implements the graphREML method for heritability partitioning and enrichment analysis using GWAS summary statistics and LDGM precision matrices.

For workflow-oriented examples, see the Heritability Estimation guide.

heritability ¶

GraphREML

ModelOptions `dataclass` ¶

ModelOptions(annotation_columns: Optional[List[str]] = None, params: Optional[ndarray] = None, sample_size: Optional[float] = None, intercept: float = 1.0, link_fn_denominator: float = 6000000.0, binary_annotations_only: bool = False)

Stores model parameters for graphREML.

Attributes:

Name	Type	Description
`annotation_columns`	`Optional[List[str]]`	names of columns to be used as annotations
`params`	`Optional[ndarray]`	Starting parameter values
`sample_size`	`Optional[float]`	GWAS sample size, only needed for heritability scaling
`intercept`	`float`	LDSC intercept or 1
`link_fn_denominator`	`float`	Scalar denominator for link function.

MethodOptions `dataclass` ¶

MethodOptions(gradient_num_samples: int = 100, gradient_seed: Optional[int] = 123, match_by_position: bool = False, num_iterations: int = 50, convergence_tol: float = 0.01, run_serial: bool = False, num_processes: Optional[int] = None, verbose: bool = False, use_surrogate_markers: bool = True, trust_region_size: float = 0.1, trust_region_rho_lb: float = 0.0001, trust_region_rho_ub: float = 0.99, trust_region_scalar: float = 5, max_trust_iterations: int = 100, minimum_likelihood_increase: float = 1e-06, convergence_window: int = 3, reset_trust_region: bool = False, num_jackknife_blocks: int = 100, max_chisq_threshold: Optional[float] = None, score_test_hdf5_file_name: Optional[str] = None, score_test_hdf5_trait_name: Optional[str] = None, surrogate_markers_path: Optional[str] = None)

Stores method parameters for graphREML.

Attributes:

Name	Type	Description
`gradient_num_samples`	`int`	Number of samples for gradient estimation
`match_by_position`	`bool`	Use position/allele instead of RSID for merging
`num_iterations`	`int`	Optimization steps
`convergence_tol`	`float`	Convergence tolerance
`run_serial`	`bool`	Run in serial rather than parallel
`num_processes`	`Optional[int]`	If None, autodetect
`verbose`	`bool`	Flag for verbose output
`use_surrogate_markers`	`bool`	Whether to use surrogate markers for missing variants
`trust_region_size`	`float`	Initial trust region size parameter
`trust_region_rho_lb`	`float`	Lower bound for trust region ratio
`trust_region_rho_ub`	`float`	Upper bound for trust region ratio
`trust_region_scalar`	`float`	Scaling factor for trust region updates
`max_trust_iterations`	`int`	Maximum number of trust region iterations
`reset_trust_region`	`bool`	Whether to reset trust region size at each iteration
`num_jackknife_blocks`	`int`	Number of blocks to use for jackknife estimation
`max_chisq_threshold`	`Optional[float]`	Maximum allowed chi^2 value in a block. Blocks with chi^2 > threshold are excluded.
`score_test_hdf5_file_name`	`Optional[str]`	Optional file name to create or append to an hdf5 file with pre-computed derivatives for the score test.
`score_test_hdf5_trait_name`	`Optional[str]`	Name of the trait's subdirectory within the score test HDF5 file.
`surrogate_markers_path`	`Optional[str]`	Optional path to an HDF5 file with per-block surrogate mappings.

GraphREML ¶

Bases: ParallelProcessor

prepare_block_data `classmethod` ¶

prepare_block_data(metadata: DataFrame, **kwargs: Any) -> list[tuple]

Prepare block-specific data for processing.

Parameters:

Name	Type	Description	Default
`metadata`	`DataFrame`	DataFrame containing LDGM metadata	required
`**kwargs`	`Any`	Additional arguments from run(), including: sumstats: DataFrame containing Z scores and variant info, optionally annotations annotation_columns: list of column names for the annotations in sumstats method: MethodOptions instance containing method parameters	`{}`

Returns:

Type	Description
`list[tuple]`	List of dictionaries containing block-specific data with keys: sumstats: DataFrame for this block, or None if max Z² exceeds threshold variant_offset: Cumulative number of variants before this block block_index: Index of this block Pz: Pre-computed precision-premultiplied Z scores for this block

Source code in src/graphld/heritability.py

@classmethod
def prepare_block_data(cls, metadata: pl.DataFrame, **kwargs: Any) -> list[tuple]:
    """Prepare block-specific data for processing.

    Args:
        metadata: DataFrame containing LDGM metadata
        **kwargs: Additional arguments from run(), including:
            sumstats: DataFrame containing Z scores and variant info, optionally annotations
            annotation_columns: list of column names for the annotations in sumstats
            method: MethodOptions instance containing method parameters

    Returns:
        List of dictionaries containing block-specific data with keys:
            sumstats: DataFrame for this block, or None if max Z² exceeds threshold
            variant_offset: Cumulative number of variants before this block
            block_index: Index of this block
            Pz: Pre-computed precision-premultiplied Z scores for this block
    """
    sumstats: pl.DataFrame = kwargs.get("sumstats")
    method: MethodOptions = kwargs.get("method")
    sumstats_blocks: list[pl.DataFrame] = partition_variants(metadata, sumstats)
    num_block_variants = sum(len(block) for block in sumstats_blocks)
    assert num_block_variants <= sumstats.shape[0], (
        f"Too many variants in blocks: {num_block_variants} > {sumstats.shape[0]}"
    )
    if method.verbose:
        print(
            f"{sum(len(block) for block in sumstats_blocks)} variants in blocks, {sumstats.shape[0]} initially"
        )
        print(
            f"Smallest block size: {min(len(block) for block in sumstats_blocks)}"
        )
        print(f"Largest block size: {max(len(block) for block in sumstats_blocks)}")

    # Filter blocks based on max Z² threshold
    if method.max_chisq_threshold is not None:
        max_z2s = [_block_max_chisq(block) for block in sumstats_blocks]
        keep_block = [max_z2 <= method.max_chisq_threshold for max_z2 in max_z2s]
        sumstats_blocks = [
            block if keep else block.head(0)
            for block, keep in zip(sumstats_blocks, keep_block, strict=False)
        ]
        if method.verbose and not all(keep_block):
            print(
                f"{len(sumstats_blocks) - sum(keep_block)} out of {len(sumstats_blocks)} blocks discarded due to\n"
                f"max chisq threshold of {method.max_chisq_threshold}"
            )
        if not any(keep_block):
            raise ValueError(
                "No LD blocks remain after applying max_chisq_threshold="
                f"{method.max_chisq_threshold}."
            )

    cumulative_num_variants = np.cumsum(
        np.array([len(df) for df in sumstats_blocks])
    )
    cumulative_num_variants = [0] + list(cumulative_num_variants[:-1])
    block_indices = list(range(len(sumstats_blocks)))
    block_Pz = [None for _ in block_indices]
    block_names = metadata.get_column("name").to_list() if len(metadata) > 0 else []

    return [
        {
            "sumstats": block,
            "variant_offset": offset,
            "block_index": index,
            "Pz": Pz,
            "block_name": block_names[index] if index < len(block_names) else None,
        }
        for block, offset, index, Pz in zip(
            sumstats_blocks,
            cumulative_num_variants,
            block_indices,
            block_Pz,
            strict=False,
        )
    ]

create_shared_memory `staticmethod` ¶

create_shared_memory(metadata: DataFrame, block_data: list[tuple], **kwargs: Any) -> SharedData

Create output array.

Parameters:

Name	Type	Description	Default
`metadata`	`DataFrame`	Metadata DataFrame containing block information	required
`block_data`	`list[tuple]`	List of block-specific sumstats DataFrames	required
`**kwargs`	`Any`	Not used	`{}`

Source code in src/graphld/heritability.py

@staticmethod
def create_shared_memory(
    metadata: pl.DataFrame, block_data: list[tuple], **kwargs: Any
) -> SharedData:
    """Create output array.

    Args:
        metadata: Metadata DataFrame containing block information
        block_data: List of block-specific sumstats DataFrames
        **kwargs: Not used
    """
    num_params = kwargs.get("num_params")
    num_blocks = len(metadata)
    num_variants = sum(
        [len(d["sumstats"]) for d in block_data if d["sumstats"] is not None]
    )

    result = SharedData(
        {
            "params": num_params,
            "variant_data": num_variants,
            "likelihood": num_blocks,
            "gradient": num_blocks * num_params,
            "hessian": num_blocks * num_params**2,
        }
    )

    return result

process_block `classmethod` ¶

process_block(ldgm: PrecisionOperator, flag: Value, shared_data: SharedData, block_offset: int, block_data: Any = None, worker_params: Tuple[ModelOptions, MethodOptions] = None)

Computes likelihood, gradient, and hessian for a single block.

Source code in src/graphld/heritability.py

@classmethod
def process_block(
    cls,
    ldgm: PrecisionOperator,
    flag: Value,
    shared_data: SharedData,
    block_offset: int,
    block_data: Any = None,
    worker_params: Tuple[ModelOptions, MethodOptions] = None,
):
    """Computes likelihood, gradient, and hessian for a single block."""
    model_options: ModelOptions
    method_options: MethodOptions
    model_options, method_options = worker_params
    seed = None
    if method_options.gradient_seed is not None:
        seed = method_options.gradient_seed + block_data["block_index"]
        np.random.seed(seed)

    sumstats: pl.DataFrame = block_data["sumstats"]
    if sumstats is None or len(sumstats) == 0:
        block_data["Pz"] = None
        return

    Pz: np.ndarray
    if flag.value == FLAGS["INITIALIZE"]:
        surrogate_map = None
        if method_options.surrogate_markers_path is not None:
            surrogate_map = cls._load_block_surrogate_map(
                method_options.surrogate_markers_path, block_data.get("block_name")
            )

        ldgm, Pz = cls._initialize_block_zscores(
            ldgm,
            sumstats,
            model_options.annotation_columns,
            method_options.match_by_position,
            method_options.verbose,
            surrogate_map=surrogate_map,
        )

        # Work in effect-size as opposed to Z score units
        if Pz is not None:
            Pz /= np.sqrt(model_options.sample_size)
            ldgm.times_scalar(model_options.intercept / model_options.sample_size)
        block_data["Pz"] = Pz
    # ldgm is modified in place and re-used in subsequent iterations

    else:
        Pz = block_data["Pz"]

    if Pz is None:
        return

    annot_indices: np.ndarray = (
        ldgm.variant_info.select("annot_indices").to_numpy().flatten()
    )
    max_index: int = np.max(annot_indices) + 1 if len(annot_indices) > 0 else 0
    variant_offset: int = block_data["variant_offset"]
    block_variants: slice = slice(variant_offset, variant_offset + max_index)
    annot: np.ndarray = ldgm.variant_info.select(
        model_options.annotation_columns
    ).to_numpy()
    params: np.ndarray = shared_data["params"].reshape(-1, 1)

    # Handle variant-specific gradient and Hessian computation
    if flag.value == FLAGS["COMPUTE_VARIANT_SCORE"]:
        del_h2i_del_xi, _ = cls._link_fn_derivatives(
            annot, params, model_options.link_fn_denominator
        )
        result = cls._compute_variant_grad(ldgm, Pz, del_h2i_del_xi, seed=seed)
        result_padded = np.zeros(max_index)
        result_padded[annot_indices] = result
        shared_data["variant_data", block_variants] = result_padded
        return

    if flag.value == FLAGS["COMPUTE_VARIANT_HESSIAN"]:
        del_h2i_del_xi, del2_h2i_del_xi2 = cls._link_fn_derivatives(
            annot, params, model_options.link_fn_denominator
        )
        del_L_del_xi = shared_data["variant_data", block_variants][
            annot_indices
        ].ravel()
        result = cls._compute_variant_hessian_diag(
            ldgm, Pz, del_L_del_xi, del_h2i_del_xi, del2_h2i_del_xi2, seed=seed
        )
        result_padded = np.zeros(max_index)
        result_padded[annot_indices] = result.ravel()
        shared_data["variant_data", block_variants] = result_padded
        return

    block_index: int = block_data["block_index"]
    num_annot = len(model_options.annotation_columns)

    old_variant_h2: np.ndarray = shared_data["variant_data", block_variants][
        annot_indices
    ]

    likelihood_only = flag.value == FLAGS["COMPUTE_LIKELIHOOD_ONLY"]
    likelihood, gradient, hessian, variant_h2 = cls._compute_block_likelihood(
        ldgm=ldgm,
        Pz=Pz,
        annotations=annot,
        params=params,
        link_fn_denominator=model_options.link_fn_denominator,
        old_variant_h2=old_variant_h2,
        num_samples=method_options.gradient_num_samples,
        likelihood_only=likelihood_only,
        seed=seed,
    )

    shared_data["likelihood", block_index] = likelihood
    variant_h2_padded = np.zeros(max_index)
    variant_h2_padded[annot_indices] = variant_h2.ravel()
    shared_data["variant_data", block_variants] = variant_h2_padded
    if likelihood_only:
        return

    gradient_slice = slice(block_index * num_annot, (block_index + 1) * num_annot)
    shared_data["gradient", gradient_slice] = gradient.flatten()

    hessian_slice = slice(
        block_index * num_annot**2, (block_index + 1) * num_annot**2
    )
    shared_data["hessian", hessian_slice] = hessian.flatten()

supervise `classmethod` ¶

supervise(manager: WorkerManager, shared_data: SharedData, block_data: list, **kwargs)

Runs graphREML. Args: manager: used to start parallel workers shared_data: used to communicate with workers block_data: annotation + gwas data passed to workers **kwargs: Additional arguments

Source code in src/graphld/heritability.py

@classmethod
def supervise(
    cls, manager: WorkerManager, shared_data: SharedData, block_data: list, **kwargs
):
    """Runs graphREML.
    Args:
        manager: used to start parallel workers
        shared_data: used to communicate with workers
        block_data: annotation + gwas data passed to workers
        **kwargs: Additional arguments
    """

    num_iterations = kwargs.get("num_iterations")
    num_params = kwargs.get("num_params")
    verbose = kwargs.get("verbose")
    method: MethodOptions = kwargs.get("method")
    model: ModelOptions = kwargs.get("model")
    trust_region_lambda = method.trust_region_size
    log_likelihood_history = []
    trust_region_history = []

    def _trust_region_step(
        gradient: np.ndarray, hessian: np.ndarray, trust_region_lambda: float
    ) -> np.ndarray:
        """Compute trust region step by solving (H + λD)x = -g."""
        hess_mod = hessian + trust_region_lambda * np.diag(
            np.diag(hessian) - np.finfo(float).eps
        )
        step = np.linalg.solve(hess_mod, -gradient)
        # predicted_increase = step.T @ gradient + 0.5 * step.T @ (hess_mod @ step)
        predicted_increase = step.T @ gradient + 0.5 * step.T @ (hessian @ step)
        assert predicted_increase > -1e-6, (
            f"Predicted increase must be greater than -epsilon but is {predicted_increase}."
        )

        return step, predicted_increase

    if model.params is not None:
        shared_data["params"] = model.params.flatten()
    else:
        shared_data["params"] = np.full(num_params, 0)

    last_step_bad = True
    for rep in range(num_iterations):
        if verbose:
            print(f"\n\tStarting iteration {rep}...")

        # Calculate likelihood, gradient, and hessian for each block
        flag = FLAGS["INITIALIZE"] if rep == 0 else FLAGS["COMPUTE_ALL"]
        manager.start_workers(flag)
        manager.await_workers()

        likelihood = cls._sum_blocks(shared_data["likelihood"], (1,))[0]
        gradient = cls._sum_blocks(shared_data["gradient"], (num_params,))
        hessian = cls._sum_blocks(shared_data["hessian"], (num_params, num_params))

        old_params = shared_data["params"].copy()
        old_likelihood = likelihood
        if verbose and rep == 0:
            print(f"Initial log likelihood: {likelihood}")

        # Reset trust region size if specified
        if method.reset_trust_region or last_step_bad:
            trust_region_lambda = method.trust_region_size

        # Trust region optimization loop
        previous_lambda = trust_region_lambda
        for trust_iter in range(method.max_trust_iterations):
            # Compute proposed step
            step, predicted_increase = _trust_region_step(
                gradient, hessian, trust_region_lambda
            )
            shared_data["params"] = old_params + step

            # Evaluate proposed step
            manager.start_workers(FLAGS["COMPUTE_LIKELIHOOD_ONLY"])
            manager.await_workers()
            new_likelihood = cls._sum_blocks(shared_data["likelihood"], (1,))[0]

            # Compute actual vs predicted increase
            actual_increase = new_likelihood - old_likelihood
            if verbose:
                print(
                    f"\tIncrease in log-likelihood: {actual_increase}, predicted increase: {predicted_increase}"
                )

            # Check if step is acceptable and update trust region size if needed
            rho = actual_increase / predicted_increase
            if rho < method.trust_region_rho_lb:
                if (
                    predicted_increase < method.minimum_likelihood_increase
                    and rho >= 0
                ):
                    if verbose:
                        print(
                            f"\tTerminated trust region size search with predicted likelihood increase {predicted_increase}."
                        )
                    break
                # Reset trust region size to initial value if its below that
                trust_region_lambda = max(
                    method.trust_region_size,
                    trust_region_lambda * method.trust_region_scalar,
                )
                shared_data["params"] = old_params  # Revert step
            elif rho > method.trust_region_rho_ub:
                trust_region_lambda /= method.trust_region_scalar
                break  # Accept step and continue to next iteration
            else:
                break  # Accept step with current trust region size

            if trust_iter == method.max_trust_iterations - 1:
                if verbose:
                    print("Warning: Maximum trust region iterations reached")

        last_step_bad = (
            trust_region_lambda > previous_lambda * method.trust_region_scalar
        )

        log_likelihood_history.append(new_likelihood)
        trust_region_history.append(trust_region_lambda)

        if verbose:
            print(f"log likelihood: {new_likelihood}")
            print(f"Trust region lambda: {trust_region_lambda}")
            if len(log_likelihood_history) >= 2:
                print(
                    f"Change in likelihood: {log_likelihood_history[-1] - log_likelihood_history[-2]}"
                )
            total_h2 = np.sum(shared_data["variant_data"])
            print(f"Total h2 at current iteration: {total_h2}")

        # Check convergence
        converged = False
        if len(log_likelihood_history) >= 1 + method.convergence_window:
            if abs(
                log_likelihood_history[-1]
                - log_likelihood_history[-method.convergence_window]
            ) < (method.convergence_window * method.convergence_tol):
                converged = True
                break

    if verbose:
        print(
            f"-----Finished optimization after {rep + 1} out of {num_iterations} steps-----"
        )

    # Point estimates
    variant_h2 = shared_data["variant_data"].copy()
    annotations = pl.concat(
        [
            dict["sumstats"].select(model.annotation_columns + SPECIAL_COLNAMES)
            for dict in block_data
            if len(dict["sumstats"]) > 0
        ]
    )
    # Enrichment is normalized against the first annotation column.
    ref_col = 0
    annotation_heritability, annotation_enrichment = cls._annotation_heritability(
        variant_h2, annotations.select(model.annotation_columns), ref_col
    )

    # Get block-wise gradients and Hessians for jackknife
    num_blocks = len(block_data)
    gradient_blocks = shared_data["gradient"].reshape((num_blocks, num_params))
    hessian_blocks = shared_data["hessian"].reshape(
        (num_blocks, num_params, num_params)
    )

    # Group blocks for jackknife
    num_jackknife_blocks = min(method.num_jackknife_blocks, num_blocks)
    jk_gradient_blocks = cls._group_blocks(gradient_blocks, num_jackknife_blocks)
    jk_hessian_blocks = cls._group_blocks(hessian_blocks, num_jackknife_blocks)

    # Compute jackknife estimates using the grouped blocks
    params = shared_data["params"].copy()
    jackknife_params = cls._compute_pseudojackknife(
        jk_gradient_blocks, jk_hessian_blocks, params
    )

    # Compute jackknife heritability estimates and standard errors
    jackknife_h2, jackknife_annot_sums = cls._compute_jackknife_heritability(
        block_data, jackknife_params, model
    )

    if method.score_test_hdf5_file_name is not None:
        if verbose:
            print("\n\tComputing variant scores...")
        manager.start_workers(FLAGS["COMPUTE_VARIANT_SCORE"])
        manager.await_workers()
        variant_score = shared_data["variant_data"].copy()

        # Jackknife block to which each variant belongs
        if verbose:
            print("\n\tComputing jackknife variant assignments...")
        block_indptrs = [dict["variant_offset"] for dict in block_data] + [
            len(variant_score)
        ]
        jackknife_variant_assignments = cls._get_variant_jackknife_assignments(
            block_indptrs, num_jackknife_blocks
        )

        # Project the annotations out of the score
        annotations_matrix = annotations.select(model.annotation_columns).to_numpy()
        _project_out(variant_score, annotations_matrix)

        if verbose:
            print("\n\tWriting variant data...")
        cls._write_variant_data(
            method.score_test_hdf5_file_name,
            annotations.select(SPECIAL_COLNAMES),
            jackknife_variant_assignments,
        )

        cls._write_trait_stats(method, variant_score)

    # Compute standard errors using jackknife formula: SE = sqrt((n-1) * var(estimates))
    n_blocks = jackknife_params.shape[0]
    params_se = np.sqrt((n_blocks - 1) * np.var(jackknife_params, axis=0, ddof=1))
    h2_se = np.sqrt((n_blocks - 1) * np.var(jackknife_h2, axis=0, ddof=1))

    # Compute normalized heritability for each jackknife estimate
    jackknife_h2_normalized = jackknife_h2 / jackknife_annot_sums

    # Compute quotient for point estimates and SE
    jackknife_enrichment_quotient = (
        jackknife_h2_normalized / jackknife_h2_normalized[:, [0]]
    )
    enrichment_se = np.sqrt(
        (n_blocks - 1) * np.var(jackknife_enrichment_quotient, axis=0, ddof=1)
    )

    # Compute difference for p-values
    jackknife_enrichment_diff = (
        jackknife_h2_normalized - jackknife_h2_normalized[:, [0]]
    )

    # Two-tailed log10(p-values) using jackknife estimates
    annotation_heritability_p = np.array(
        [
            cls._wald_log10pvalue(jackknife_h2[:, i])
            for i in range(jackknife_h2.shape[1])
        ]
    )
    # Use differences as opposed to quotients for log10(p-values)
    annotation_enrichment_p = np.array(
        [
            cls._wald_log10pvalue(jackknife_enrichment_diff[:, i])
            for i in range(jackknife_enrichment_diff.shape[1])
        ]
    )
    params_p = np.array(
        [
            cls._wald_log10pvalue(jackknife_params[:, i])
            for i in range(jackknife_params.shape[1])
        ]
    )

    likelihood_changes = [
        a - b
        for a, b in zip(
            log_likelihood_history[1:], log_likelihood_history[:-1], strict=False
        )
    ]

    if verbose:
        num_annotations = len(annotation_heritability)
        print(f"Heritability: {annotation_heritability[: min(5, num_annotations)]}")
        print(f"Enrichment: {annotation_enrichment[: min(5, num_annotations)]}")
        print(
            f"Enrichment -log10(p-values): {-annotation_enrichment_p[: min(5, num_annotations)]}"
        )

    return {
        "parameters": params,
        "parameters_se": params_se,
        "parameters_log10pval": params_p,
        "heritability": annotation_heritability,
        "heritability_se": h2_se,
        "heritability_log10pval": annotation_heritability_p,
        "enrichment": annotation_enrichment,
        "enrichment_se": enrichment_se,
        "enrichment_log10pval": annotation_enrichment_p,
        "likelihood_history": log_likelihood_history,
        "jackknife_h2": jackknife_h2,
        "jackknife_params": jackknife_params,
        "jackknife_enrichment": jackknife_enrichment_quotient,
        "variant_h2": variant_h2,
        "log": {
            "converged": converged,
            "num_iterations": rep + 1,
            "likelihood_changes": likelihood_changes,
            "final_likelihood": log_likelihood_history[-1],
            "trust_region_lambdas": trust_region_history,
        },
    }

load_ldgm ¶

load_ldgm(filepath: str, snplist_path: Optional[str] = None, population: Optional[str] = 'EUR', snps_only: bool = False) -> Union['PrecisionOperator', List['PrecisionOperator']]

Load an LDGM from a single LD block's edgelist and snplist files.

Parameters:

Name	Type	Description	Default
`filepath`	`str`	Path to the .edgelist file or directory containing it	required
`snplist_path`	`Optional[str]`	Optional path to .snplist file or directory. If None, uses filepath	`None`
`population`	`Optional[str]`	Optional population name to filter files and set allele frequency column. Defaults to "EUR"	`'EUR'`
`snps_only`	`bool`	Import snplist data for SNPs only (smaller memory usage)	`False`

Returns:

Type	Description
`Union['PrecisionOperator', List['PrecisionOperator']]`	If filepath is a directory: List of PrecisionOperator instances, one for each edgelist file
`Union['PrecisionOperator', List['PrecisionOperator']]`	If filepath is a file: Single PrecisionOperator instance with loaded precision matrix and variant info

Source code in src/graphld/io.py

def load_ldgm(filepath: str, snplist_path: Optional[str] = None, population: Optional[str] = "EUR",
              snps_only: bool = False) -> Union["PrecisionOperator", List["PrecisionOperator"]]:
    """
    Load an LDGM from a single LD block's edgelist and snplist files.

    Args:
        filepath: Path to the .edgelist file or directory containing it
        snplist_path: Optional path to .snplist file or directory. If None, uses filepath
        population: Optional population name to filter files and set allele frequency
            column. Defaults to "EUR"
        snps_only: Import snplist data for SNPs only (smaller memory usage)

    Returns:
        If filepath is a directory:
            List of PrecisionOperator instances, one for each edgelist file
        If filepath is a file:
            Single PrecisionOperator instance with loaded precision matrix and variant info
    """
    from scipy.sparse import csc_matrix

    from .precision import PrecisionOperator

    # Handle directory vs file input
    filepath = Path(filepath)
    if filepath.is_dir():
        pattern = "*.edgelist"
        if population:
            pattern = f"*{population}*.edgelist"
        edgelist_files = list(filepath.glob(pattern))
        if not edgelist_files:
            raise FileNotFoundError(f"No edgelist files found in {filepath}")

        # Load each file and return a list of PrecisionOperators
        operators = []
        for edgelist_file in edgelist_files:
            operator = load_ldgm(edgelist_file, snplist_path, population, snps_only)
            operators.append(operator)
        return operators

    # Use provided snplist path or find corresponding snplist file
    if snplist_path is None:
        snplist_path = filepath.parent
        pattern = filepath.stem.split('.')[0]  # Remove all extensions
        if pattern.endswith(f".{population}"):
            pattern = pattern[:-len(f".{population}")]
        snplist_files = list(Path(snplist_path).glob(f"{pattern}*.snplist"))
        if not snplist_files:
            raise FileNotFoundError(f"No matching snplist file found for {filepath}")
        snplist_file = snplist_files[0]
    else:
        snplist_file = Path(snplist_path)
        if not snplist_file.exists():
            raise FileNotFoundError(f"Snplist file not found: {snplist_file}")

    # Load edgelist data
    edgelist = pl.read_csv(filepath, separator=',', has_header=False,
                          new_columns=['i', 'j', 'value'])

    # Create sparse matrix
    matrix = csc_matrix(
        (edgelist['value'].to_numpy(),
         (edgelist['i'].to_numpy(), edgelist['j'].to_numpy()))
    )

    # Make matrix symmetric
    matrix_t = matrix.T
    diag_vals = matrix.diagonal().copy()
    matrix = matrix + matrix_t
    matrix.setdiag(diag_vals, k=0)

    # Verify diagonal values
    assert np.allclose(matrix.diagonal(), diag_vals), "Diagonal values not set correctly"

    # Create mask for rows/cols with nonzeros on diagonal
    diag = matrix.diagonal()
    nonzero_where = np.where(diag != 0)[0]
    n_nonzero = len(nonzero_where)

    # Load variant info
    variant_info = pl.read_csv(snplist_file, separator=',')
    num_rows = variant_info['index'].max() + 1

    # Create mapping from old indices to new indices
    rows = np.full(num_rows, -1)
    rows[nonzero_where] = np.arange(n_nonzero)

    # If population is specified and exists as a column, rename it to 'af'
    if population and population in variant_info.columns:
        variant_info = variant_info.rename({population: 'af'})
    elif 'af' not in variant_info.columns:
        available_cols = ", ".join(variant_info.columns)
        raise ValueError(
            f"Neither 'af' column nor '{population}' column found in snplist. "
            f"Available columns: {available_cols}"
        )

    # Store original indices and update with new mapping
    variant_info = variant_info.with_columns([
        pl.col('index').alias('original_index'),
        pl.col('index').map_elements(lambda x: rows[x], return_dtype=pl.Int64).alias('index')
    ])

    # Filter out variants with no corresponding matrix row
    variant_info = variant_info.filter(pl.col('index') >= 0)
    if snps_only:
        required_cols = {"anc_alleles", "deriv_alleles"}
        missing_cols = required_cols - set(variant_info.columns)
        if missing_cols:
            raise ValueError(
                "snps_only=True requires snplist columns: "
                + ", ".join(sorted(required_cols))
            )
        variant_info = variant_info.filter(
            (pl.col('anc_alleles').str.len_chars() == 1)
            & (pl.col('deriv_alleles').str.len_chars() == 1)
        )

    # Subset matrix to rows/cols with nonzero diagonal
    matrix = matrix[nonzero_where][:, nonzero_where]

    return PrecisionOperator(matrix, variant_info)

merge_alleles ¶

merge_alleles(anc_alleles: Series, deriv_alleles: Series, ref_alleles: Series, alt_alleles: Series) -> pl.Series

Compare alleles between two sources and return phase information.

Parameters:

Name	Type	Description	Default
`anc_alleles`	`Series`	Ancestral alleles from PrecisionOperator	required
`deriv_alleles`	`Series`	Derived alleles from PrecisionOperator	required
`ref_alleles`	`Series`	Reference alleles from summary statistics	required
`alt_alleles`	`Series`	Alternative alleles from summary statistics	required

Returns:

Type	Description
`Series`	Series of integers indicating phase, where 1 means alleles match exactly,
`Series`	-1 means alleles match but are swapped, and 0 means alleles do not match.

Source code in src/graphld/io.py

def merge_alleles(anc_alleles: pl.Series, deriv_alleles: pl.Series,
                  ref_alleles: pl.Series, alt_alleles: pl.Series) -> pl.Series:
    """Compare alleles between two sources and return phase information.

    Args:
        anc_alleles: Ancestral alleles from PrecisionOperator
        deriv_alleles: Derived alleles from PrecisionOperator
        ref_alleles: Reference alleles from summary statistics
        alt_alleles: Alternative alleles from summary statistics

    Returns:
        Series of integers indicating phase, where 1 means alleles match exactly,
        -1 means alleles match but are swapped, and 0 means alleles do not match.
    """
    # Convert to numpy arrays for faster comparison
    anc = anc_alleles.to_numpy()
    der = deriv_alleles.to_numpy()
    ref = ref_alleles.to_numpy()
    alt = alt_alleles.to_numpy()

    # Make case-insensitive
    anc = np.char.lower(anc.astype(str))
    der = np.char.lower(der.astype(str))
    ref = np.char.lower(ref.astype(str))
    alt = np.char.lower(alt.astype(str))

    # Check matches
    exact_match = (anc == ref) & (der == alt)
    flipped_match = (anc == alt) & (der == ref)

    # Null alleles are given NaN phase so that corresponding Z scores will be NaN
    null_match = (ref_alleles.is_null().to_numpy() & alt_alleles.is_null().to_numpy())

    # Convert to phase
    phase = np.zeros(len(anc), dtype=np.float32)
    phase[exact_match] = 1
    phase[flipped_match] = -1
    phase[null_match] = np.nan

    return pl.Series(phase)

merge_snplists ¶

merge_snplists(precision_op: 'PrecisionOperator', sumstats: DataFrame, *, variant_id_col: str = 'SNP', ref_allele_col: str = 'REF', alt_allele_col: str = 'ALT', match_by_position: bool = False, pos_col: str = 'POS', table_format: str = '', add_cols: list[str] = None, add_allelic_cols: list[str] = None, representatives_only: bool = False, modify_in_place: bool = False) -> Tuple['PrecisionOperator', np.ndarray]

Merge a PrecisionOperator instance with summary statistics DataFrame.

Parameters:

Name	Type	Description	Default
`precision_op`	`'PrecisionOperator'`	PrecisionOperator instance	required
`sumstats`	`DataFrame`	Summary statistics DataFrame	required
`variant_id_col`	`str`	Column name containing variant IDs	`'SNP'`
`ref_allele_col`	`str`	Column name containing reference allele	`'REF'`
`alt_allele_col`	`str`	Column name containing alternative allele	`'ALT'`
`match_by_position`	`bool`	Whether to match SNPs by position instead of ID	`False`
`pos_col`	`str`	Column name containing position	`'POS'`
`table_format`	`str`	Optional file format specification (e.g., 'vcf')	`''`
`add_cols`	`list[str]`	Optional list of column names from sumstats to append to variant_info	`None`
`add_allelic_cols`	`list[str]`	Optional list of column names from sumstats to append to variant_info, multiplied by the phase (-1 or 1) to align with ancestral/derived alleles. If no alleles are provided, these are added without sign-flipping.	`None`
`modify_in_place`	`bool`	Whether to modify the PrecisionOperator in place	`False`

Returns:

Type	Description
`'PrecisionOperator'`	Tuple containing:
`ndarray`	Modified PrecisionOperator with merged variant info and appended columns
`Tuple['PrecisionOperator', ndarray]`	Array of indices into sumstats DataFrame indicating which rows were successfully merged

Source code in src/graphld/io.py

def merge_snplists(precision_op: "PrecisionOperator",
                   sumstats: pl.DataFrame, *,
                   variant_id_col: str = 'SNP',
                   ref_allele_col: str = 'REF',
                   alt_allele_col: str = 'ALT',
                   match_by_position: bool = False,
                   pos_col: str = 'POS',
                   table_format: str = '',
                   add_cols: list[str] = None,
                   add_allelic_cols: list[str] = None,
                   representatives_only: bool = False,
                   modify_in_place: bool = False) -> Tuple["PrecisionOperator", np.ndarray]:
    """Merge a PrecisionOperator instance with summary statistics DataFrame.

    Args:
        precision_op: PrecisionOperator instance
        sumstats: Summary statistics DataFrame
        variant_id_col: Column name containing variant IDs
        ref_allele_col: Column name containing reference allele
        alt_allele_col: Column name containing alternative allele
        match_by_position: Whether to match SNPs by position instead of ID
        pos_col: Column name containing position
        table_format: Optional file format specification (e.g., 'vcf')
        add_cols: Optional list of column names from sumstats to append to variant_info
        add_allelic_cols: Optional list of column names from sumstats to append to variant_info,
            multiplied by the phase (-1 or 1) to align with ancestral/derived alleles.
            If no alleles are provided, these are added without sign-flipping.
        modify_in_place: Whether to modify the PrecisionOperator in place

    Returns:
        Tuple containing:
        - Modified PrecisionOperator with merged variant info and appended columns
        - Array of indices into sumstats DataFrame indicating which rows were successfully merged
    """
    # Handle VCF format
    if table_format.lower() == 'vcf':
        match_by_position = True
        pos_col = 'POS'
        ref_allele_col = 'REF'
        alt_allele_col = 'ALT'
    elif table_format.lower() == 'ldsc':
        match_by_position = False
        ref_allele_col = 'A2'
        alt_allele_col = 'A1'

    # Validate inputs
    if match_by_position:
        pos_options = ['position', 'POS', 'BP']
        if pos_col is not None:
            pos_options.insert(0, pos_col)
        pos_col = next((col for col in pos_options if col in sumstats.columns), None)
        if pos_col is None:
            raise ValueError(
                f"Could not find position column. Tried: {', '.join(pos_options)}"
            )
        if pos_col not in sumstats.columns:
            msg = (f"Summary statistics must contain {pos_col} column "
                  f"for position matching. Found columns: {', '.join(sumstats.columns)}")
            raise ValueError(msg)
    else:
        if variant_id_col not in sumstats.columns:
            msg = (f"Summary statistics must contain {variant_id_col} column. "
                  f"Found columns: {', '.join(sumstats.columns)}")
            raise ValueError(msg)

    # Match variants
    match_by = ('position', pos_col) if match_by_position else ('site_ids', variant_id_col)
    row_nr_col = _temporary_column(
        precision_op.variant_info.columns + sumstats.columns, "row_nr"
    )
    merged = precision_op.variant_info.join(
        sumstats.with_row_index(name=row_nr_col),
        left_on=[match_by[0]],
        right_on=[match_by[1]],
        suffix="_sumstats",
        how='inner'
    )

    # Check alleles if provided
    phase = 1
    if all(col in sumstats.columns for col in [ref_allele_col, alt_allele_col]):
        phase = merge_alleles(
            merged['anc_alleles'],
            merged['deriv_alleles'],
            merged[ref_allele_col],
            merged[alt_allele_col]
        ).alias('phase')
        merged = merged.with_columns(phase)

        # Update indices to only include variants with matching alleles
        merged = merged.filter(pl.col('phase') != 0)
        phase = merged['phase'].to_numpy()

    add_cols = add_cols or []
    add_allelic_cols = add_allelic_cols or []
    new_cols = {}

    # Check all columns exist
    missing_cols = [col for col in add_cols + add_allelic_cols if col not in sumstats.columns]
    if missing_cols:
        msg = (f"Requested columns not found in sumstats: {', '.join(missing_cols)}. "
              f"Available columns: {', '.join(sumstats.columns)}")
        raise ValueError(msg)

    # Add columns with appropriate transformations
    for col in add_cols:
        new_cols[col] = pl.col(col)

    for col in add_allelic_cols:
        new_cols[col] = pl.col(col) * phase

    # Add all new columns at once if any
    if new_cols:
        merged = merged.with_columns(**new_cols)

    # Sort by index and add is_representative column
    merged = (
        merged
        .sort('index')
        .with_columns(
            pl.col('index').is_first_distinct().cast(pl.Int8).alias('is_representative')
        )
    )

    # Create new PrecisionOperator with merged variant info
    unique_indices = np.unique(merged['index'].to_numpy())
    if modify_in_place:
        precision_op.set_which_indices(unique_indices)
    else:
        precision_op = precision_op[unique_indices]

    # Create mapping from old indices to new contiguous ones
    index_map = {old_idx: new_idx for new_idx, old_idx in enumerate(unique_indices)}

    # Update indices in merged data to be contiguous using efficient replace_strict
    merged = merged.with_columns(
        pl.col('index').replace_strict(index_map).alias('index')
    )

    if representatives_only:
        merged = merged.filter(pl.col('is_representative') == 1)

    precision_op.variant_info = merged
    sumstat_indices = merged.select(row_nr_col).to_numpy().flatten().astype(int)

    return precision_op, sumstat_indices

partition_variants ¶

partition_variants(ldgm_metadata: DataFrame, variant_data: DataFrame, *, chrom_col: Optional[str] = None, pos_col: Optional[str] = None) -> List[pl.DataFrame]

Partition variant data according to LDGM blocks.

Parameters:

Name	Type	Description	Default
`ldgm_metadata`	`DataFrame`	DataFrame from read_ldgm_metadata containing block info	required
`variant_data`	`DataFrame`	DataFrame containing variant information	required
`chrom_col`	`Optional[str]`	Optional name of chromosome column. If None, tries common names	`None`
`pos_col`	`Optional[str]`	Optional name of position column. If None, tries common names	`None`

Returns:

Type	Description
`List[DataFrame]`	List of DataFrames, one per row in ldgm_metadata, containing variants
`List[DataFrame]`	that fall within each block's coordinates. Variants within each
`List[DataFrame]`	returned DataFrame are sorted by chromosome and position rather than
`List[DataFrame]`	preserving the input row order.

Raises:

Type	Description
`ValueError`	If the file does not have the expected columns: variant_id, chromosome, position, ...

Source code in src/graphld/io.py

def partition_variants(
    ldgm_metadata: pl.DataFrame,
    variant_data: pl.DataFrame,
    *,
    chrom_col: Optional[str] = None,
    pos_col: Optional[str] = None
) -> List[pl.DataFrame]:
    """Partition variant data according to LDGM blocks.

    Args:
        ldgm_metadata: DataFrame from read_ldgm_metadata containing block info
        variant_data: DataFrame containing variant information
        chrom_col: Optional name of chromosome column. If None, tries common names
        pos_col: Optional name of position column. If None, tries common names

    Returns:
        List of DataFrames, one per row in ldgm_metadata, containing variants
        that fall within each block's coordinates. Variants within each
        returned DataFrame are sorted by chromosome and position rather than
        preserving the input row order.

    Raises:
        ValueError: If the file does not have the expected columns: variant_id,
            chromosome, position, ...
    """
    # Find chromosome column
    chrom_options = ['chrom', 'chromosome', 'CHR']
    if chrom_col is not None:
        chrom_options.insert(0, chrom_col)
    chrom_col = next((col for col in chrom_options if col in variant_data.columns), None)
    if chrom_col is None:
        raise ValueError(
            f"Could not find chromosome column. Tried: {', '.join(chrom_options)}"
        )

    # Find position column
    pos_options = ['position', 'POS', 'BP']
    if pos_col is not None:
        pos_options.insert(0, pos_col)
    pos_col = next((col for col in pos_options if col in variant_data.columns), None)
    if pos_col is None:
        raise ValueError(
            f"Could not find position column. Tried: {', '.join(pos_options)}"
        )

    # Convert chromosome column to integer if needed
    if variant_data[chrom_col].dtype != pl.Int64:
        variant_data = variant_data.with_columns(
            pl.col(chrom_col).cast(pl.Int64).alias(chrom_col)
        )

    # First sort variants by chromosome and position
    sorted_variants = variant_data.sort([chrom_col, pos_col])

    partitioned = []
    chrom_variant_cache = {}
    for block in ldgm_metadata.iter_rows(named=True):
        chrom = block['chrom']
        if chrom not in chrom_variant_cache:
            chrom_variant_cache[chrom] = sorted_variants.filter(pl.col(chrom_col) == chrom)
        chrom_variants = chrom_variant_cache[chrom]
        if len(chrom_variants) == 0:
            partitioned.append(pl.DataFrame())
            continue

        positions = chrom_variants.get_column(pos_col).to_numpy()
        start_idx = np.searchsorted(positions, block['chromStart'])
        end_idx = np.searchsorted(positions, block['chromEnd'])
        partitioned.append(chrom_variants.slice(start_idx, end_idx - start_idx))

    return partitioned

create_ldgm_metadata ¶

create_ldgm_metadata(directory: Union[str, Path], output_file: Optional[str] = None) -> pl.DataFrame

Create metadata file for LDGM files in a directory.

Parameters:

Name	Type	Description	Default
`directory`	`Union[str, Path]`	Directory containing .snplist and .edgelist files	required
`output_file`	`Optional[str]`	Optional path to write CSV file. If None, only returns DataFrame	`None`

Returns:

Type	Description
`DataFrame`	Polars DataFrame containing metadata for each LDGM file

Source code in src/graphld/io.py

def create_ldgm_metadata(
    directory: Union[str, Path], output_file: Optional[str] = None
) -> pl.DataFrame:
    """Create metadata file for LDGM files in a directory.

    Args:
        directory: Directory containing .snplist and .edgelist files
        output_file: Optional path to write CSV file. If None, only returns DataFrame

    Returns:
        Polars DataFrame containing metadata for each LDGM file
    """
    directory = Path(directory)
    if not directory.exists():
        raise FileNotFoundError(f"Directory not found: {directory}")

    # Find all edgelist files
    edgelist_files = list(directory.glob("*.edgelist"))
    if not edgelist_files:
        raise FileNotFoundError(f"No .edgelist files found in {directory}")

    # Process each file
    data = []
    for edgefile in edgelist_files:
        # Parse filename to get info
        name = edgefile.name
        parts = name.split('_')  # e.g. 1kg_chr1_2888443_4320284.EUR.edgelist
        if len(parts) < 4:
            print(f"Skipping {name}: unexpected filename format")
            continue

        # Get chromosome and positions
        try:
            chrom = int(parts[1].replace('chr', ''))
            chromStart = int(parts[2])
            chromEnd = int(parts[3].split('.')[0])  # Remove population/extension
        except (ValueError, IndexError):
            print(f"Skipping {name}: could not parse chromosome/position")
            continue

        # Get population
        try:
            population = name.split('.')[-2]  # Second to last part
        except IndexError:
            print(f"Skipping {name}: could not parse population")
            continue

        # Find corresponding snplist file
        base_name = name.split('.')[0]  # Remove population and extension
        snplist_files = list(directory.glob(f"{base_name}.snplist"))
        if not snplist_files:
            print(f"Skipping {name}: no matching .snplist file")
            continue
        snplist_name = snplist_files[0].name

        # Count variants in snplist
        try:
            snplist_df = pl.read_csv(snplist_files[0])
            num_variants = len(snplist_df)
        except Exception as e:
            print(f"Skipping {name}: error reading snplist: {e}")
            continue

        # Count entries in edgelist
        try:
            edgelist_df = pl.read_csv(edgefile, has_header=False,
                                    new_columns=['i', 'j', 'value'])

            # Count unique diagonal indices
            diag_mask = edgelist_df['i'] == edgelist_df['j']
            num_indices = len(edgelist_df.filter(diag_mask)['i'].unique())

            # Total number of entries
            num_entries = len(edgelist_df)

        except Exception as e:
            print(f"Skipping {name}: error reading edgelist: {e}")
            continue

        # Add row to metadata
        data.append({
            'chrom': chrom,
            'chromStart': chromStart,
            'chromEnd': chromEnd,
            'name': name,
            'snplistName': snplist_name,
            'population': population,
            'numVariants': num_variants,
            'numIndices': num_indices,
            'numEntries': num_entries,
            'info': ''
        })

    # Create DataFrame
    if not data:
        raise ValueError("No valid LDGM files found")

    df = pl.DataFrame(data)

    # Sort by chromosome and start position
    df = df.sort(['chrom', 'chromStart'])

    # Write to file if requested
    if output_file:
        df.write_csv(output_file)

    return df

read_ldgm_metadata ¶

read_ldgm_metadata(filepath: Union[str, Path], *, populations: Optional[Union[str, List[str]]] = None, chromosomes: Optional[Union[int, List[int]]] = None, max_blocks: Optional[int] = None) -> pl.DataFrame

Read LDGM metadata from CSV file.

Parameters:

Name	Type	Description	Default
`filepath`	`Union[str, Path]`	Path to metadata CSV file	required
`populations`	`Optional[Union[str, List[str]]]`	Optional population(s) to filter by	`None`
`chromosomes`	`Optional[Union[int, List[int]]]`	Optional chromosome(s) to filter by	`None`
`max_blocks`	`Optional[int]`	Optional maximum number of blocks to return	`None`

Returns:

Type	Description
`DataFrame`	Polars DataFrame containing LDGM metadata, filtered by population and chromosome
`DataFrame`	if specified, and limited to max_blocks if specified

Source code in src/graphld/io.py

def read_ldgm_metadata(
    filepath: Union[str, Path],
    *,
    populations: Optional[Union[str, List[str]]] = None,
    chromosomes: Optional[Union[int, List[int]]] = None,
    max_blocks: Optional[int] = None
) -> pl.DataFrame:
    """Read LDGM metadata from CSV file.

    Args:
        filepath: Path to metadata CSV file
        populations: Optional population(s) to filter by
        chromosomes: Optional chromosome(s) to filter by
        max_blocks: Optional maximum number of blocks to return

    Returns:
        Polars DataFrame containing LDGM metadata, filtered by population and chromosome
        if specified, and limited to max_blocks if specified
    """
    try:
        df = pl.read_csv(filepath)
        required_cols = [
            'chrom', 'chromStart', 'chromEnd', 'name', 'snplistName',
            'population', 'numVariants', 'numIndices', 'numEntries', 'info'
        ]
        missing = [col for col in required_cols if col not in df.columns]
        if missing:
            raise ValueError(f"Missing required columns: {', '.join(missing)}")

        # Filter by population if specified
        if populations is not None:
            if isinstance(populations, str):
                populations = [populations]
            df = df.filter(pl.col('population').is_in(populations))
            if len(df) == 0:
                raise ValueError(f"No blocks found for populations: {populations}")

        # Filter by chromosome if specified
        if chromosomes is not None:
            if isinstance(chromosomes, int):
                chromosomes = [chromosomes]
            df = df.filter(pl.col('chrom').is_in(chromosomes))
            if len(df) == 0:
                raise ValueError(f"No blocks found for chromosomes: {chromosomes}")

        # Sort by chromosome and position
        df = df.sort(['chrom', 'chromStart'])

        # Limit number of blocks if specified
        if max_blocks is not None:
            df = df.head(max_blocks)

        return df

    except Exception as e:
        raise ValueError(f"Error reading metadata file: {e}") from e

load_annotations ¶

load_annotations(annot_path: str, chromosome: Optional[int] = None, infer_schema_length: int = 100000, add_alleles: bool = False, add_positions: bool = True, positions_file: str = POSITIONS_FILE, file_pattern: str = '*.{chrom}.annot', exclude_bed: bool = False) -> pl.DataFrame

Load annotation data for specified chromosome(s) and merge with LDGMs data.

Parameters:

Name	Type	Description	Default
`annot_path`	`str`	Path to directory containing annotation files	required
`chromosome`	`Optional[int]`	Specific chromosome number, or None for all chromosomes	`None`
`infer_schema_length`	`int`	Number of rows to infer schema from. Runs faster if this is smaller, but will throw an error if too small because floating-point columns will be cast as integers.	`100000`
`file_pattern`	`str`	Filename pattern to match, with {chrom} as a placeholder for chromosome number	`'*.{chrom}.annot'`
`exclude_bed`	`bool`	If True, skip loading .bed files from the annotations directory	`False`

Returns:

Type	Description
`DataFrame`	DataFrame containing annotations

Raises:

Type	Description
`ValueError`	If no matching annotation files are found

Source code in src/graphld/io.py

def load_annotations(annot_path: str,
                    chromosome: Optional[int] = None,
                    infer_schema_length: int = 100_000,
                    add_alleles: bool = False,
                    add_positions: bool = True,
                    positions_file: str = POSITIONS_FILE,
                    file_pattern: str = "*.{chrom}.annot",
                    exclude_bed: bool = False
                    ) -> pl.DataFrame:
    """Load annotation data for specified chromosome(s) and merge with LDGMs data.

    Args:
        annot_path: Path to directory containing annotation files
        chromosome: Specific chromosome number, or None for all chromosomes
        infer_schema_length: Number of rows to infer schema from. Runs faster if this is
            smaller, but will throw an error if too small because floating-point columns
            will be cast as integers.
        file_pattern: Filename pattern to match, with {chrom} as a placeholder for chromosome number
        exclude_bed: If True, skip loading .bed files from the annotations directory

    Returns:
        DataFrame containing annotations

    Raises:
        ValueError: If no matching annotation files are found
    """
    # Determine which chromosomes to process
    if chromosome is not None:
        chromosomes = [chromosome]
    else:
        chromosomes = range(1, 23)  # Assuming chromosomes 1-22

    snplist_data = None

    # Find matching files
    annotations = []
    for chromosome in chromosomes:
        chromosome_pattern = file_pattern.format(chrom=chromosome)
        matching_files = sorted(Path(annot_path).glob(chromosome_pattern))

        # Read all matching files for this chromosome
        dfs = []
        seen_cols = set()
        for file_path in matching_files:
            df = pl.scan_csv(file_path, separator='\t', infer_schema_length=infer_schema_length)
            # Drop columns that were already seen in previous files to avoid duplicates
            schema_names = df.collect_schema().names()
            cols_to_keep = [col for col in schema_names if col not in seen_cols]
            if cols_to_keep:
                df = df.select(cols_to_keep)
                seen_cols.update(cols_to_keep)
                dfs.append(df)

        # Horizontally concatenate all dataframes for this chromosome
        if dfs:
            combined_df = pl.concat(dfs, how="horizontal").collect()
            if add_positions or add_alleles:
                if snplist_data is None:
                    snplist_data = _read_annotation_positions(positions_file, add_alleles)
                combined_df = _attach_annotation_positions(
                    combined_df,
                    snplist_data,
                    chromosome,
                    add_positions=add_positions,
                    add_alleles=add_alleles,
                )
            annotations.append(combined_df)

    # Check if any files were found
    if not annotations:
        raise ValueError(
            f"No annotation files found in {annot_path} matching pattern {file_pattern}"
        )

    # Concatenate all chromosome dataframes and handle different schemas
    annotations = pl.concat(annotations, how="diagonal_relaxed")

    # Convert binary columns to boolean to save memory
    numeric_types = (pl.Int8, pl.Int16, pl.Int32, pl.Int64, pl.Float32, pl.Float64)
    binary_cols = []
    for col in annotations.columns:
        # Skip non-numeric columns
        if not isinstance(annotations[col].dtype, numeric_types):
            continue
        # Check if column only contains 0, 1 and null values
        unique_vals = set(annotations[col].unique().drop_nulls())
        if unique_vals == {0, 1}:
            binary_cols.append(col)

    # Convert binary columns to boolean
    if binary_cols:
        bool_exprs = [pl.col(col).cast(pl.Boolean) for col in binary_cols]
        annotations = annotations.with_columns(bool_exprs)

    if not add_positions and 'BP' in annotations.columns:
        annotations = annotations.rename({'BP': 'POS'})

    bed_files = list_bed_files(annot_path)
    if exclude_bed or not bed_files:
        return annotations

    return add_bed_annotations(annotations, bed_files)

read_concat_snplists ¶

read_concat_snplists(ldgm_metadata: DataFrame, parent_dir: Path) -> pl.LazyFrame

Read and concatenate snplists from LDGM metadata.

Parameters:

Name	Type	Description	Default
`ldgm_metadata`	`DataFrame`	DataFrame from read_ldgm_metadata containing block info	required

Returns:

Type	Description
`LazyFrame`	LazyFrame containing variant information concatenated across blocks.

Source code in src/graphld/io.py

def read_concat_snplists(ldgm_metadata: pl.DataFrame, parent_dir: Path) -> pl.LazyFrame:
    """Read and concatenate snplists from LDGM metadata.

    Args:
        ldgm_metadata: DataFrame from read_ldgm_metadata containing block info

    Returns:
        LazyFrame containing variant information concatenated across blocks.
    """
    ldgms = [
        load_ldgm(parent_dir / Path(row["name"])) for row in ldgm_metadata.iter_rows(named=True)
    ]
    lazy_frames = [
        ldgm.variant_info.lazy().with_columns(pl.lit(i).alias('block'))
        for i, ldgm in enumerate(ldgms)
    ]
    return pl.concat(lazy_frames)

gaussian_likelihood ¶

gaussian_likelihood(pz: ndarray, M: PrecisionOperator) -> float

Compute log-likelihood of GWAS summary statistics under a Gaussian model.

The model is

beta ~ MVN(0, D) z|beta ~ MVN(sqrt(n)Rbeta, R) where R is the LD matrix, n the sample size pz = inv(R) * z / sqrt(n) M = cov(pz) = D + inv(R)/n

Parameters:

Name	Type	Description	Default
`pz`	`ndarray`	Array of precision-premultiplied GWAS effect size estimates	required
`M`	`PrecisionOperator`	PrecisionOperator. This should be the covariance of pz.	required

Returns:

Type	Description
`float`	Log-likelihood value

Source code in src/graphld/likelihood.py

def gaussian_likelihood(
    pz: np.ndarray,
    M: PrecisionOperator,
) -> float:
    """Compute log-likelihood of GWAS summary statistics under a Gaussian model.

    The model is:
        beta ~ MVN(0, D)
        z|beta ~ MVN(sqrt(n)*R*beta, R) where R is the LD matrix, n the sample size
        pz = inv(R) * z / sqrt(n)
        M = cov(pz) = D + inv(R)/n

    Args:
        pz: Array of precision-premultiplied GWAS effect size estimates
        M: PrecisionOperator. This should be the covariance of pz.

    Returns:
        Log-likelihood value

    """
    # Following scipy's convention:
    # log_pdf = -0.5 * (n * log(2π) + log|Σ| + x^T Σ^{-1} x)
    n = len(pz)
    logdet = M.logdet()

    # Compute quadratic form
    b = M.solve(pz)
    quad = np.sum(pz * b)

    # Compute log likelihood
    ll = -0.5 * (n * np.log(2 * np.pi) + logdet + quad)

    return ll

gaussian_likelihood_gradient ¶

gaussian_likelihood_gradient(pz: ndarray, M: PrecisionOperator, del_M_del_a: Optional[ndarray] = None, n_samples: int = 10, seed: Optional[int] = None, trace_estimator: Optional[str] = 'xdiag') -> np.ndarray

Computes the score under a Gaussian model.

The model is:
beta ~ MVN(0, D)
z|beta ~ MVN(sqrt(n)*R*beta, R) where R is the LD matrix, n the sample size
pz = inv(R) * z / sqrt(n)
M = cov(pz) = D + inv(R)/n

Parameters:

Name	Type	Description	Default
`pz`	`ndarray`	Array of precision-premultiplied GWAS effect size estimates	required
`M`	`PrecisionOperator`	PrecisionOperator. This should be the covariance of pz.	required
`del_M_del_a`	`Optional[ndarray]`	Matrix of derivatives of M's diagonal elements wrt parameters a	`None`
`n_samples`	`int`	Number of probe vectors for Hutchinson's method or xdiag	`10`
`seed`	`Optional[int]`	Random seed for generating probe vectors	`None`
`trace_estimator`	`Optional[str]`	Method for computing the trace estimator. Options: "exact", "hutchinson", "xdiag"	`'xdiag'`

Returns:

Type	Description
`ndarray`	Array of diagonal elements of the gradient wrt M's diagonal elements,
`ndarray`	or with respect to parameters a if del_M_del_a is provided

Source code in src/graphld/likelihood.py

def gaussian_likelihood_gradient(
    pz: np.ndarray,
    M: PrecisionOperator,
    del_M_del_a: Optional[np.ndarray] = None,
    n_samples: int = 10,
    seed: Optional[int] = None,
    trace_estimator: Optional[str] = "xdiag",
) -> np.ndarray:
    """Computes the score under a Gaussian model.

        The model is:
        beta ~ MVN(0, D)
        z|beta ~ MVN(sqrt(n)*R*beta, R) where R is the LD matrix, n the sample size
        pz = inv(R) * z / sqrt(n)
        M = cov(pz) = D + inv(R)/n

    Args:
        pz: Array of precision-premultiplied GWAS effect size estimates
        M: PrecisionOperator. This should be the covariance of pz.
        del_M_del_a: Matrix of derivatives of M's diagonal elements wrt parameters a
        n_samples: Number of probe vectors for Hutchinson's method or xdiag
        seed: Random seed for generating probe vectors
        trace_estimator: Method for computing the trace estimator.
            Options: "exact", "hutchinson", "xdiag"

    Returns:
        Array of diagonal elements of the gradient wrt M's diagonal elements,
        or with respect to parameters a if del_M_del_a is provided
    """
    # Compute b = M^(-1) * pz
    b = M.solve(pz)

    # Compute diagonal elements of M^(-1)
    minv_diag = M.inverse_diagonal(method=trace_estimator,
                                    n_samples=n_samples,
                                    seed=seed)

    # Compute gradient diagonal elements
    node_grad = -0.5 * (minv_diag.flatten() - b.flatten()**2)

    if del_M_del_a is None:
        return node_grad
    with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
        return node_grad @ del_M_del_a

gaussian_likelihood_hessian ¶

gaussian_likelihood_hessian(pz: ndarray, M: PrecisionOperator, del_M_del_a: Optional[ndarray] = None, trace_estimator: str = _TRACE_ESTIMATOR_DEFAULT, n_samples: int = 100, seed: Optional[int] = None, *, diagonal_method: Optional[str] = None) -> np.ndarray

Computes the average information matrix of the Gaussian log-likelihood.

The model is

beta ~ MVN(0, D) z|beta ~ MVN(sqrt(n)Rbeta, R) where R is the LD matrix, n the sample size pz = inv(R) * z / sqrt(n) M = cov(pz) = D + inv(R)/n

Parameters:

Name	Type	Description	Default
`pz`	`ndarray`	Array of precision-premultiplied GWAS effect size estimates	required
`M`	`PrecisionOperator`	PrecisionOperator. This should be the covariance of pz.	required
`del_M_del_a`	`Optional[ndarray]`	Matrix of derivatives of M's diagonal elements wrt parameters a. If None, only the diagonal elements are computed.	`None`
`trace_estimator`	`str`	Method for computing the diagonal trace estimator when del_M_del_a is None. Options accepted by PrecisionOperator include "exact", "hutchinson", and "xdiag". Defaults to "xdiag".	`_TRACE_ESTIMATOR_DEFAULT`
`n_samples`	`int`	Number of probe vectors for Hutchinson's method or xdiag	`100`
`seed`	`Optional[int]`	Random seed for generating probe vectors	`None`
`diagonal_method`	`Optional[str]`	Deprecated alias for trace_estimator.	`None`

Returns:

Type	Description
`ndarray`	Matrix of second derivatives wrt parameters a, or array of diagonal elements
`ndarray`	if del_M_del_a is None

Source code in src/graphld/likelihood.py

def gaussian_likelihood_hessian(
    pz: np.ndarray,
    M: PrecisionOperator,
    del_M_del_a: Optional[np.ndarray] = None,
    trace_estimator: str = _TRACE_ESTIMATOR_DEFAULT,
    n_samples: int = 100,
    seed: Optional[int] = None,
    *,
    diagonal_method: Optional[str] = None,
) -> np.ndarray:
    """Computes the average information matrix of the Gaussian log-likelihood.

    The model is:
        beta ~ MVN(0, D)
        z|beta ~ MVN(sqrt(n)*R*beta, R) where R is the LD matrix, n the sample size
        pz = inv(R) * z / sqrt(n)
        M = cov(pz) = D + inv(R)/n

    Args:
        pz: Array of precision-premultiplied GWAS effect size estimates
        M: PrecisionOperator. This should be the covariance of pz.
        del_M_del_a: Matrix of derivatives of M's diagonal elements wrt parameters a.
            If None, only the diagonal elements are computed.
        trace_estimator: Method for computing the diagonal trace estimator when
            del_M_del_a is None. Options accepted by PrecisionOperator include
            "exact", "hutchinson", and "xdiag". Defaults to "xdiag".
        n_samples: Number of probe vectors for Hutchinson's method or xdiag
        seed: Random seed for generating probe vectors
        diagonal_method: Deprecated alias for trace_estimator.

    Returns:
        Matrix of second derivatives wrt parameters a, or array of diagonal elements
        if del_M_del_a is None
    """
    trace_estimator_was_supplied = trace_estimator is not _TRACE_ESTIMATOR_DEFAULT
    if diagonal_method is not None:
        if trace_estimator_was_supplied and trace_estimator != diagonal_method:
            raise ValueError(
                "trace_estimator and diagonal_method specify different Hessian "
                "diagonal estimators; use trace_estimator only."
            )
        warnings.warn(
            "diagonal_method is deprecated; use trace_estimator instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        trace_estimator = diagonal_method
    elif not trace_estimator_was_supplied:
        trace_estimator = "xdiag"

    # Compute b = M^(-1) * pz
    b = M.solve(pz)
    if b.ndim == 1:
        b = b.reshape(-1, 1)

    # If del_M_del_a is None, compute only the diagonal of the Hessian
    if del_M_del_a is None:
        minv_diag = M.inverse_diagonal(
            method=trace_estimator, n_samples=n_samples, seed=seed
        )
        hess_diag = -0.5 * minv_diag.flatten() * b.flatten()**2
        return hess_diag

    # Compute b_scaled = b .* del_sigma_del_a
    b_scaled = b * del_M_del_a

    # Compute M^(-1) * b_scaled
    minv_b_scaled = M.solve(b_scaled)

    # Compute Hessian: -1/2 * b_scaled^T * M^(-1) * b_scaled
    with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
        hess = -0.5 * (b_scaled.T @ minv_b_scaled)

    return hess

softmax_robust ¶

softmax_robust(x: ndarray) -> np.ndarray

Numerically stable softplus-like link implementation.

Source code in src/graphld/heritability.py

def softmax_robust(x: np.ndarray) -> np.ndarray:
    """Numerically stable softplus-like link implementation."""
    with np.errstate(over="ignore"):
        y = x + np.log1p(np.exp(-x))
        mask = x < 0
        y[mask] = np.log1p(np.exp(x[mask]))
    return y

run_graphREML ¶

run_graphREML(model_options: ModelOptions, method_options: MethodOptions, summary_stats: DataFrame, annotation_data: DataFrame, ldgm_metadata_path: str, populations: Union[str, List[str]] = None, chromosomes: Optional[Union[int, List[int]]] = None) -> dict[str, Any]

Run GraphREML heritability partitioning.

Parameters:

Name	Type	Description	Default
`model_options`	`ModelOptions`	Model options containing annotations and parameterization.	required
`method_options`	`MethodOptions`	Method options controlling optimization and processing.	required
`summary_stats`	`DataFrame`	DataFrame containing GWAS summary statistics	required
`annotation_data`	`DataFrame`	DataFrame containing variant annotations	required
`ldgm_metadata_path`	`str`	Path to LDGM metadata file	required
`populations`	`Union[str, List[str]]`	Optional list of populations to include	`None`
`chromosomes`	`Optional[Union[int, List[int]]]`	Optional list of chromosomes to include	`None`

Returns:

Name	Type	Description
`dict`	`dict[str, Any]`	Estimated parameters, heritability estimates, standard errors, and convergence diagnostics.

Source code in src/graphld/heritability.py

def run_graphREML(
    model_options: ModelOptions,
    method_options: MethodOptions,
    summary_stats: pl.DataFrame,
    annotation_data: pl.DataFrame,
    ldgm_metadata_path: str,
    populations: Union[str, List[str]] = None,
    chromosomes: Optional[Union[int, List[int]]] = None,
) -> dict[str, Any]:
    """Run GraphREML heritability partitioning.

    Args:
        model_options: Model options containing annotations and parameterization.
        method_options: Method options controlling optimization and processing.
        summary_stats: DataFrame containing GWAS summary statistics
        annotation_data: DataFrame containing variant annotations
        ldgm_metadata_path: Path to LDGM metadata file
        populations: Optional list of populations to include
        chromosomes: Optional list of chromosomes to include

    Returns:
        dict: Estimated parameters, heritability estimates, standard errors,
            and convergence diagnostics.
    """
    if populations is None:
        raise ValueError("Populations must be provided")

    if model_options.binary_annotations_only:
        original_annotation_columns = list(model_options.annotation_columns)
        annotation_data, model_options.annotation_columns = _filter_binary_annotations(
            annotation_data, model_options.annotation_columns, method_options.verbose
        )
        _subset_model_options_to_annotations(
            model_options, original_annotation_columns, model_options.annotation_columns
        )
    else:
        _validate_model_params(model_options)

    # Merge summary stats with annotations
    merge_how = "right" if method_options.use_surrogate_markers else "inner"
    if not method_options.match_by_position:
        join_cols = ["SNP"]
        merged_data = summary_stats.join(annotation_data, on=join_cols, how=merge_how)

        # Handle column renaming - ensure we have consistent CHR and POS columns
        if "CHR_right" in merged_data.columns and "POS_right" in merged_data.columns:
            # Drop existing CHR and POS columns if they exist to avoid duplicates
            merged_data = (
                merged_data.drop("CHR")
                .drop("POS")
                .rename({"CHR_right": "CHR", "POS_right": "POS"})
            )

    else:
        join_cols = ["CHR", "POS"]
        merged_data = summary_stats.join(annotation_data, on=join_cols, how=merge_how)

    # Deduplicate chr/pos pairs with multiple entries in the summary statistics
    merged_data = merged_data.unique(subset=join_cols, keep="first")

    if merged_data.is_empty():
        raise ValueError(
            "No overlapping variants found between summary statistics and annotations."
        )

    # Print shape of each DataFrame
    if method_options.verbose:
        print(f"Summary stats shape: {summary_stats.shape}")
        print(f"Annotation data shape: {annotation_data.shape}")
        print(f"Merged data shape: {merged_data.shape}")
        mean_chisq = (merged_data["Z"] ** 2).mean()
        print(f"Mean chisq: {mean_chisq}")
        max_chisq = (merged_data["Z"] ** 2).max()
        print(f"Max chisq: {max_chisq}")

    _infer_or_default_sample_size(model_options, merged_data, method_options.verbose)

    run_fn = GraphREML.run_serial if method_options.run_serial else GraphREML.run
    return run_fn(
        ldgm_metadata_path,
        populations=populations,
        chromosomes=chromosomes,
        sumstats=merged_data,
        num_processes=method_options.num_processes,
        worker_params=(model_options, method_options),
        num_params=len(model_options.annotation_columns),
        model=model_options,
        method=method_options,
        num_iterations=method_options.num_iterations,
        verbose=method_options.verbose,
        convergence_tol=method_options.convergence_tol,
        sample_size=model_options.sample_size,
    )

graphld.heritability¶

heritability ¶

ModelOptions dataclass ¶

MethodOptions dataclass ¶

GraphREML ¶

prepare_block_data classmethod ¶

create_shared_memory staticmethod ¶

process_block classmethod ¶

supervise classmethod ¶

load_ldgm ¶

merge_alleles ¶

merge_snplists ¶

partition_variants ¶

create_ldgm_metadata ¶

read_ldgm_metadata ¶

load_annotations ¶

read_concat_snplists ¶

gaussian_likelihood ¶

gaussian_likelihood_gradient ¶

gaussian_likelihood_hessian ¶

softmax_robust ¶

run_graphREML ¶

ModelOptions `dataclass` ¶

MethodOptions `dataclass` ¶

prepare_block_data `classmethod` ¶

create_shared_memory `staticmethod` ¶

process_block `classmethod` ¶

supervise `classmethod` ¶