Skip to content

graphld.simulate

Simulation of GWAS summary statistics.

Summary statistics can be simulated from their asymptotic distribution without individual-level genotype data. Effect sizes are drawn from a flexible mixture distribution, with Python API support for annotation-dependent effect-size scaling and frequency-dependent architectures. Annotation-dependent polygenicity is reserved for future support.

For usage examples, see the Simulation guide.

simulate

Simulate GWAS summary statistics.

Simulate dataclass

Simulate(sample_size: int, heritability: float = 0.5, component_variance: Union[ndarray, List[float]] = None, component_weight: Union[ndarray, List[float]] = None, alpha_param: float = -1, annotation_dependent_polygenicity: bool = False, link_fn: Callable[[ndarray], ndarray] = _default_link_fn, random_seed: Optional[int] = None, annotation_columns: Optional[List[str]] = None)

Bases: ParallelProcessor, _SimulationSpecification

Parallel processor for simulating GWAS summary statistics.

create_shared_memory staticmethod

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

Create shared memory arrays for simulation.

Parameters:

Name Type Description Default
metadata DataFrame

Metadata DataFrame containing block information

required
block_data list[tuple]

List of tuples containing block-specific annotation DataFrames

required
**kwargs Any

Additional keyword arguments

{}
Source code in src/graphld/simulate.py
@staticmethod
def create_shared_memory(
    metadata: pl.DataFrame,
    block_data: list[tuple],
    **kwargs: Any
) -> SharedData:
    """Create shared memory arrays for simulation.

    Args:
        metadata: Metadata DataFrame containing block information
        block_data: List of tuples containing block-specific annotation DataFrames
        **kwargs: Additional keyword arguments
    """
    # Get total number of variants and indices
    num_variants = np.array([len(df) for df, _ in block_data])
    total_variants = int(sum(num_variants))  # Convert to Python int

    # Create shared arrays sized according to metadata
    shared = SharedData({
        'beta': total_variants,    # Causal effect sizes (one per index)
        'alpha': total_variants,   # Marginal effect sizes (one per index)
        'h2': total_variants,      # Per-variant heritability
        'noise': total_variants,   # Noise component
        'scale_param': 1,         # Single value
    })

    # Initialize arrays with zeros
    shared['beta'][:] = 0
    shared['alpha'][:] = 0
    shared['noise'][:] = 0
    shared['h2'][:] = 0
    shared['scale_param'][:] = 0

    return shared

prepare_block_data classmethod

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

Prepare block-specific data for processing.

Parameters:

Name Type Description Default
metadata DataFrame

DataFrame containing LDGM metadata

required
**kwargs Any

Additional arguments from run(), including: annotations: Optional DataFrame containing variant annotations

{}

Returns:

Type Description
list[tuple]

List of block-specific annotation DataFrames, or None if no annotations

Source code in src/graphld/simulate.py
@classmethod
def prepare_block_data(cls, metadata: pl.DataFrame, **kwargs: Any) -> list[tuple]:
    """Prepare block-specific data for processing.

    Args:
        metadata: DataFrame containing LDGM metadata
        **kwargs: Additional arguments from run(), including:
            annotations: Optional DataFrame containing variant annotations

    Returns:
        List of block-specific annotation DataFrames, or None if no annotations
    """
    annotations = kwargs.get('annotations')
    if annotations is None:
        block_annotations = _create_block_annotations(
            metadata,
            kwargs.get('ldgm_metadata_path_duplicate')
        )
    else:
        block_annotations = partition_variants(metadata, annotations)

    cumulative_num_variants = np.cumsum(np.array([len(df) for df in block_annotations]))
    cumulative_num_variants = [0] + list(cumulative_num_variants[:-1])

    return list(zip(block_annotations, cumulative_num_variants, strict=False))

process_block classmethod

process_block(ldgm: PrecisionOperator, flag: Value, shared_data: SharedData, block_offset: int, block_data: Optional[tuple] = None, worker_params: Optional[Dict] = None) -> None

Process a single block.

Source code in src/graphld/simulate.py
@classmethod
def process_block(cls, ldgm: PrecisionOperator, flag: Value,
                 shared_data: SharedData, block_offset: int,
                 block_data: Optional[tuple] = None,
                 worker_params: Optional[Dict] = None) -> None:
    """Process a single block."""
    # If we have block_data, merge it with LDGM variant info
    if block_data is not None:
        assert isinstance(block_data, tuple), "block_data must be a tuple"
        annotations, variant_offset = block_data
        num_variants = len(annotations)

        # Merge annotations with LDGM variant info and get indices of merged variants
        from .io import merge_snplists
        ldgm, sumstat_indices = merge_snplists(
            ldgm, annotations,
            match_by_position=True,
            pos_col='POS',
            ref_allele_col='A2',
            alt_allele_col='A1'
        )
        phase = (
            ldgm.variant_info.get_column('phase').to_numpy()
            if 'phase' in ldgm.variant_info.columns else 1
        )
    else:
        variant_offset = block_offset
        num_variants = ldgm.shape[0]
        sumstat_indices = range(num_variants)
        phase = 1

    # Get block slice using the number of indices
    block_slice = slice(variant_offset, variant_offset + num_variants)

    block_random_seed = None if worker_params.random_seed is None \
        else worker_params.random_seed + variant_offset

    # Simulate effect sizes using the merged data
    beta, alpha = _simulate_beta_block(
        ldgm,
        worker_params,
        random_seed=block_random_seed,
    )
    noise = _simulate_noise_block(ldgm, random_seed=block_random_seed)

    # Create zero-filled arrays for all variants in sumstats
    beta_reshaped = np.zeros(num_variants)
    alpha_reshaped = np.zeros(num_variants)
    noise_reshaped = np.zeros(num_variants)

    # Fill in values for successfully merged variants
    beta_reshaped[sumstat_indices] = beta.flatten() * phase
    alpha_reshaped[sumstat_indices] = alpha.flatten() * phase
    noise_reshaped[sumstat_indices] = noise.flatten() * phase

    # Update the shared memory arrays
    block_slice = slice(variant_offset, variant_offset + num_variants)
    shared_data['beta', block_slice] = beta_reshaped
    shared_data['alpha', block_slice] = alpha_reshaped
    shared_data['noise', block_slice] = noise_reshaped

supervise classmethod

supervise(manager: WorkerManager, shared_data: Dict[str, Any], block_data: list, **kwargs: Any) -> pl.DataFrame

Supervise worker processes and collect results.

Parameters:

Name Type Description Default
manager WorkerManager

Worker manager

required
shared_data Dict[str, Any]

Dictionary of shared memory arrays

required
**kwargs Any

Additional arguments

{}

Returns:

Type Description
DataFrame

DataFrame containing simulated summary statistics

Source code in src/graphld/simulate.py
@classmethod
def supervise(
    cls,
    manager: WorkerManager,
    shared_data: Dict[str, Any],
    block_data: list,
    **kwargs: Any
) -> pl.DataFrame:
    """Supervise worker processes and collect results.

    Args:
        manager: Worker manager
        shared_data: Dictionary of shared memory arrays
        **kwargs: Additional arguments

    Returns:
        DataFrame containing simulated summary statistics
    """
    manager.start_workers()
    manager.await_workers()
    beta, alpha, noise = shared_data['beta'], shared_data['alpha'], shared_data['noise']

    # Compute scaling parameter to achieve desired heritability
    spec = kwargs['spec']
    h2 = spec.heritability
    current_h2 = np.dot(beta,alpha)
    assert current_h2 >= 0, "beta'*R*beta should be non-negative"
    scaling_param = np.sqrt(h2 / current_h2) if current_h2 > 0 else 1
    beta *= scaling_param
    alpha *= scaling_param

    # Concatenate block annotations and add simulation results
    result = pl.concat([
        df.select(['CHR', 'SNP', 'POS', 'A1', 'A2'])
        for df, _ in block_data
    ])

    return result.with_columns([
        pl.Series('Z', noise + np.sqrt(spec.sample_size) * alpha),
        pl.Series('beta', beta),
        pl.Series('beta_marginal', alpha),
        pl.Series('N', np.full(len(result), spec.sample_size)),
    ])

simulate

simulate(ldgm_metadata_path: str = 'data/ldgms/metadata.csv', populations: Optional[Union[str, List[str]]] = 'EUR', chromosomes: Optional[Union[int, List[int]]] = None, run_in_serial: bool = False, num_processes: Optional[int] = None, annotations: Optional[DataFrame] = None, verbose: bool = False) -> pl.DataFrame

Simulate genetic data.

Parameters:

Name Type Description Default
ldgm_metadata_path str

Path to LDGM metadata file

'data/ldgms/metadata.csv'
populations Optional[Union[str, List[str]]]

Population(s) to filter

'EUR'
chromosomes Optional[Union[int, List[int]]]

Chromosome(s) to filter

None
run_in_serial bool

Whether to run in serial mode

False
annotations Optional[DataFrame]

Optional variant annotations

None

Returns:

Type Description
DataFrame

Simulated genetic data DataFrame

Source code in src/graphld/simulate.py
def simulate(
    self,
    ldgm_metadata_path: str = 'data/ldgms/metadata.csv',
    populations: Optional[Union[str, List[str]]] = 'EUR',
    chromosomes: Optional[Union[int, List[int]]] = None,
    run_in_serial: bool = False,
    num_processes: Optional[int] = None,
    annotations: Optional[pl.DataFrame] = None,
    verbose: bool = False,
) -> pl.DataFrame:
    """Simulate genetic data.

    Args:
        ldgm_metadata_path: Path to LDGM metadata file
        populations: Population(s) to filter
        chromosomes: Chromosome(s) to filter
        run_in_serial: Whether to run in serial mode
        annotations: Optional variant annotations

    Returns:
        Simulated genetic data DataFrame
    """
    # Check if link_fn is picklable when using multiprocessing
    if not run_in_serial and self.link_fn is not _default_link_fn:
        import pickle
        try:
            pickle.dumps(self.link_fn)
        except (pickle.PicklingError, AttributeError, TypeError) as e:
            raise ValueError(
                f"link_fn must be picklable to work with multiprocessing. "
                f"Lambda functions and nested functions cannot be pickled. "
                f"Define your link function at module level or use run_in_serial=True. "
                f"Original error: {e}"
            ) from e

    run_fn = self.run_serial if run_in_serial else self.run
    result = run_fn(
        ldgm_metadata_path=ldgm_metadata_path,
        populations=populations,
        chromosomes=chromosomes,
        worker_params=self,  # Use instance itself as spec
        spec=self,
        annotations=annotations,
        # Pass path to prepare_block_data
        ldgm_metadata_path_duplicate=ldgm_metadata_path,
        num_processes=num_processes,
    )

    if verbose:
        print(f"Number of variants in summary statistics: {len(result)}")
        nonzero_count = (result['beta'] != 0).sum()
        print(f"Number of variants with nonzero beta: {nonzero_count}")

    return result

run_simulate

run_simulate(sample_size: int, heritability: float = 0.5, component_variance: Optional[Union[ndarray, List[float]]] = None, component_weight: Optional[Union[ndarray, List[float]]] = None, alpha_param: float = -1, annotation_dependent_polygenicity: bool = False, link_fn: Callable[[ndarray], ndarray] = _default_link_fn, random_seed: Optional[int] = None, annotation_columns: Optional[List[str]] = None, ldgm_metadata_path: str = 'data/ldgms/metadata.csv', populations: Optional[Union[str, List[str]]] = 'EUR', chromosomes: Optional[Union[int, List[int]]] = None, run_in_serial: bool = False, num_processes: Optional[int] = None, annotations: Optional[DataFrame] = None, verbose: bool = False) -> pl.DataFrame

Run GWAS summary statistics simulation with specified parameters.

Parameters:

Name Type Description Default
sample_size int

Sample size for the population

required
heritability float

Total heritability (h2) for the trait

0.5
component_variance Optional[Union[ndarray, List[float]]]

Per-allele effect size variance for each mixture component

None
component_weight Optional[Union[ndarray, List[float]]]

Mixture weight for each component (must sum to ≤ 1)

None
alpha_param float

Alpha parameter for allele frequency-dependent architecture

-1
annotation_dependent_polygenicity bool

Reserved for future support for using annotations to modify the proportion of causal variants. Currently raises NotImplementedError when enabled.

False
link_fn Callable[[ndarray], ndarray]

Function mapping annotation vector to relative per-variant heritability. Default is the softplus-like mapping x -> log(1 + exp(sum(x))). Must be defined at module level (not as lambda or nested function) to work with multiprocessing

_default_link_fn
random_seed Optional[int]

Random seed for reproducibility

None
annotation_columns Optional[List[str]]

List of column names to use as annotations

None
ldgm_metadata_path str

Path to LDGM metadata file

'data/ldgms/metadata.csv'
populations Optional[Union[str, List[str]]]

Population(s) to filter

'EUR'
chromosomes Optional[Union[int, List[int]]]

Chromosome(s) to filter

None
run_in_serial bool

Whether to run in serial mode

False
num_processes Optional[int]

Number of processes for parallel execution

None
annotations Optional[DataFrame]

Optional variant annotations DataFrame

None
verbose bool

Whether to print progress information

False

Returns:

Type Description
DataFrame

DataFrame containing simulated summary statistics

Source code in src/graphld/simulate.py
def run_simulate(
    sample_size: int,
    heritability: float = 0.5,
    component_variance: Optional[Union[np.ndarray, List[float]]] = None,
    component_weight: Optional[Union[np.ndarray, List[float]]] = None,
    alpha_param: float = -1,
    annotation_dependent_polygenicity: bool = False,
    link_fn: Callable[[np.ndarray], np.ndarray] = _default_link_fn,
    random_seed: Optional[int] = None,
    annotation_columns: Optional[List[str]] = None,
    ldgm_metadata_path: str = 'data/ldgms/metadata.csv',
    populations: Optional[Union[str, List[str]]] = 'EUR',
    chromosomes: Optional[Union[int, List[int]]] = None,
    run_in_serial: bool = False,
    num_processes: Optional[int] = None,
    annotations: Optional[pl.DataFrame] = None,
    verbose: bool = False,
) -> pl.DataFrame:
    """Run GWAS summary statistics simulation with specified parameters.

    Args:
        sample_size: Sample size for the population
        heritability: Total heritability (h2) for the trait
        component_variance: Per-allele effect size variance for each mixture component
        component_weight: Mixture weight for each component (must sum to ≤ 1)
        alpha_param: Alpha parameter for allele frequency-dependent architecture
        annotation_dependent_polygenicity: Reserved for future support for using
            annotations to modify the proportion of causal variants. Currently raises
            NotImplementedError when enabled.
        link_fn: Function mapping annotation vector to relative per-variant heritability.
            Default is the softplus-like mapping x -> log(1 + exp(sum(x))).
            Must be defined at module level (not as lambda or nested function)
            to work with multiprocessing
        random_seed: Random seed for reproducibility
        annotation_columns: List of column names to use as annotations
        ldgm_metadata_path: Path to LDGM metadata file
        populations: Population(s) to filter
        chromosomes: Chromosome(s) to filter
        run_in_serial: Whether to run in serial mode
        num_processes: Number of processes for parallel execution
        annotations: Optional variant annotations DataFrame
        verbose: Whether to print progress information

    Returns:
        DataFrame containing simulated summary statistics
    """
    sim = Simulate(
        sample_size=sample_size,
        heritability=heritability,
        component_variance=component_variance,
        component_weight=component_weight,
        alpha_param=alpha_param,
        annotation_dependent_polygenicity=annotation_dependent_polygenicity,
        link_fn=link_fn,
        random_seed=random_seed,
        annotation_columns=annotation_columns,
    )

    return sim.simulate(
        ldgm_metadata_path=ldgm_metadata_path,
        populations=populations,
        chromosomes=chromosomes,
        run_in_serial=run_in_serial,
        num_processes=num_processes,
        annotations=annotations,
        verbose=verbose,
    )