agricola Python API

def step1(
    datasets: list[PgenData],
    Y: ArrayLike,
    X: ArrayLike | None,
    phenotypes: list[str],
    train_mask: ArrayLike,
    test_mask: ArrayLike,
    h2_prior: ArrayLike,
    trait_type: str | TraitType,
    loocv: bool = False,
    B: int = 1000,
    idx_sample: ArrayLike | None = None,
    variants: list[str] | None = None,
    level0_dir: str | None = None,
    prune_blocks: bool = True,
    key: ArrayLike | None = None,
    memory_mode: str = "standard",
) -> dict[str, DataFrame]:
    """Perform agricola step 1

    Args:
        datasets: A list of PgenData objects (either single object or one
            per-chromosome)
        Y: A (N, P) jax array of phenotypes
        X: A (N, C) jax array of covariates (no intercept)
        phenotypes: A list of phenotype names, ordered as the columns of Y
        train_mask: A (N, K) jax array indicating training set status for each set k in 1, ..., K
        test_mask: A (N, K) jax array indicating test set status for each set k in 1, ..., K
        h2_prior: A 1D jax array of prior values for snp heritability
        trait_type: Either "qt" or "bt"
        loocv: A boolean indicating whether to perform LOOCV instead of standard
            cross validation. Ignored for trait_type="qt".
        B: The number of variants per block
        idx_sample: An optional (N_sub,) jax array with indices of samples to include
        variants: A list of variant IDs to include in the analysis. If not provided,
            all variants are used
        level0_dir: The directory where level 0 predictions are written
        prune_blocks: Whether to sample variants in a dataset in level 0 so that
            n_variants (mod B) = 0. This will improve speed with JIT compilations
        key: Optional JAX PRNG key used for level 0 variant subsampling

    Returns:
        A dict where keys are chromosomes and values are (N, P) pandas DataFrames
            of level 1 predictions
    """
    directory_context = (
        tempfile.TemporaryDirectory() if level0_dir is None else nullcontext(level0_dir)
    )
    with directory_context as working_dir:
        level0_path = Path(working_dir)
        level0_path.mkdir(parents=True, exist_ok=True)

        level0_files = level0(
            datasets,
            Y,
            X,
            phenotypes,
            train_mask,
            test_mask,
            h2_prior,
            B,
            idx_sample,
            variants,
            str(level0_path),
            prune_blocks,
            key,
            memory_mode,
        )

        return level1(
            level0_files,
            Y,
            X,
            phenotypes,
            train_mask,
            test_mask,
            h2_prior,
            trait_type,
            loocv,
            memory_mode,
        )

def step2(
    datasets: list[LancData],
    Y: ArrayLike,
    X: ArrayLike | None,
    step1_predictions: dict[str, pd.DataFrame] | None,
    outdir: str | Path,
    phenotypes: list[str],
    trait_type: str | TraitType = TraitType.QT,
    test_type: str | TestType = TestType.SCORE,
    chrom: str | None = None,
    B: int = 1000,
    min_ac: int = 1,
    idx_sample: ArrayLike | None = None,
    variants: list[str] | None = None,
    adjust_lanc: bool = True,
    impute: bool = False,
    overwrite: bool = True,
    partition_phenotype: bool = True,
    max_rows: int | None = None,
) -> None:
    """Perform agricola step 2

    Args:
        datasets: A list of LancData objects (either single object or one
            per-chromosome)
        Y: A (N, P) jax array of outcomes
        X: A (N, C) jax array of covariates
        step1_predictions: An optional dict with LOCO linear predictions from step 1.
            The values are (N, P) NumPy arrays
        outdir: Outputs will be written to {output_prefix}_{phenotype}.parquet
        phenotypes: A list of phenotype names
        trait_type: either "qt" or "bt"
        test_type: Either "score" or "wald"
        B: The block size (max number of variants to read at once)
        min_ac: the minimum allele count threshold
        idx_sample: An optional numpy array with ordered indices of samples (in
            the psam file) to retain
        variants: An optional list of variant IDs to retain
        adjust_lanc: A boolean indicating whether to adjust tests for local ancestry
        impute: Whether to impute the phenotype. Much faster, but only available
            for qt traits. If all phenotypes are non-missing, this is ignored.
        overwrite: Whether to overwrite the outdir if it already exists
        partition_phenotype: Whether to partition output parquet files by phenotype
        max_rows: Max number of rows/variants per phenotype to keep in memory
            before writing an output file. Defaults to 5000000 / len(phenotypes)
    """
    ## Create writer
    outdir_path = Path(outdir)
    if overwrite and outdir_path.exists():
        shutil.rmtree(outdir_path)

    if max_rows:
        max_rows_total = max_rows * len(phenotypes)
    else:
        max_rows_total = 5000000

    with ParquetRotatingWriter(outdir_path, partition_phenotype, max_rows_total) as writer:
        if impute:
            M = jnp.ones(shape=jnp.asarray(Y).shape)
        else:
            M = (~jnp.isnan(jnp.asarray(Y))).astype(float)

        Y, X, step1_predictions_np, idx_sample, test_type_enum, trait_type_enum = (
            validate_step2_inputs(
                datasets,
                Y,
                X,
                phenotypes,
                step1_predictions,
                B,
                idx_sample,
                variants,
                test_type,
                trait_type,
            )
        )

        ## Adjust phenotype for covariates to match step 1
        if trait_type_enum == TraitType.QT:
            Q, _ = jnp.linalg.qr(X, mode="reduced")
            Y = stdize(Y - (Q @ (Q.T @ Y)))
            if (M == 1).all():
                impute = True
        elif impute:
            raise ValueError("impute must be False for binary traits")

        time_total_start = time.perf_counter()
        for dataset in datasets:
            pgen_path = dataset.plink_prefix + ".pgen"
            logger.info("Testing associations for file: %s", pgen_path)

            time_ds_start = time.perf_counter()
            _step2_dataset(
                dataset,
                writer,
                Y,
                M,
                step1_predictions_np,
                X,
                idx_sample,
                phenotypes,
                trait_type_enum,
                test_type_enum,
                chrom,
                B,
                min_ac,
                variants,
                adjust_lanc,
                impute,
            )
            time_ds = str(timedelta(seconds=int(time.perf_counter() - time_ds_start)))
            logger.info("Elapsed time: %s", time_ds)

        time_total = str(timedelta(seconds=int(time.perf_counter() - time_total_start)))
        logger.info("Step 2 completed in: %s", time_total)