def test_no_values(spark):
variant_df = spark.createDataFrame([__construct_row([0, 1])]).limit(0)
sample_ids = ["a", "b"]
with pytest.raises(Exception):
functions.block_variants_and_samples(variant_df,
sample_ids,
variants_per_block=10,
sample_block_count=2)
def test_missing_values(spark):
variant_df = spark.read.format("vcf").load(
"test-data/combined.chr20_18210071_18210093.g.vcf")
sample_ids = ["HG00096", "HG00268", "NA19625"]
with pytest.raises(AnalysisException):
functions.block_variants_and_samples(variant_df,
sample_ids,
variants_per_block=10,
sample_block_count=2)
def test_mismatch_num_values_sample_ids(spark):
variant_df = spark.createDataFrame(
[__construct_row([0, 1]),
__construct_row([1, 1])])
sample_ids = ["a", "b", "c"]
with pytest.raises(Exception):
functions.block_variants_and_samples(variant_df,
sample_ids,
variants_per_block=10,
sample_block_count=2)
def test_inconsistent_num_values(spark):
variant_df = spark.createDataFrame(
[__construct_row([0, 1, 1]),
__construct_row([1, 2])])
sample_ids = ["a", "b", "c"]
block_df, index_map = functions.block_variants_and_samples(
variant_df, sample_ids, variants_per_block=10, sample_block_count=2)
with pytest.raises(Exception):
block_df.collect(
) # Number of values is checked lazily within the transformer
def test_block_variants_and_samples(spark):
variant_df = spark.read.format("vcf") \
.load("test-data/combined.chr20_18210071_18210093.g.vcf") \
.withColumn("values", expr("genotype_states(genotypes)"))
sample_ids = ["HG00096", "HG00268", "NA19625"]
block_gt, index_map = functions.block_variants_and_samples(
variant_df, sample_ids, variants_per_block=10, sample_block_count=2)
expected_block_gt = glow.transform("block_variants_and_samples",
variant_df,
variants_per_block=10,
sample_block_count=2)
assert block_gt.collect() == expected_block_gt.collect()
assert index_map == {"1": ["HG00096", "HG00268"], "2": ["NA19625"]}
def run(
plink_path: str,
traits_path: str,
covariates_path: str,
variants_per_block: int,
sample_block_count: int,
output_dir: str,
plink_fam_sep: str = "\t",
plink_bim_sep: str = "\t",
alphas: Optional[list] = None,
contigs: List[str] = None,
):
"""Run Glow WGR"""
output_path = Path(output_dir)
if output_path.exists():
shutil.rmtree(output_path)
output_path.mkdir(parents=True, exist_ok=False)
if alphas is None:
alphas = np.array([])
else:
alphas = np.array(alphas).astype(float)
spark = spark_session()
logger.info(
f"Loading PLINK dataset at {plink_path} (fam sep = {plink_fam_sep}, bim sep = {plink_bim_sep}, alphas = {alphas})"
)
df = (spark.read.format("plink").option(
"bimDelimiter",
plink_bim_sep).option("famDelimiter", plink_fam_sep).option(
"includeSampleIds", True).option("mergeFidIid",
False).load(plink_path))
variant_df = df.withColumn(
"values", mean_substitute(genotype_states(F.col("genotypes")))).filter(
F.size(F.array_distinct("values")) > 1)
if contigs is not None:
variant_df = variant_df.filter(F.col("contigName").isin(contigs))
sample_ids = get_sample_ids(variant_df)
logger.info(
f"Found {len(sample_ids)} samples, first 10: {sample_ids[:10]}")
###########
# Stage 1 #
###########
logger.info(HR)
logger.info("Calculating variant/sample block info")
block_df, sample_blocks = block_variants_and_samples(
variant_df,
sample_ids,
variants_per_block=variants_per_block,
sample_block_count=sample_block_count,
)
label_df = pd.read_csv(traits_path, index_col="sample_id")
label_df = (label_df - label_df.mean()) / label_df.std(ddof=0)
logger.info(HR)
logger.info("Trait info:")
logger.info(_info(label_df))
cov_df = pd.read_csv(covariates_path, index_col="sample_id")
cov_df = (cov_df - cov_df.mean()) / cov_df.std(ddof=0)
logger.info(HR)
logger.info("Covariate info:")
logger.info(_info(cov_df))
stack = RidgeReducer(alphas=alphas)
reduced_block_df = stack.fit_transform(block_df, label_df, sample_blocks,
cov_df)
logger.info(HR)
logger.info("Stage 1: Reduced block schema:")
logger.info(_schema(reduced_block_df))
path = output_path / "reduced_blocks.parquet"
reduced_block_df.write.parquet(str(path), mode="overwrite")
logger.info(f"Stage 1: Reduced blocks written to {path}")
# Flatten to scalars for more convenient access w/o Spark
flat_reduced_block_df = spark.read.parquet(str(path))
path = output_path / "reduced_blocks_flat.csv.gz"
flat_reduced_block_df = _flatten_reduced_blocks(flat_reduced_block_df)
flat_reduced_block_df = flat_reduced_block_df.toPandas()
flat_reduced_block_df.to_csv(path, index=False)
# flat_reduced_block_df.write.parquet(str(path), mode='overwrite')
logger.info(f"Stage 1: Flattened reduced blocks written to {path}")
###########
# Stage 2 #
###########
# Monkey-patch this in until there's a glow release beyond 0.5.0
if glow_version != "0.5.0":
raise NotImplementedError(
f"Must remove adjustements for glow != 0.5.0 (found {glow_version})"
)
# Remove after glow update
RidgeRegression.transform_loco = transform_loco
estimator = RidgeRegression(alphas=alphas)
model_df, cv_df = estimator.fit(reduced_block_df, label_df, sample_blocks,
cov_df)
logger.info(HR)
logger.info("Stage 2: Model schema:")
logger.info(_schema(model_df))
logger.info("Stage 2: CV schema:")
logger.info(_schema(cv_df))
y_hat_df = estimator.transform(reduced_block_df, label_df, sample_blocks,
model_df, cv_df, cov_df)
logger.info(HR)
logger.info("Stage 2: Prediction info:")
logger.info(_info(y_hat_df))
logger.info(y_hat_df.head(5))
path = output_path / "predictions.csv"
y_hat_df.reset_index().to_csv(path, index=False)
logger.info(f"Stage 2: Predictions written to {path}")
y_hat_df_loco = estimator.transform_loco(reduced_block_df, label_df,
sample_blocks, model_df, cv_df,
cov_df)
path = output_path / "predictions_loco.csv"
y_hat_df_loco.reset_index().to_csv(path, index=False)
logger.info(f"Stage 2: LOCO Predictions written to {path}")
###########
# Stage 3 #
###########
# Do this to correct for the error in Glow at https://github.com/projectglow/glow/issues/257
if glow_version != "0.5.0":
raise NotImplementedError(
f"Must remove adjustements for glow != 0.5.0 (found {glow_version})"
)
cov_arr = cov_df.to_numpy()
cov_arr = cov_arr.T.ravel(order="C").reshape(cov_arr.shape)
# Convert the pandas dataframe into a Spark DataFrame
adjusted_phenotypes = reshape_for_gwas(spark, label_df - y_hat_df)
# Run GWAS w/o LOCO (this could be for a much larger set of variants)
wgr_gwas = (variant_df.withColumnRenamed("values", "callValues").crossJoin(
adjusted_phenotypes.withColumnRenamed(
"values", "phenotypeValues")).select(
"start",
"names",
"label",
expand_struct(
linear_regression_gwas(F.col("callValues"),
F.col("phenotypeValues"),
F.lit(cov_arr))),
))
logger.info(HR)
logger.info("Stage 3: GWAS (no LOCO) schema:")
logger.info(_schema(wgr_gwas))
# Convert to pandas
wgr_gwas = wgr_gwas.toPandas()
logger.info(HR)
logger.info("Stage 3: GWAS (no LOCO) info:")
logger.info(_info(wgr_gwas))
logger.info(wgr_gwas.head(5))
path = output_path / "gwas.csv"
wgr_gwas.to_csv(path, index=False)
logger.info(f"Stage 3: GWAS (no LOCO) results written to {path}")
logger.info(HR)
logger.info("Done")
# TODO: Enable this once WGR is fully released
# See: https://github.com/projectglow/glow/issues/256)
# Run GWAS w/ LOCO
adjusted_phenotypes = reshape_for_gwas(spark, label_df - y_hat_df_loco)
wgr_gwas = (variant_df.withColumnRenamed("values", "callValues").join(
adjusted_phenotypes.withColumnRenamed("values", "phenotypeValues"),
["contigName"],
).select(
"contigName",
"start",
"names",
"label",
expand_struct(
linear_regression_gwas(F.col("callValues"),
F.col("phenotypeValues"), F.lit(cov_arr))),
))
# Convert to pandas
wgr_gwas = wgr_gwas.toPandas()
logger.info(HR)
logger.info("Stage 3: GWAS (with LOCO) info:")
logger.info(_info(wgr_gwas))
logger.info(wgr_gwas.head(5))
path = output_path / "gwas_loco.csv"
wgr_gwas.to_csv(path, index=False)
logger.info(f"Stage 3: GWAS (with LOCO) results written to {path}")
logger.info(HR)
logger.info("Done")