score_test

Enrichment score test module for fast annotation testing.

Submodules

• score_test.score_test: Core score test implementation
• score_test.score_test_io: I/O functions for score test files
• score_test.meta_analysis: Meta-analysis across traits
• score_test.convert_scores: Convert between variant and gene-level scores
• score_test.genesets: Gene set utilities

score_test

GraphLD Score Test - Fast heritability enrichment testing for genomic annotations.

This module provides a fast score test for testing genomic or gene annotations for heritability enrichment conditional upon a null model. The test produces Z scores where positive values indicate enrichment and negative values indicate depletion.

Main Functions

run_score_test: Run the enrichment score test on annotations load_variant_data: Load precomputed row-level score statistics load_trait_data: Load trait-specific score statistics load_annotations: Load annotation data from LDSC format files load_gene_table: Load gene position table load_gene_sets_from_gmt: Load gene sets from GMT format files gene_variant_matrix: Build variant-to-gene nearest-gene weight matrices convert_gene_to_variant_annotations: Convert gene-level to variant-level annotations

Example::

from score_test import load_trait_data, load_variant_data, run_score_test
from score_test.score_test_io import load_variant_annotations

row_data = load_variant_data("path/to/scores.h5")
trait_data = load_trait_data("path/to/scores.h5", "height", row_data)
annot = load_variant_annotations("path/to/annot_dir/", variant_table=row_data)
results = run_score_test(trait_data, annot)

run_score_test

run_score_test(trait_data: TraitData, annot: Annot) -> Tuple[np.ndarray, np.ndarray]

Run approximate score test for hypothesis testing of new annotation or functional category.

This function tests annotations against precomputed gradient statistics. It does not adjust for uncertainty in fitted model parameters.

Parameters:

Name	Type	Description	Default
trait_data	TraitData	TraitData object containing variant data with gradients	required
annot	Annot	Annot object (VariantAnnot or GeneAnnot) containing annotations to test	required

Returns:

Type	Description
Tuple[ndarray, ndarray]	Tuple of (point_estimates, jackknife_estimates)

Source code in src/score_test/score_test.py

def run_score_test(trait_data: TraitData,
    annot: Annot,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Run approximate score test for hypothesis testing of new annotation or functional category.

    This function tests annotations against precomputed gradient statistics. It
    does not adjust for uncertainty in fitted model parameters.

    Args:
        trait_data: TraitData object containing variant data with gradients
        annot: Annot object (VariantAnnot or GeneAnnot) containing annotations to test

    Returns:
        Tuple of (point_estimates, jackknife_estimates)
    """
    # Merge trait data with annotations
    grad, _, _, test_annot, block_boundaries = annot.merge(trait_data)

    noBlocks = len(block_boundaries) - 1

    # Compute single-block derivatives
    U_block = [] # Derivative of the log-likelihood for each test annotation
    with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
        for i in range(noBlocks):
            block = range(block_boundaries[i], block_boundaries[i+1])
            Ui = grad[block].reshape(1,-1) @ test_annot[block, :]
            U_block.append(Ui)
    U_total = sum(U_block)

    # Compute leave-one-out derivatives
    U_jackknife = np.zeros((noBlocks, U_total.shape[1]))
    for i in range(noBlocks):
        block = range(block_boundaries[i], block_boundaries[i+1])

        # deduct the score contributed from one LD block
        U_jackknife[i, :] = U_total - U_block[i].ravel()

    return U_total, U_jackknife

load_trait_data

load_trait_data(hdf5_path: str, trait_name: str, variant_table: DataFrame) -> TraitData

Load trait data and combine with variant data into a TraitData object.

Parameters:

Name	Type	Description	Default
hdf5_path	str	Path to HDF5 file containing variant and trait data	required
trait_name	str	Name of the trait to load	required
variant_table	DataFrame	Pre-loaded variant table DataFrame	required

Returns:

Type	Description
TraitData	TraitData object containing variant dataframe with gradients/corrections and parameters

Source code in src/score_test/score_test_io.py

def load_trait_data(
    hdf5_path: str, trait_name: str, variant_table: pl.DataFrame
) -> "TraitData":
    """
    Load trait data and combine with variant data into a TraitData object.

    Args:
        hdf5_path: Path to HDF5 file containing variant and trait data
        trait_name: Name of the trait to load
        variant_table: Pre-loaded variant table DataFrame

    Returns:
        TraitData object containing variant dataframe with gradients/corrections and parameters
    """
    # Import at runtime to avoid circular import
    try:
        from .score_test import TraitData
    except ImportError:
        from score_test import TraitData

    trait_hdf5 = load_trait_hdf5(hdf5_path, trait_name)

    # Determine primary key column
    possible_keys = ["RSID", "gene_id", "gene_name", "CHR", "POS"]
    keys = [key for key in possible_keys if key in variant_table.columns]
    if len(keys) == 0:
        raise ValueError("variant_table must have one of: " + ", ".join(possible_keys))

    # Add trait-specific columns (gradient, hessian, etc.)
    new_columns = []
    for key_name, value in trait_hdf5.items():
        if isinstance(value, np.ndarray):
            decoded_value = _decode_bytes_array(value)
            new_columns.append(pl.Series(name=key_name, values=decoded_value))

    df = variant_table.with_columns(new_columns)

    if "parameters" in trait_hdf5:
        params = trait_hdf5["parameters"]["parameters"]
        jk_params = trait_hdf5["parameters"]["jackknife_parameters"]
    else:
        params, jk_params = None, None

    # Create TraitData - it will compute exclude_cols and annot_names via properties
    trait_data = TraitData(
        df=df,
        params=params,
        jk_params=jk_params,
        keys=keys,
    )

    # Compute annotation names using the property
    trait_data.annot_names = [
        col for col in df.columns if col not in trait_data.exclude_cols
    ]

    return trait_data

load_annotations

load_annotations(annot_path: str, chromosome: Optional[int] = None, infer_schema_length: int = 100000, add_positions: bool = True, variant_table: DataFrame | None = None, exclude_bed: bool = False) -> pl.DataFrame

Load annotation data for specified chromosome(s).

Parameters:

Name	Type	Description	Default
annot_path	str	Path to directory containing .annot and/or .bed 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	100000
add_positions	bool	If True, rename BP to POS	True
variant_table	DataFrame | None	Variant table to use as the coordinate universe for BED-only annotation directories	None
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/score_test/score_test_io.py

def load_annotations(
    annot_path: str,
    chromosome: Optional[int] = None,
    infer_schema_length: int = 100_000,
    add_positions: bool = True,
    variant_table: pl.DataFrame | None = None,
    exclude_bed: bool = False,
) -> pl.DataFrame:
    """Load annotation data for specified chromosome(s).

    Args:
        annot_path: Path to directory containing .annot and/or .bed files
        chromosome: Specific chromosome number, or None for all chromosomes
        infer_schema_length: Number of rows to infer schema from
        add_positions: If True, rename BP to POS
        variant_table: Variant table to use as the coordinate universe for
            BED-only annotation directories
        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
    """
    bed_files = list_bed_files(annot_path)

    # Determine which chromosomes to process
    if chromosome is not None:
        chromosomes = [chromosome]
    else:
        chromosomes = range(1, 23)  # Assuming chromosomes 1-22

    # Find matching files
    annotations = []
    for chromosome in chromosomes:
        file_pattern = f"*.{chromosome}.annot"
        matching_files = sorted(Path(annot_path).glob(file_pattern))

        # Read all matching files for this chromosome
        dfs = []
        cols_found = set()
        for file_path in matching_files:
            df = pl.scan_csv(
                file_path, separator="\t", infer_schema_length=infer_schema_length
            )
            schema_names = df.collect_schema().names()
            cols_to_keep = [col for col in schema_names if col not in cols_found]
            if cols_to_keep:
                df = df.select(cols_to_keep)
                cols_found.update(cols_to_keep)
                dfs.append(df)

        # Horizontally concatenate all dataframes for this chromosome
        if dfs:
            combined_df = pl.concat(dfs, how="horizontal")
            combined_df = combined_df.select(sorted(combined_df.collect_schema().names()))
            annotations.append(combined_df)

    if not annotations:
        if bed_files and not exclude_bed:
            annotations = _variant_table_as_annotation_table(variant_table, chromosomes)
        else:
            raise ValueError(
                f"No annotation files found in {annot_path} matching "
                "pattern *.{chrom}.annot or *.bed"
            )

    if isinstance(annotations, list):
        # Concatenate all chromosome dataframes vertically
        annotations = pl.concat(annotations, how="vertical").collect()

    if not isinstance(annotations, pl.DataFrame):
        raise ValueError(
            f"No annotation files found in {annot_path} matching pattern *.{{chrom}}.annot"
        )

    # 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 add_positions:
        annotations = annotations.rename({"BP": "POS"})

    if exclude_bed or not bed_files:
        return annotations

    position_col = "POS" if add_positions else "BP"
    return add_bed_annotations(annotations, bed_files, position_col=position_col)

convert_gene_to_variant_annotations

convert_gene_to_variant_annotations(gene_annot: object, variant_table: DataFrame, gene_table: DataFrame, nearest_weights: ndarray) -> object

Convert gene annotations to variant-level annotations.

Can accept either: - GeneAnnot object with gene_sets and annot_names attributes - dict[str, list[str]] mapping gene set names to gene lists

Parameters:

Name	Type	Description	Default
gene_annot	object	GeneAnnot object or dict mapping set names to lists of genes (symbols or IDs)	required
variant_table	DataFrame	Variant table DataFrame with CHR, POS, RSID columns	required
gene_table	DataFrame	Gene table DataFrame	required
nearest_weights	ndarray	Weights for k-nearest genes	required

Returns:

Type	Description
object	VariantAnnot object with variant-level annotations, if gene_annot is GeneAnnot.
object	DataFrame with variant-level annotations in LDSC format, if gene_annot is dict.

Source code in src/score_test/genesets.py

def convert_gene_to_variant_annotations(gene_annot: object,
                                        variant_table: pl.DataFrame,
                                        gene_table: pl.DataFrame,
                                        nearest_weights: np.ndarray) -> object:
    """Convert gene annotations to variant-level annotations.

    Can accept either:
    - GeneAnnot object with gene_sets and annot_names attributes
    - dict[str, list[str]] mapping gene set names to gene lists

    Args:
        gene_annot: GeneAnnot object or dict mapping set names to lists of genes (symbols or IDs)
        variant_table: Variant table DataFrame with CHR, POS, RSID columns
        gene_table: Gene table DataFrame
        nearest_weights: Weights for k-nearest genes

    Returns:
        VariantAnnot object with variant-level annotations, if gene_annot is GeneAnnot.
        DataFrame with variant-level annotations in LDSC format, if gene_annot is dict.
    """
    # Import at runtime to avoid circular import
    try:
        from .score_test import VariantAnnot
    except ImportError:
        from score_test import VariantAnnot

    # Handle both GeneAnnot objects and plain dicts
    if hasattr(gene_annot, 'gene_sets'):
        gene_sets = gene_annot.gene_sets
        annot_names = gene_annot.annot_names
        return_variant_annot = True
    else:
        gene_sets = gene_annot
        annot_names = None
        return_variant_annot = False

    df_annot = gene_sets_to_variant_annotation_frame(
        gene_sets,
        variant_table,
        gene_table,
        nearest_weights,
        variant_id_col="RSID",
        output_id_col="RSID",
    )

    if return_variant_annot:
        return VariantAnnot(df_annot, annot_names)
    else:
        return df_annot

gene_variant_matrix

gene_variant_matrix(variant_table: DataFrame, gene_table: DataFrame, nearest_weights: ndarray, dtype: dtype = np.float64) -> csr_matrix

Compute a chromosome-aware variants x genes weighted matrix.

Parameters:

Name	Type	Description	Default
variant_table	DataFrame	DataFrame with CHR, POS columns	required
gene_table	DataFrame	DataFrame with CHR, POS columns (POS is midpoint for genes)	required
nearest_weights	ndarray	Weights for k-nearest genes	required
dtype	dtype	Data type of the sparse matrix values	float64

Returns:

Type	Description
csr_matrix	Sparse matrix mapping variants to genes with shape (nvariants, ngenes)

Source code in src/score_test/genesets.py

def gene_variant_matrix(
    variant_table: pl.DataFrame,
    gene_table: pl.DataFrame,
    nearest_weights: np.ndarray,
    dtype: np.dtype = np.float64,
) -> csr_matrix:
    """Compute a chromosome-aware variants x genes weighted matrix.

    Args:
        variant_table: DataFrame with CHR, POS columns
        gene_table: DataFrame with CHR, POS columns (POS is midpoint for genes)
        nearest_weights: Weights for k-nearest genes
        dtype: Data type of the sparse matrix values

    Returns:
        Sparse matrix mapping variants to genes with shape (nvariants, ngenes)
    """
    nvar, ngene = len(variant_table), len(gene_table)
    num_nearest = len(nearest_weights)
    if num_nearest <= 0 or num_nearest > ngene:
        raise ValueError("nearest_weights length must be between 1 and the number of genes")

    nearest = _get_chromosome_aware_nearest_genes(variant_table, gene_table, num_nearest)
    rows = np.repeat(np.arange(nvar, dtype=np.int32), num_nearest)
    cols = nearest.ravel()
    data = np.tile(nearest_weights, nvar)
    return csr_matrix((data, (rows, cols)), shape=(nvar, ngene), dtype=dtype)

score_test

TraitData dataclass

TraitData(df: DataFrame, params: ndarray | None = None, jk_params: ndarray | None = None, annot_names: list[str] | None = None, keys: list[str] | None = None)

exclude_cols property

exclude_cols: set[str]

Columns to exclude when determining annotation names.

Annot

Annot(annot_names: List[str], other_key: str | list[str])

Base class for annotations that can be merged with TraitData.

Parameters:

Name	Type	Description	Default
annot_names	List[str]	List of annotation column names to test	required
other_key	str | list[str]	Column name(s) to use for merging (e.g., 'RSID', 'gene_id', or ['CHR', 'POS'])	required

Source code in src/score_test/score_test.py

def __init__(self, annot_names: List[str], other_key: str | list[str]):
    """
    Args:
        annot_names: List of annotation column names to test
        other_key: Column name(s) to use for merging (e.g., 'RSID', 'gene_id',
            or ['CHR', 'POS'])
    """
    self.annot_names = annot_names
    self.other_key = other_key

merge

merge(trait_data: TraitData) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None, np.ndarray, np.ndarray]

Merge with TraitData and return extracted arrays.

Returns:

Type	Description
ndarray	Tuple of (grad, hessian, model_annot, test_annot, block_boundaries).
ndarray | None	hessian and model_annot can be None when unavailable.

Source code in src/score_test/score_test.py

def merge(
    self, trait_data: TraitData
) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None, np.ndarray, np.ndarray]:
    """Merge with TraitData and return extracted arrays.

    Returns:
        Tuple of (grad, hessian, model_annot, test_annot, block_boundaries).
        hessian and model_annot can be None when unavailable.
    """
    raise NotImplementedError("Subclasses must implement merge()")

VariantAnnot

VariantAnnot(df: DataFrame, annot_names: List[str], other_key: str | list[str] = 'RSID')

Bases: Annot

Variant-level annotations.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with variant annotations	required
annot_names	List[str]	List of annotation column names to test	required
other_key	str | list[str]	Column name(s) to use for merging with trait data	'RSID'

Source code in src/score_test/score_test.py

def __init__(
    self,
    df: pl.DataFrame,
    annot_names: List[str],
    other_key: str | list[str] = 'RSID',
):
    """
    Args:
        df: DataFrame with variant annotations
        annot_names: List of annotation column names to test
        other_key: Column name(s) to use for merging with trait data
    """
    super().__init__(annot_names, other_key=other_key)
    self.df = df

perturb

perturb(fraction: float, seed: int | None = None)

Perturb binary annotations.

Source code in src/score_test/score_test.py

def perturb(self, fraction: float, seed: int | None = None):
    """Perturb binary annotations."""

    def _perturb_binary_vector(vals: np.ndarray, fraction: float, rng: np.random.Generator) -> np.ndarray:
        p = np.mean(vals)
        assert 0 <= p <= 1, f"p must be between 0 and 1, got {p}"
        perturb_mask = (rng.random(len(vals)) < fraction)
        vals[perturb_mask] = rng.binomial(1, p, np.sum(perturb_mask)).astype(vals.dtype)
        return vals

    rng = np.random.default_rng(seed)
    new_cols = []
    kept_names = []
    for col in self.annot_names:
        vals = self.df[col].to_numpy().copy()

        if np.any(np.isnan(vals)):
            raise ValueError(f"Annotation '{col}' contains NaNs")

        unique = np.unique(vals)
        if np.all(np.isin(unique, [0, 1])):
            vals = _perturb_binary_vector(vals, fraction, rng)
            new_cols.append(pl.Series(col, vals))
            kept_names.append(col)

    if new_cols:
        self.df = self.df.with_columns(new_cols)

    self.annot_names = kept_names

merge

merge(trait_data: TraitData) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None, np.ndarray, np.ndarray]

Merge variant annotations with TraitData.

Returns:

Type	Description
ndarray	Tuple of (grad, hessian, model_annot, test_annot, block_boundaries).
ndarray | None	hessian is None if not available; model_annot is None if the
ndarray | None	fitted model has no conditioning annotations.

Source code in src/score_test/score_test.py

def merge(
    self, trait_data: TraitData
) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None, np.ndarray, np.ndarray]:
    """Merge variant annotations with TraitData.

    Returns:
        Tuple of (grad, hessian, model_annot, test_annot, block_boundaries).
        hessian is None if not available; model_annot is None if the
        fitted model has no conditioning annotations.
    """
    # Check if trait_data has the required key
    merge_keys = [self.other_key] if isinstance(self.other_key, str) else self.other_key
    missing_trait_keys = [key for key in merge_keys if key not in trait_data.keys]
    if missing_trait_keys:
        raise ValueError(
            f"TraitData does not have required key(s) {missing_trait_keys}. "
            f"Available keys: {trait_data.keys}"
        )

    # Verify merge key exists in annotation DataFrame
    missing_annot_keys = [key for key in merge_keys if key not in self.df.columns]
    if missing_annot_keys:
        raise ValueError(
            f"Merge key(s) {missing_annot_keys} not found in annotation DataFrame. "
            f"Available columns: {self.df.columns}"
        )

    df_merged = trait_data.df.join(
        self.df,
        left_on=merge_keys,
        right_on=merge_keys,
        how='inner',
        maintain_order='left'
    )
    block_boundaries = get_block_boundaries(df_merged['jackknife_blocks'].to_numpy())

    # Extract arrays from merged dataframe
    grad = df_merged['gradient'].to_numpy().astype(np.float64)
    hessian = df_merged['hessian'].to_numpy().astype(np.float64) if 'hessian' in df_merged.columns else None
    model_annot = df_merged[trait_data.annot_names].to_numpy().astype(np.float64) if trait_data.annot_names else None
    test_annot = df_merged[self.annot_names].to_numpy().astype(np.float64)

    return grad, hessian, model_annot, test_annot, block_boundaries

GeneAnnot

GeneAnnot(gene_sets: dict[str, list[str]])

Bases: Annot

Gene-level annotations.

Parameters:

Name	Type	Description	Default
gene_sets	dict[str, list[str]]	Dictionary mapping set names to lists of genes	required

Source code in src/score_test/score_test.py

def __init__(self, gene_sets: dict[str, list[str]]):
    """
    Args:
        gene_sets: Dictionary mapping set names to lists of genes
    """
    self.gene_sets = gene_sets

    # Determine gene key from first gene in first set
    # Import helper function
    try:
        from .genesets import _is_gene_id
    except ImportError:
        from genesets import _is_gene_id

    first_gene = next(iter(next(iter(gene_sets.values()))))
    self.other_key = 'gene_id' if _is_gene_id(first_gene) else 'gene_name'

    super().__init__(list(gene_sets.keys()), self.other_key)

merge

merge(trait_data: TraitData) -> tuple[np.ndarray, None, None, np.ndarray, np.ndarray]

Merge gene annotations with gene-level TraitData.

Creates one-hot encodings for each gene set and returns unmodified gradients from the TraitData.

Returns:

Type	Description
ndarray	Tuple of (grad, hessian, model_annot, test_annot, block_boundaries).
None	hessian and model_annot are currently None for gene-set tests.

Source code in src/score_test/score_test.py

def merge(
    self, trait_data: TraitData
) -> tuple[np.ndarray, None, None, np.ndarray, np.ndarray]:
    """Merge gene annotations with gene-level TraitData.

    Creates one-hot encodings for each gene set and returns unmodified
    gradients from the TraitData.

    Returns:
        Tuple of (grad, hessian, model_annot, test_annot, block_boundaries).
        hessian and model_annot are currently None for gene-set tests.
    """
    # Check if trait_data has the required key
    if self.other_key not in trait_data.keys:
        raise ValueError(f"TraitData does not have required key '{self.other_key}'. Available keys: {trait_data.keys}")
    # Use the appropriate key for merging (gene_id or gene_name)
    # For gene-level data, we need to match using the same identifier type as the gene sets
    merge_key = self.other_key

    # Verify merge key exists in trait_data
    if merge_key not in trait_data.df.columns:
        raise ValueError(f"Merge key '{merge_key}' not found in trait_data.df")

    # Get gene identifiers from trait_data
    gene_ids = trait_data.df[merge_key].to_list()

    # Create one-hot encoding for each gene set
    test_annot_dict = {}
    for set_name, gene_list in self.gene_sets.items():
        # Create binary indicator: 1 if gene is in set, 0 otherwise
        gene_set = set(gene_list)
        test_annot_dict[set_name] = [1.0 if gene in gene_set else 0.0 for gene in gene_ids]

    # Create DataFrame with one-hot encodings
    test_annot_df = pl.DataFrame(test_annot_dict)
    test_annot = test_annot_df.to_numpy().astype(np.float64)

    # Extract grad and correction (unmodified from TraitData)
    grad = trait_data.df['gradient'].to_numpy().astype(np.float64)

    # Get block boundaries
    block_boundaries = get_block_boundaries(trait_data.df['jackknife_blocks'].to_numpy())

    return grad, None, None, test_annot, block_boundaries

run_score_test

run_score_test(trait_data: TraitData, annot: Annot) -> Tuple[np.ndarray, np.ndarray]

Run approximate score test for hypothesis testing of new annotation or functional category.

This function tests annotations against precomputed gradient statistics. It does not adjust for uncertainty in fitted model parameters.

Parameters:

Name	Type	Description	Default
trait_data	TraitData	TraitData object containing variant data with gradients	required
annot	Annot	Annot object (VariantAnnot or GeneAnnot) containing annotations to test	required

Returns:

Type	Description
Tuple[ndarray, ndarray]	Tuple of (point_estimates, jackknife_estimates)

Source code in src/score_test/score_test.py

def run_score_test(trait_data: TraitData,
    annot: Annot,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Run approximate score test for hypothesis testing of new annotation or functional category.

    This function tests annotations against precomputed gradient statistics. It
    does not adjust for uncertainty in fitted model parameters.

    Args:
        trait_data: TraitData object containing variant data with gradients
        annot: Annot object (VariantAnnot or GeneAnnot) containing annotations to test

    Returns:
        Tuple of (point_estimates, jackknife_estimates)
    """
    # Merge trait data with annotations
    grad, _, _, test_annot, block_boundaries = annot.merge(trait_data)

    noBlocks = len(block_boundaries) - 1

    # Compute single-block derivatives
    U_block = [] # Derivative of the log-likelihood for each test annotation
    with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
        for i in range(noBlocks):
            block = range(block_boundaries[i], block_boundaries[i+1])
            Ui = grad[block].reshape(1,-1) @ test_annot[block, :]
            U_block.append(Ui)
    U_total = sum(U_block)

    # Compute leave-one-out derivatives
    U_jackknife = np.zeros((noBlocks, U_total.shape[1]))
    for i in range(noBlocks):
        block = range(block_boundaries[i], block_boundaries[i+1])

        # deduct the score contributed from one LD block
        U_jackknife[i, :] = U_total - U_block[i].ravel()

    return U_total, U_jackknife

main

main(variant_stats_hdf5, output_fp, variant_annot_dir, gene_annot_dir, random_genes, random_variants, gene_table, nearest_weights, annotations, trait_name, verbose, seed, perturb_annot)

Run score test for annotation enrichment.

Source code in src/score_test/score_test.py

@click.command(context_settings={'help_option_names': ['-h', '--help']})
@click.argument('variant_stats_hdf5', type=click.Path(exists=True))
@click.argument('output_fp', required=False, default=None)
@click.option('-a', '--variant-annot-dir', 'variant_annot_dir', type=click.Path(exists=True),
              help="Directory containing variant-level annotation files (.annot and/or .bed).")
@click.option('-g', '--gene-annot-dir', 'gene_annot_dir', type=click.Path(exists=True),
              help="Directory containing gene-level annotations to convert to variant-level.")
@click.option('--random-genes', 'random_genes',
              help="Comma-separated probabilities (0-1) for random gene-level annotations (e.g., '0.1,0.01').")
@click.option('--random-variants', 'random_variants',
              help="Comma-separated probabilities (0-1) for random variant-level annotations (e.g., '0.1,0.01').")
@click.option('--gene-table', default='data/genes.tsv', type=click.Path(),
              help="Path to gene table TSV file (required for gene-level options).")
@click.option('--nearest-weights', default='0.4,0.2,0.1,0.1,0.1,0.05,0.05',
              help="Comma-separated weights for k-nearest genes (for gene-level options).")
@click.option('--annotations',
              help="Optional comma-separated list of specific annotation names to test.")
@click.option('-n', '--name', 'trait_name',
              help="Specific trait name to process from HDF5 file. If omitted, all traits are processed.")
@click.option('-v', '--verbose', is_flag=True,
              help='Enable verbose output (log messages and results to console).')
@click.option('--seed', type=int, default=None,
              help='Seed for generating random annotations.')
@click.option('--perturb-annot', type=click.FloatRange(0, 1), default=0,
              help='Fraction of variants to perturb for calibration testing.')
def main(variant_stats_hdf5, output_fp, variant_annot_dir, gene_annot_dir, random_genes,
         random_variants, gene_table, nearest_weights, annotations, trait_name, verbose,
         seed, perturb_annot):
    """Run score test for annotation enrichment."""

    _setup_logging(output_fp, verbose)
    logging.info(f"{time.strftime('%Y-%m-%d %H:%M:%S')}")
    logging.info(f"Command: {' '.join(sys.argv)}")
    start_time = time.time()

    # Set random seed if provided
    if seed:
        np.random.seed(seed)

    # Parse annotation names if provided
    annot_names_filter = [a.strip() for a in annotations.split(',')] if annotations else None

    # Detect if this is gene-level or variant-level data
    is_gene_level = is_gene_level_hdf5(variant_stats_hdf5)
    data_type = "gene" if is_gene_level else "variant"

    # Load row data (variants or genes)
    data_table = load_row_data(variant_stats_hdf5)
    logging.info(f"Loaded {len(data_table)} {data_type}s from {variant_stats_hdf5}")

    # Load annotations based on source type
    weights = np.array([float(w) for w in nearest_weights.split(',')], dtype=np.float64)

    num_provided = 0
    annot = None

    if variant_annot_dir:
        if is_gene_level:
            raise click.UsageError("Cannot use --variant-annot-dir with gene-level HDF5 file")
        annot = load_variant_annotations(
            variant_annot_dir,
            annot_names_filter,
            variant_table=data_table,
        )
        logging.info(f"Loaded {len(annot.annot_names)} variant annotations from {variant_annot_dir}")
        num_provided += 1

    if gene_annot_dir:
        gene_annot: GeneAnnot = load_gene_annotations(
            gene_annot_dir, annot_names_filter
        )
        if is_gene_level:
            # For gene-level data, use gene annotations directly
            annot = gene_annot
            logging.info(f"Loaded {len(annot.annot_names)} gene annotations from {gene_annot_dir}")
        else:
            # For variant-level data, convert gene to variant annotations
            _require_gene_table(gene_table, "--gene-annot-dir")
            chromosomes = data_table['CHR'].unique().sort().to_list()
            gene_table_df = load_gene_table(gene_table, chromosomes)
            annot = convert_gene_to_variant_annotations(gene_annot, data_table, gene_table_df, weights)
            logging.info(f"Loaded {len(annot.annot_names)} gene annotations from {gene_annot_dir}")
        num_provided += 1

    if random_genes:
        probs = _parse_probs(random_genes)
        if not is_gene_level:
            _require_gene_table(gene_table, "--random-genes")
        gene_annot: GeneAnnot = create_random_gene_annotations(
            data_table, gene_table, probs
        )
        if is_gene_level:
            # For gene-level data, use gene annotations directly
            annot = gene_annot
            logging.info(f"Created {len(annot.annot_names)} random gene annotations")
        else:
            # For variant-level data, convert gene to variant annotations
            chromosomes = data_table['CHR'].unique().sort().to_list()
            gene_table_df = load_gene_table(gene_table, chromosomes)
            annot = convert_gene_to_variant_annotations(gene_annot, data_table, gene_table_df, weights)
            logging.info(f"Created {len(annot.annot_names)} random gene annotations")
        num_provided += 1

    if random_variants:
        if is_gene_level:
            raise click.UsageError("Cannot use --random-variants with gene-level HDF5 file")
        probs = _parse_probs(random_variants)
        annot = create_random_variant_annotations(data_table, probs)
        logging.info(f"Created {len(annot.annot_names)} random variant annotations")
        num_provided += 1

    if num_provided != 1:
        msg = "Must specify exactly one of: --variant-annot-dir, " + \
            "--gene-annot-dir, --random-genes, --random-variants"
        raise click.UsageError(msg)

    # Apply perturbation if requested
    if perturb_annot > 0:
        if isinstance(annot, VariantAnnot):
             logging.info(f"Perturbing annotations with fraction {perturb_annot}")
             annot.perturb(perturb_annot, seed)
        else:
             logging.warning("Perturbation only supported for VariantAnnot")

    # Run the score test
    results_dict = {'annotation' : annot.annot_names}
    trait_names = get_trait_names(variant_stats_hdf5, trait_name)
    if trait_name:
        available_trait_names = set(get_trait_names(variant_stats_hdf5))
        trait_names = [trait for trait in trait_names if trait in available_trait_names]

    # Store results for each trait for meta-analysis
    trait_results = {}

    for trait in trait_names:
        trait_data = load_trait_data(variant_stats_hdf5, trait_name=trait, variant_table=data_table)

        point_estimates, jackknife_estimates = run_score_test(
                trait_data=trait_data,
                annot=annot,
            )

        # Store for meta-analysis
        trait_results[trait] = (point_estimates, jackknife_estimates)

        # Compute Z-scores
        std_dev = np.std(jackknife_estimates, axis=0)
        z_col = f"{trait}_Z"
        z_scores = point_estimates.ravel() / std_dev / np.sqrt(jackknife_estimates.shape[0] - 1)
        results_dict[z_col] = z_scores

    # Load trait groups and perform meta-analyses
    trait_groups = get_trait_groups(variant_stats_hdf5)
    for group_name, group_traits in trait_groups.items():
        # Filter to traits that were actually processed
        group_traits = [t for t in group_traits if t in trait_results]
        z_col = f"{group_name}_Z"
        if not group_traits:
            warning = (
                f"Meta-analysis group '{group_name}' has no processed traits; "
                f"filling {z_col} with NaN"
            )
            logging.warning(warning)
            click.echo(f"Warning: {warning}", err=True)
            results_dict[z_col] = np.full(len(results_dict["annotation"]), np.nan)
            continue

        meta = MetaAnalysis()
        for trait in group_traits:
            meta.update(*trait_results[trait])
        results_dict[z_col] = meta.z_scores.ravel()
        logging.info(f"Computed meta-analysis for group '{group_name}' with {len(group_traits)} traits")

    results_df = pl.DataFrame(results_dict)

    if verbose or not output_fp:
        with pl.Config(tbl_rows=-1, tbl_cols=-1):
            print(results_df)

    if random_genes or random_variants or perturb_annot > 0:
        print("\nRoot mean squared Z-scores:")
        for col in [c for c in results_df.columns if c.endswith('_Z')]:
            print(f"{col}: {np.sqrt(np.mean(results_df[col].to_numpy()**2)):.4f}")

    if output_fp:
        results_df.write_csv(output_fp + ".txt", separator='\t')
        logging.info(f'Results written to {output_fp}.txt')

    logging.info(f'Total time: {time.time()-start_time:.2f}s')

score_test_io

I/O operations for score test.

is_gene_level_hdf5

is_gene_level_hdf5(hdf5_path: str) -> bool

Check if HDF5 file contains gene-level data.

Parameters:

Name	Type	Description	Default
hdf5_path	str	Path to HDF5 file	required

Returns:

Type	Description
bool	True if file contains gene-level data, False for variant-level

Source code in src/score_test/score_test_io.py

def is_gene_level_hdf5(hdf5_path: str) -> bool:
    """Check if HDF5 file contains gene-level data.

    Args:
        hdf5_path: Path to HDF5 file

    Returns:
        True if file contains gene-level data, False for variant-level
    """
    with h5py.File(hdf5_path, "r") as f:
        # Check metadata attribute first
        if "data_type" in f.attrs:
            return f.attrs["data_type"] == "gene"
        # Fallback: check if row_data contains gene_id
        if "row_data" not in f:
            raise ValueError("HDF5 file must contain 'row_data' group")
        return "gene_id" in f["row_data"]

load_row_data

load_row_data(hdf5_path: str) -> pl.DataFrame

Load row data table from HDF5 file format.

Parameters:

Name	Type	Description	Default
hdf5_path	str	Path to the HDF5 file	required

Returns:

Type	Description
DataFrame	Polars DataFrame containing row data (variants or genes)

Source code in src/score_test/score_test_io.py

def load_row_data(hdf5_path: str) -> pl.DataFrame:
    """
    Load row data table from HDF5 file format.

    Args:
        hdf5_path: Path to the HDF5 file

    Returns:
        Polars DataFrame containing row data (variants or genes)
    """

    if not os.path.exists(hdf5_path):
        raise FileNotFoundError(f"HDF5 file not found: {hdf5_path}")

    with h5py.File(hdf5_path, "r") as f:
        if "row_data" not in f:
            raise ValueError("HDF5 file must contain 'row_data' group")

        # Use helper to load all datasets
        data = _load_hdf5_group(f["row_data"])

        # Remove attributes (prefixed with @)
        data = {k: v for k, v in data.items() if not k.startswith("@")}

        # Process arrays: flatten 2D columns with shape (n, 1) and convert bytes to strings
        for key, arr in data.items():
            if hasattr(arr, "shape"):
                # Flatten 2D arrays with single column
                if len(arr.shape) == 2 and arr.shape[1] == 1:
                    data[key] = arr.ravel()
                    arr = data[key]  # update reference

                # Convert bytes to strings
                data[key] = _decode_bytes_array(arr)

    return pl.DataFrame(data)

get_trait_names

get_trait_names(hdf5_path: str, trait_name: Optional[str] = None) -> List[str]

Get list of trait names from HDF5 file.

Parameters:

Name	Type	Description	Default
hdf5_path	str	Path to the HDF5 file containing trait data	required
trait_name	Optional[str]	Optional specific trait name or meta-analysis name. If provided, returns list with just this trait, or expands meta-analysis to its traits. If None, returns all trait names in the file.	None

Returns:

Type	Description
List[str]	List of trait names

Source code in src/score_test/score_test_io.py

def get_trait_names(hdf5_path: str, trait_name: Optional[str] = None) -> List[str]:
    """
    Get list of trait names from HDF5 file.

    Args:
        hdf5_path: Path to the HDF5 file containing trait data
        trait_name: Optional specific trait name or meta-analysis name.
                   If provided, returns list with just this trait, or expands meta-analysis to its traits.
                   If None, returns all trait names in the file.

    Returns:
        List of trait names
    """
    with h5py.File(hdf5_path, "r") as f:
        if trait_name:
            # Check if it's a trait first
            if trait_name in f["traits"].keys():
                return [trait_name]

            # Check if it's a meta-analysis group
            groups = get_trait_groups(hdf5_path)
            if trait_name in groups:
                return groups[trait_name]

            # Not found as either trait or meta-analysis
            raise ValueError(
                f"Trait or meta-analysis '{trait_name}' not found in HDF5 file"
            )
        else:
            return list(f["traits"].keys())

load_trait_hdf5

load_trait_hdf5(hdf5_path: str, trait_name: str) -> dict

Load trait data from HDF5 file format.

Returns:

Type	Description
dict	Dictionary with keys: gradient (required), and optionally parameters, jackknife_parameters, other datasets

Source code in src/score_test/score_test_io.py

def load_trait_hdf5(hdf5_path: str, trait_name: str) -> dict:
    """
    Load trait data from HDF5 file format.

    Returns:
        Dictionary with keys: gradient (required), and optionally parameters, jackknife_parameters, other datasets
    """

    with h5py.File(hdf5_path, "r") as f:
        trait_group = f[f"traits/{trait_name}"]

        # Use helper to load everything recursively
        data = _load_hdf5_group(trait_group)

    return data

load_trait_data

load_trait_data(hdf5_path: str, trait_name: str, variant_table: DataFrame) -> TraitData

Load trait data and combine with variant data into a TraitData object.

Parameters:

Name	Type	Description	Default
hdf5_path	str	Path to HDF5 file containing variant and trait data	required
trait_name	str	Name of the trait to load	required
variant_table	DataFrame	Pre-loaded variant table DataFrame	required

Returns:

Type	Description
TraitData	TraitData object containing variant dataframe with gradients/corrections and parameters

Source code in src/score_test/score_test_io.py

def load_trait_data(
    hdf5_path: str, trait_name: str, variant_table: pl.DataFrame
) -> "TraitData":
    """
    Load trait data and combine with variant data into a TraitData object.

    Args:
        hdf5_path: Path to HDF5 file containing variant and trait data
        trait_name: Name of the trait to load
        variant_table: Pre-loaded variant table DataFrame

    Returns:
        TraitData object containing variant dataframe with gradients/corrections and parameters
    """
    # Import at runtime to avoid circular import
    try:
        from .score_test import TraitData
    except ImportError:
        from score_test import TraitData

    trait_hdf5 = load_trait_hdf5(hdf5_path, trait_name)

    # Determine primary key column
    possible_keys = ["RSID", "gene_id", "gene_name", "CHR", "POS"]
    keys = [key for key in possible_keys if key in variant_table.columns]
    if len(keys) == 0:
        raise ValueError("variant_table must have one of: " + ", ".join(possible_keys))

    # Add trait-specific columns (gradient, hessian, etc.)
    new_columns = []
    for key_name, value in trait_hdf5.items():
        if isinstance(value, np.ndarray):
            decoded_value = _decode_bytes_array(value)
            new_columns.append(pl.Series(name=key_name, values=decoded_value))

    df = variant_table.with_columns(new_columns)

    if "parameters" in trait_hdf5:
        params = trait_hdf5["parameters"]["parameters"]
        jk_params = trait_hdf5["parameters"]["jackknife_parameters"]
    else:
        params, jk_params = None, None

    # Create TraitData - it will compute exclude_cols and annot_names via properties
    trait_data = TraitData(
        df=df,
        params=params,
        jk_params=jk_params,
        keys=keys,
    )

    # Compute annotation names using the property
    trait_data.annot_names = [
        col for col in df.columns if col not in trait_data.exclude_cols
    ]

    return trait_data

load_annotations

load_annotations(annot_path: str, chromosome: Optional[int] = None, infer_schema_length: int = 100000, add_positions: bool = True, variant_table: DataFrame | None = None, exclude_bed: bool = False) -> pl.DataFrame

Load annotation data for specified chromosome(s).

Parameters:

Name	Type	Description	Default
annot_path	str	Path to directory containing .annot and/or .bed 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	100000
add_positions	bool	If True, rename BP to POS	True
variant_table	DataFrame | None	Variant table to use as the coordinate universe for BED-only annotation directories	None
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/score_test/score_test_io.py

def load_annotations(
    annot_path: str,
    chromosome: Optional[int] = None,
    infer_schema_length: int = 100_000,
    add_positions: bool = True,
    variant_table: pl.DataFrame | None = None,
    exclude_bed: bool = False,
) -> pl.DataFrame:
    """Load annotation data for specified chromosome(s).

    Args:
        annot_path: Path to directory containing .annot and/or .bed files
        chromosome: Specific chromosome number, or None for all chromosomes
        infer_schema_length: Number of rows to infer schema from
        add_positions: If True, rename BP to POS
        variant_table: Variant table to use as the coordinate universe for
            BED-only annotation directories
        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
    """
    bed_files = list_bed_files(annot_path)

    # Determine which chromosomes to process
    if chromosome is not None:
        chromosomes = [chromosome]
    else:
        chromosomes = range(1, 23)  # Assuming chromosomes 1-22

    # Find matching files
    annotations = []
    for chromosome in chromosomes:
        file_pattern = f"*.{chromosome}.annot"
        matching_files = sorted(Path(annot_path).glob(file_pattern))

        # Read all matching files for this chromosome
        dfs = []
        cols_found = set()
        for file_path in matching_files:
            df = pl.scan_csv(
                file_path, separator="\t", infer_schema_length=infer_schema_length
            )
            schema_names = df.collect_schema().names()
            cols_to_keep = [col for col in schema_names if col not in cols_found]
            if cols_to_keep:
                df = df.select(cols_to_keep)
                cols_found.update(cols_to_keep)
                dfs.append(df)

        # Horizontally concatenate all dataframes for this chromosome
        if dfs:
            combined_df = pl.concat(dfs, how="horizontal")
            combined_df = combined_df.select(sorted(combined_df.collect_schema().names()))
            annotations.append(combined_df)

    if not annotations:
        if bed_files and not exclude_bed:
            annotations = _variant_table_as_annotation_table(variant_table, chromosomes)
        else:
            raise ValueError(
                f"No annotation files found in {annot_path} matching "
                "pattern *.{chrom}.annot or *.bed"
            )

    if isinstance(annotations, list):
        # Concatenate all chromosome dataframes vertically
        annotations = pl.concat(annotations, how="vertical").collect()

    if not isinstance(annotations, pl.DataFrame):
        raise ValueError(
            f"No annotation files found in {annot_path} matching pattern *.{{chrom}}.annot"
        )

    # 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 add_positions:
        annotations = annotations.rename({"BP": "POS"})

    if exclude_bed or not bed_files:
        return annotations

    position_col = "POS" if add_positions else "BP"
    return add_bed_annotations(annotations, bed_files, position_col=position_col)

load_variant_annotations

load_variant_annotations(annot_dir: str, annot_names: list[str] | None = None, variant_table: DataFrame | None = None) -> VariantAnnot

Load variant-level annotations from directory.

Parameters:

Name	Type	Description	Default
annot_dir	str	Directory containing .annot and/or .bed annotation files	required
annot_names	list[str] | None	Optional list of specific annotation names to load	None
variant_table	DataFrame | None	Variant table to use as the coordinate universe for BED-only annotation directories	None

Returns:

Type	Description
VariantAnnot	VariantAnnot object

Source code in src/score_test/score_test_io.py

def load_variant_annotations(
    annot_dir: str,
    annot_names: list[str] | None = None,
    variant_table: pl.DataFrame | None = None,
) -> "VariantAnnot":
    """Load variant-level annotations from directory.

    Args:
        annot_dir: Directory containing .annot and/or .bed annotation files
        annot_names: Optional list of specific annotation names to load
        variant_table: Variant table to use as the coordinate universe for
            BED-only annotation directories

    Returns:
        VariantAnnot object
    """
    # Import at runtime to avoid circular import
    try:
        from .score_test import VariantAnnot
    except ImportError:
        from score_test import VariantAnnot

    df_annot = load_annotations(
        annot_dir, add_positions=False, variant_table=variant_table
    )

    # Rename SNP to RSID for consistency
    if "SNP" in df_annot.columns:
        df_annot = df_annot.rename({"SNP": "RSID"})

    merge_key: str | list[str] = "RSID"
    if "RSID" not in df_annot.columns and {"CHR", "BP"}.issubset(df_annot.columns):
        df_annot = df_annot.rename({"BP": "POS"})
        merge_key = ["CHR", "POS"]

    # Exclude positional and identifier columns from annotations
    exclude_cols = ["CHR", "BP", "POS", "RSID", "CM"]

    if annot_names:
        available = [col for col in df_annot.columns if col not in exclude_cols]
        annot_names = [name for name in annot_names if name in available]
    else:
        annot_names = [col for col in df_annot.columns if col not in exclude_cols]

    return VariantAnnot(df_annot, annot_names, other_key=merge_key)

load_gene_annotations

load_gene_annotations(gene_annot_dir: str, annot_names: list[str] | None = None) -> GeneAnnot

Load GMT gene-set files from a directory into a GeneAnnot object.

Parameters:

Name	Type	Description	Default
gene_annot_dir	str	Directory containing GMT files with gene sets	required
annot_names	list[str] | None	Optional list of specific annotation names to load	None

Returns:

Type	Description
GeneAnnot	GeneAnnot object

Source code in src/score_test/score_test_io.py

def load_gene_annotations(
    gene_annot_dir: str,
    annot_names: list[str] | None = None,
) -> "GeneAnnot":
    """Load GMT gene-set files from a directory into a GeneAnnot object.

    Args:
        gene_annot_dir: Directory containing GMT files with gene sets
        annot_names: Optional list of specific annotation names to load

    Returns:
        GeneAnnot object
    """
    # Import at runtime to avoid circular import
    try:
        from .score_test import GeneAnnot
    except ImportError:
        from score_test import GeneAnnot

    gene_sets = load_gene_sets_from_gmt(gene_annot_dir)

    if annot_names:
        gene_sets = {
            name: genes for name, genes in gene_sets.items() if name in annot_names
        }

    return GeneAnnot(gene_sets)

create_random_gene_annotations

create_random_gene_annotations(data_table: DataFrame, gene_table_path: str, probs: List[float]) -> GeneAnnot

Create random gene-level annotations as a GeneAnnot object.

Parameters:

Name	Type	Description	Default
data_table	DataFrame	Data table DataFrame (variants or genes)	required
gene_table_path	str	Path to gene table TSV file	required
probs	List[float]	List of probabilities for random gene selection	required

Returns:

Type	Description
GeneAnnot	GeneAnnot object with random gene sets

Source code in src/score_test/score_test_io.py

def create_random_gene_annotations(
    data_table: pl.DataFrame,
    gene_table_path: str,
    probs: List[float],
) -> "GeneAnnot":
    """Create random gene-level annotations as a GeneAnnot object.

    Args:
        data_table: Data table DataFrame (variants or genes)
        gene_table_path: Path to gene table TSV file
        probs: List of probabilities for random gene selection

    Returns:
        GeneAnnot object with random gene sets
    """
    # Import at runtime to avoid circular import
    try:
        from .score_test import GeneAnnot
    except ImportError:
        from score_test import GeneAnnot

    # Check if this is gene-level data (has gene_id column)
    if "gene_id" in data_table.columns:
        # For gene-level data, use genes directly from the data table
        gene_names = data_table["gene_name"].to_list()
    else:
        # For variant-level data, load genes from file
        chromosomes = data_table["CHR"].unique().sort().to_list()
        gene_table = load_gene_table(gene_table_path, chromosomes)
        gene_names = gene_table["gene_name"].to_list()

    # Create random gene sets
    gene_sets = {}
    for i, p in enumerate(probs):
        set_name = f"random_gene_{i}"
        n_genes = int(len(gene_names) * p)
        gene_sets[set_name] = list(
            np.random.choice(gene_names, size=n_genes, replace=False)
        )

    return GeneAnnot(gene_sets)

create_random_variant_annotations

create_random_variant_annotations(variant_table: DataFrame, probs: List[float]) -> VariantAnnot

Create random variant-level annotations as a VariantAnnot object.

Parameters:

Name	Type	Description	Default
variant_table	DataFrame	Variant table DataFrame	required
probs	List[float]	List of probabilities for random annotation	required

Returns:

Type	Description
VariantAnnot	VariantAnnot object with random variant annotations

Source code in src/score_test/score_test_io.py

def create_random_variant_annotations(
    variant_table: pl.DataFrame,
    probs: List[float],
) -> "VariantAnnot":
    """Create random variant-level annotations as a VariantAnnot object.

    Args:
        variant_table: Variant table DataFrame
        probs: List of probabilities for random annotation

    Returns:
        VariantAnnot object with random variant annotations
    """
    # Import at runtime to avoid circular import
    try:
        from .score_test import VariantAnnot
    except ImportError:
        from score_test import VariantAnnot

    num_variants = len(variant_table)

    # Create random annotations
    variant_annots = {}
    annot_names = []

    for i, p in enumerate(probs):
        col_name = f"random_variant_{i}"
        annot_names.append(col_name)
        variant_annots[col_name] = np.random.binomial(1, p, size=num_variants).astype(
            np.float64
        )

    # Create output DataFrame
    # Use RSID for variants
    if "RSID" not in variant_table.columns:
        raise ValueError(
            "variant_table must have 'RSID' column for variant annotations"
        )

    df_annot = pl.DataFrame(
        {
            "CHR": variant_table["CHR"],
            "BP": variant_table["POS"],
            "RSID": variant_table["RSID"],
            "CM": pl.Series([0.0] * len(variant_table)),
            **variant_annots,
        }
    )

    return VariantAnnot(df_annot, annot_names)

get_trait_groups

get_trait_groups(hdf5_path: str) -> dict[str, list[str]]

Get trait groups for meta-analysis from HDF5 file.

Parameters:

Name	Type	Description	Default
hdf5_path	str	Path to HDF5 file containing trait groups	required

Returns:

Type	Description
dict[str, list[str]]	Dictionary mapping group names to lists of trait names.
dict[str, list[str]]	Returns empty dict if no groups are defined.

Source code in src/score_test/score_test_io.py

def get_trait_groups(hdf5_path: str) -> dict[str, list[str]]:
    """Get trait groups for meta-analysis from HDF5 file.

    Args:
        hdf5_path: Path to HDF5 file containing trait groups

    Returns:
        Dictionary mapping group names to lists of trait names.
        Returns empty dict if no groups are defined.
    """
    with h5py.File(hdf5_path, "r") as f:
        if "groups" not in f:
            return {}

        groups = {}
        groups_group = f["groups"]
        for group_name in groups_group.keys():
            trait_list = groups_group[group_name][:]
            # Convert bytes to strings if necessary
            if trait_list.dtype.kind == "S":
                trait_list = [t.decode("utf-8") for t in trait_list]
            else:
                trait_list = trait_list.tolist()
            groups[group_name] = trait_list

        return groups

save_trait_groups

save_trait_groups(hdf5_path: str, groups: dict[str, list[str]]) -> None

Save trait groups for meta-analysis to HDF5 file.

Parameters:

Name	Type	Description	Default
hdf5_path	str	Path to HDF5 file to create/update	required
groups	dict[str, list[str]]	Dictionary mapping group names to lists of trait names	required

Source code in src/score_test/score_test_io.py

def save_trait_groups(hdf5_path: str, groups: dict[str, list[str]]) -> None:
    """Save trait groups for meta-analysis to HDF5 file.

    Args:
        hdf5_path: Path to HDF5 file to create/update
        groups: Dictionary mapping group names to lists of trait names
    """
    with h5py.File(hdf5_path, "a") as f:
        if "groups" in f:
            del f["groups"]

        groups_group = f.create_group("groups")

        for group_name, trait_list in groups.items():
            # Convert strings to bytes for HDF5 storage
            trait_array = np.array(trait_list, dtype="S")
            groups_group.create_dataset(group_name, data=trait_array)

save_trait_data

save_trait_data(trait_data: TraitData, hdf5_path: str, trait_name: str) -> None

Save trait data to HDF5 file in new format.

Parameters:

Name	Type	Description	Default
trait_data	TraitData	TraitData object to save	required
hdf5_path	str	Path to HDF5 file to create/update	required
trait_name	str	Name of the trait	required

Source code in src/score_test/score_test_io.py

def save_trait_data(
    trait_data: "TraitData",
    hdf5_path: str,
    trait_name: str,
) -> None:
    """Save trait data to HDF5 file in new format.

    Args:
        trait_data: TraitData object to save
        hdf5_path: Path to HDF5 file to create/update
        trait_name: Name of the trait
    """
    with h5py.File(hdf5_path, "a") as f:
        # Create metadata attribute if it doesn't exist
        if "metadata" not in f.attrs:
            f.attrs["metadata"] = ""

        # Set data_type and keys attributes
        if "gene_id" in trait_data.df.columns:
            f.attrs["data_type"] = "gene"
            f.attrs["keys"] = ["gene_id", "gene_name"]
        else:
            f.attrs["data_type"] = "variant"
            f.attrs["keys"] = ["RSID", "POS"]

        # Create row_data group if it doesn't exist, or recreate if size mismatch
        if "row_data" in f:
            existing_size = len(f["row_data"][list(f["row_data"].keys())[0]])
            expected_size = len(trait_data.df)
            if existing_size != expected_size:
                # Size mismatch - delete and recreate
                del f["row_data"]

        if "row_data" not in f:
            data_group = f.create_group("row_data")

            # Save all columns except trait-specific ones (gradient, hessian, etc.)
            # Use TraitData's exclude_cols property
            row_data_cols = set(trait_data.df.columns) - trait_data.exclude_cols
            # Add back standard columns that should be in row_data
            row_data_cols.update({"CHR", "POS", "jackknife_blocks"})
            if "RSID" in trait_data.df.columns:
                row_data_cols.add("RSID")
            if "gene_id" in trait_data.df.columns:
                row_data_cols.update({"gene_id", "gene_name"})

            for col in trait_data.df.columns:
                if col in row_data_cols:
                    data = trait_data.df[col].to_numpy()
                    # Handle string columns
                    if data.dtype == object:
                        data = data.astype("S")
                    data_group.create_dataset(col, data=data)

        # Create traits group if it doesn't exist
        if "traits" not in f:
            f.create_group("traits")

        trait_path = f"traits/{trait_name}"
        if trait_path in f:
            del f[trait_path]

        trait_group = f.create_group(trait_path)

        # Save parameters in parameters/ subgroup
        if trait_data.params is not None or trait_data.jk_params is not None:
            params_group = trait_group.create_group("parameters")
            if trait_data.params is not None:
                params_group.create_dataset("parameters", data=trait_data.params)
            if trait_data.jk_params is not None:
                params_group.create_dataset(
                    "jackknife_parameters", data=trait_data.jk_params
                )

        # Save gradient (required)
        trait_group.create_dataset(
            "gradient", data=trait_data.df["gradient"].to_numpy()
        )

        # Save any other trait-specific datasets (e.g., hessian)
        standard_cols = {
            "CHR",
            "POS",
            "jackknife_blocks",
            "RSID",
            "gene_id",
            "gene_name",
            "gradient",
        }
        annot_cols = set(trait_data.annot_names) if trait_data.annot_names else set()
        for col in trait_data.df.columns:
            if col not in standard_cols and col not in annot_cols:
                trait_group.create_dataset(col, data=trait_data.df[col].to_numpy())