graphld.io

Core I/O functions for loading LDGMs and working with variant data.

For end-to-end loading and merge examples, see the I/O and Merging guide.

io

Input/output operations for LDGM files.

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)