<jupyter_start><jupyter_text>**Step 1**<jupyter_code>WITH_COVARIATE = True
import glow
glow.register(spark, new_session=False)

import pandas as pd 
from pathlib import Path

from pyspark.sql.functions import col

import xarray as xr
import zarr
# Download the zip file mentioned in the link below to target directory and modify the path below as needed:
# https://github.com/pystatgen/sgkit/issues/622
# Direct link: https://github.com/pystatgen/sgkit/files/6779496/regenie_simulations.zip
dsdir = Path('../../../../sgkit/tests/test_regenie/dataset/sim_sm_02')

genotypes_vcf = dsdir / 'genotypes.bed'
covariates_csv = dsdir / 'covariates.csv'
continuous_phenotypes_csv = dsdir / 'traits.csv'
variants = spark.read.format('plink').load(str(genotypes_vcf), famDelimiter='\t')
genotypes = variants.withColumn('values', glow.mean_substitute(glow.genotype_states(col('genotypes'))))

label_df = pd.read_csv(continuous_phenotypes_csv, index_col='sample_id') #[['Continuous_Trait_1', 'Continuous_Trait_2']]
label_df = (label_df - label_df.mean()) / label_df.std(ddof=0)

if WITH_COVARIATE is True:
    covariate_df = pd.read_csv(covariates_csv, index_col='sample_id')
    covariate_df = (covariate_df - covariate_df.mean()) / covariate_df.std(ddof=0)
    covariate_suffix = ""
else:
    covariate_df = pd.DataFrame()
    covariate_suffix = "_nocovariate"
variants_per_block = 10
sample_block_count = 5

sample_ids = glow.get_sample_ids(genotypes)
sample_ids = [sample_id[2:] for sample_id in sample_ids]

block_df, sample_blocks = glow.block_variants_and_samples(genotypes, sample_ids, variants_per_block, sample_block_count)
reduction = glow.RidgeReduction(block_df, label_df, sample_blocks, covariate_df)
reduced_block_df = reduction.fit_transform()
regression = glow.RidgeRegression.from_ridge_reduction(reduction)
model_df, cv_df = regression.fit()
continuous_offsets = regression.transform_loco()<jupyter_output><jupyter_text>**Export Offsets**<jupyter_code>co_df = continuous_offsets.reset_index(level=['sample_id', 'contigName'])
co_df['sample_id']=co_df['sample_id'].apply(
    lambda m: int(str(m)[-2:])-1
)
co_df['contigName'] = co_df['contigName'].apply(
    lambda m: int(str(m))-1
)
co_df = co_df.rename(columns= {
    'sample_id': 'samples',
    'contigName': 'contigs'
})

co_df = co_df.set_index(['samples', 'contigs'])


co = co_df.to_xarray()

co=co.assign(regenie_loco_prediction=xr.concat([co.Y0000, co.Y0001, co.Y0002, co.Y0003, co.Y0004], 'outcomes'))
co['regenie_loco_prediction'] = co['regenie_loco_prediction'].transpose()


glow_offset_path = Path(f'../../../../sgkit/tests/test_regenie/dataset/sim_sm_02/glow_offsets{covariate_suffix}.zarr.zip')
store = zarr.ZipStore(str(glow_offset_path), mode='w')
co.to_zarr(store)
store.close()<jupyter_output><empty_output><jupyter_text>**Step 2**<jupyter_code>genotypes = (glow.transform('split_multiallelics', variants)
    .withColumn('gt', glow.mean_substitute(glow.genotype_states(col('genotypes'))))
    .select('contigName', 'start', 'names', 'gt')
    .cache())
lin_reg_df = glow.gwas.linear_regression(
    genotypes,
    label_df,
    covariate_df,
    offset_df=continuous_offsets,
    values_column='gt'
  )
output_path = Path('../../../../sgkit/tests/test_regenie/result/sim_sm_02-wgr_02')
# Convert to pandas
wgr_gwas = lin_reg_df.toPandas()

path = output_path / f'gwas_loco{covariate_suffix}.csv'
wgr_gwas.to_csv(path, index=False)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14514252472367656, 'fraction_numerical': 0.014834205933682374, 'mean_word_length': 6.379828326180258, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 12, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '176136', 'license': 'Apache-2.0', 'n_tokens_mistral': 1381, 'n_tokens_neox': 1299, 'n_words': 190, 'path': 'validation/gwas/method/regenie_loco_regression/GlowGR_continuous.ipynb', 'repo_name': 'pystatgen/sgk', 'size': '3438'}