#!/usr/bin/env python
# coding: utf-8

# このノートブックでは、
# 染色体ごとに保存されているゲノムデータ (VCF形式)をマージし、matrix table 形式で保存します。
# 
# その後、imputation quality の低いバリアント (DR2 < 0.3) を除外します。
# 
# またその後のゲノムデータを用いて、PRS の計算手順を説明します。
# 
# 2種類の脳梗塞 PRS モデルを事例として用います。
# 
# - PGS002724 モデルに含まれるバリアント数 1213574個
# - PGS002725 モデルに含まれるバリアント数 6010730個
# 

# ## このノートブックを実行する前に下記をご注意ください
# 
# - このノートブック用の jupyter サービスを立ち上げる前に `export PYSPARK_SUBMIT_ARGS='--driver-memory 48g --executor-memory 48g pyspark-shell'` のようにこの計算用のメモリを大きく確保してください。(この場合48ギガのメモリを割り当てています。) これを行っていただかないと、いずれかのセルで OutOfMemory エラーが発生し、それ以降のセル実行は機能しなくなります。
# - hail の背後では spark が働いており、普段の jupyter notebook のようにセル実行を停止しても、その計算は動き続けます。その状態で新たなセルを実行すると予期せぬエラーが発生する場合があります。そのため「セルの実行がなかなか終わらないな」とお思いになられても、実行状態が完了するまで停止操作は行わないことをおすすめします。
# - 本ノートブックは https://github.com/hacchy1983/prs-on-hail-public に変更を加えたものになります。
# 

# ## Step 1 hail など、必要なモジュールを読み込みます
# 下記のコードを実行してください。 ページ上側のメニューバーにある 実行 ボタンを押下することで、実行することができます。

# In[1]:


import hail as hl
hl.init()
from hail.plot import show
from pprint import pprint
hl.plot.output_notebook()


# In[57]:


get_ipython().system('ls outputs/')


# ## Step 2 ゲノムデータのファイル形式を変換します
# 
# ゲノムデータは、1番染色体から22番染色体まで、染色体ごとに異なるファイルに保存されています。
# 例えば、1番染色体のゲノムデータは、 `outputs/chr1.beagle.vcf.gz` に保存されています。
# また、2番染色体のゲノムデータは、 `outputs/chr2.beagle.vcf.gz` に 22番染色体のゲノムデータは、 `outputs/chr22.beagle.vcf.gz` に保存されています。
# 
# ゲノムデータは、よく利用されるファイル形式である VCF フォーマットで保存されています。
# これを、matrix table 形式に変換して、保存します。
# 
# 下記コマンドを実行してください。
# このセルの実行にはかなりの待ち時間が生じますが「このノートブックを実行する前に下記をご注意ください」のように途中でセル実行を停止することはおすすめしません。

# In[58]:


for chr in range(1,23):
    infile = 'outputs/chr' + str(chr) + '.beagle.vcf.gz'
    outfile = 'outputs/chr' + str(chr) + '.mt'
    print(chr, infile, outfile)
    hl.import_vcf(infile, force_bgz=True).write(outfile, overwrite=True)


# ## Step 3 ゲノムデータをマージします
# 続いて、1番染色体から22番染色体の matrix table 形式のファイルを、ひとつの matrix table 形式のファイルにマージします。
# 
# 下記コマンドを実行してください。

# In[8]:


chr = 1
file = 'outputs/chr' + str(chr) + '.mt'
mt = hl.read_matrix_table(file)
print(chr, file, mt.count(), mt.count())
for chr in range(2,23):
    file = 'outputs/chr' + str(chr) + '.mt'
    tmpmt = hl.read_matrix_table(file)
    mt = mt.union_rows(tmpmt)
    print(chr, file, tmpmt.count(), mt.count())


# ゲノムデータのマージが完了しました。
# マージ後のゲノムデータのバリアント数を表示します。

# In[9]:


mt.count()


# これは、次のことを意味します。
# 
# - 研究対象者の人数が 2318 名
# - バリアントの個数が 30761499 個
# 
# ## Step 4 マージしたゲノムデータを保存します
# 今後のステップの実行時間を短縮するため、マージしたゲノムデータを保存し、再度読み込みます。
# 
# 下記のコードは、マージしたゲノムデータを outputs/chrAll.mt に保存します。
# 

# In[10]:


mt.write('outputs/chrAll.mt', overwrite=True)


# 下記のコードは、 `outputs/chrAll.mt` を読み込みます。

# In[2]:


mt = hl.read_matrix_table('outputs/chrAll.mt')


# ## Step 5 ゲノムデータに variantID を追加します
# 下記のコードは、ゲノムデータのバリアント情報に `variantID` を追加します。

# In[3]:


mt = mt.annotate_rows(variantID = (hl.str(mt.locus.contig) + ":" + hl.str(mt.locus.position)) )


# 下記のコードは、 `variantID` を追加した後のバリアント情報を表示します。
# `show(5)` は、先頭の 5 個のバリアントのみを表示する、ことを意味します。

# In[4]:


mt.rows().show(5)


# `variantID` のカラムが追加されていることが分かります。
# 
# ## Step 6 imputation quality に基づいてゲノムデータをフィルタリングします
# 
# mt には multi-allelic site が含まれます。次にその割合を確認してみます。

# In[5]:


mt1 = mt.filter_rows(hl.len(mt.info.DR2) == 1)
mt1.count()


# In[6]:


mtnot1 = mt.filter_rows(hl.len(mt.info.DR2) > 1)
mtnot1.count()


# 1% ほどが multi-allelic site であることがわかります。 1% は無視するにはやや多いですが、
# このチュートリアルでは内容をわかりやすくするために除外したもの(すなわち`mt1`)を今後用います。
# 下記のコードでは、各バリアントの `imputation quality` の分布を表示します。 

# In[7]:


p = hl.plot.histogram(mt1.info.DR2.first(), title='Imputation Quality Histogram', legend='Imputation Quality (DR2)')
show(p)


# `imputation quality` が低いバリアントが少しあることが分かります。
# 下記のコードは、`imputation quality` が低い（DR2 < 0.3）バリアントを除外します。

# In[8]:


mt1_filt = mt1.filter_rows(mt1.info.DR2.first()>=0.3)


# 下記のコードは、`imputation quality` が低いバリアントを除外した後の分布を表示します。

# In[12]:


p = hl.plot.histogram(mt1_filt.info.DR2.first(), title='Imputation Quality Histogram', legend='Imputation Quality (DR2)')
show(p)


# `imputation quality` が低いバリアントが除外されたことが分かります。
# 
# 下記のコードは、`imputation quality` が低いバリアントを除外した後のバリアントの個数を表示します。

# In[10]:


mt1_filt.count()


# (29799034, 2318) と表示されました。
# 
# これは、次のことを意味します。
# 
# 研究対象者の人数が 2318 名
# バリアントの個数が 29799034 個
# 下記のコードは、`imputation quality` が低いバリアントを除外した後の allele frequency の分布を表示します。

# In[11]:


p = hl.plot.histogram(mt1_filt.info.AF.first(), title='AF Histogram', legend='AF', bins=50)
show(p)


# AF<1% のバリアントが多くあることが分かります。

# ## Step 7 フィルタリング後のゲノムデータを保存します
# 下記のコードは、マージしたゲノムデータを `outputs/chrAll.filtered.mt` に保存します。
# ここまでのセルには実行完了までに時間がかかるものが存在します。
# ここをチェックポイントとして同じことに時間をかけずに再開できるようにしておきましょう。

# In[12]:


mt1_filt.write('outputs/chrAll.filtered.mt', overwrite=True)


# ## Step 8 ゲノムデータを読み込みます
# 下記のコードでは、Step 7 で保存したゲノムデータを読み込みます。

# In[ ]:


mt1_filt = hl.read_matrix_table('outputs/chrAll.filtered.mt')


# In[67]:


get_ipython().system('ls -alh outputs')


# ## Step 9-1 PRSモデルを読み込みます (PGS002724, PGS002725)
# 下記のコードを実行すると、`prs-models/PGS002724.txt` と `prs-models/PGS002724.txt` のデータが読み込まれます。

# In[34]:


get_ipython().system('wget https://ftp.ebi.ac.uk/pub/databases/spot/pgs/scores/PGS002724/ScoringFiles/PGS002724.txt.gz')


# In[35]:


get_ipython().system('wget https://ftp.ebi.ac.uk/pub/databases/spot/pgs/scores/PGS002725/ScoringFiles/PGS002725.txt.gz')


# In[36]:


get_ipython().system('gunzip PGS002724.txt.gz')


# In[37]:


get_ipython().system('gunzip PGS002725.txt.gz')


# In[38]:


get_ipython().system('ls')


# In[28]:


get_ipython().system('mkdir prs-models')


# In[39]:


get_ipython().system('mv PGS002724.txt prs-models/')


# In[40]:


get_ipython().system('mv PGS002725.txt prs-models/')


# In[13]:


get_ipython().system('ls prs-models')


# In[15]:


get_ipython().system('head prs-models/PGS002724.txt -n 30')


# In[14]:


model_PGS002724 = hl.import_table('prs-models/PGS002724.txt', impute=True, force=True, comment='#')


# In[15]:


model_PGS002725 = hl.import_table('prs-models/PGS002725.txt', impute=True, force=True, comment='#')


# 下記のコードは、PRSモデルに含まれるバリアントの個数を表示します。

# In[16]:


model_PGS002724.count()


# In[17]:


model_PGS002725.count()


# `1213574` と `6010730` と表示されました。
# これは、PRSモデル 002724 と 002725 に含まれるバリアントの個数がそれぞれ 1213574, 6010730 個であることを意味します。
# 
# 下記のコードは、読み込んだ PRS モデルの最初の 5 行を表示します。

# In[18]:


model_PGS002724.show(5)


# In[19]:


model_PGS002725.show(5)


# 下記のコードは、読み込んだ PRS モデルに `variantID` を追加します。

# In[20]:


model_PGS002724 = model_PGS002724.annotate(
    variantID = hl.str(model_PGS002724.chr_name) + ":" + hl.str(model_PGS002724.chr_position) 
)


# In[21]:


model_PGS002725 = model_PGS002725.annotate(
    variantID = hl.str(model_PGS002725.chr_name) + ":" + hl.str(model_PGS002725.chr_position) 
)


# 下記のコードは `variantID` を追加した後の最初の 5 行を表示します。

# In[22]:


model_PGS002724.show(5)


# In[23]:


model_PGS002725.show(5)


# `variantID` のカラムが追加されていることが分かります。

# ## Step 9-2 ゲノムデータとPRSモデルに共通するバリアントを抽出します
# 下記のコードは、PRSモデルのバリアント情報を variantID で検索できるようにします。

# In[24]:


model_PGS002724 = model_PGS002724.key_by('variantID')


# In[25]:


model_PGS002725 = model_PGS002725.key_by('variantID')


# In[26]:


mt1_filt.rows().show(5)


# 下記のコードは、ゲノムデータと PRS モデルに共通するバリアントを抽出します。

# In[27]:


mt_match_PGS002724 = mt1_filt.annotate_rows(**model_PGS002724[mt1_filt.variantID])
mt_match_PGS002724 = mt_match_PGS002724.filter_rows(hl.is_defined(mt_match_PGS002724.effect_weight))


# In[28]:


mt_match_PGS002725 = mt1_filt.annotate_rows(**model_PGS002725[mt1_filt.variantID])
mt_match_PGS002725 = mt_match_PGS002725.filter_rows(hl.is_defined(mt_match_PGS002725.effect_weight))


# In[29]:


mt_match_PGS002724.rows().show(5)


# In[27]:


model_PGS002724.filter(model_PGS002724.chr_name==1).count()


# In[37]:


model_PGS002725.filter(model_PGS002725.chr_name==1).count()


# 今後のステップの実行時間を短縮するため(ここを第2のチェックポイントとするため)、一旦抽出したゲノムデータを保存しておきます。
# 
# 下記のコードは、抽出したゲノムデータを `outputs/chrAll.filtered.matched_PGS002724.mt`, `outputs/chrAll.filtered.matched_PGS002725.mt` に保存します。
# 

# In[29]:


mt_match_PGS002724.write('outputs/chrAll.filtered.matched_PGS002724.mt', overwrite=True)


# In[38]:


mt_match_PGS002725.write('outputs/chrAll.filtered.matched_PGS002725.mt', overwrite=True)


# In[30]:


len(dict(mt_match_PGS002724.aggregate_rows(hl.agg.counter(mt_match_PGS002724.variantID))))


# In[31]:


len(dict(mt_match_PGS002725.aggregate_rows(hl.agg.counter(mt_match_PGS002725.variantID))))


# ## Step 9-3 PRSを計算して保存します
# 下記のコードは、ゲノムデータのアリル情報とPRSモデルのアリル情報を比較し、合致しているかをチェックします。

# In[32]:


flip_PGS002724 = hl.case().when( 
    (mt_match_PGS002724.effect_allele == mt_match_PGS002724.alleles[0])
    & (mt_match_PGS002724.other_allele == mt_match_PGS002724.alleles[1]), True ).when( 
    (mt_match_PGS002724.effect_allele == mt_match_PGS002724.alleles[1])
    & (mt_match_PGS002724.other_allele == mt_match_PGS002724.alleles[0]), False ).or_missing()


# In[33]:


flip_PGS002725 = hl.case().when( 
    (mt_match_PGS002725.effect_allele == mt_match_PGS002725.alleles[0]) 
    & (mt_match_PGS002725.other_allele == mt_match_PGS002725.alleles[1]), True ).when( 
    (mt_match_PGS002725.effect_allele == mt_match_PGS002725.alleles[1])
    & (mt_match_PGS002725.other_allele == mt_match_PGS002725.alleles[0]), False ).or_missing()


# In[34]:


mt_match_PGS002724 = mt_match_PGS002724.annotate_rows(flip=flip_PGS002724)


# In[35]:


mt_match_PGS002725 = mt_match_PGS002725.annotate_rows(flip=flip_PGS002725)


# 下記のコードは、各々のバリアントについて研究対象者の持っているアリル数（`mt_match_PGS002724.DS`）とバリアントの重み（`mt_match_PGS002724.effect_weight`）を掛け合わせて、ゲノムデータとPRSモデルの共通するバリアントについて足し合わせます。 これにより、PRSを計算することができます。

# In[36]:


prs_PGS002724 = hl.agg.sum(hl.float64(mt_match_PGS002724.effect_weight) * 
                    hl.if_else( mt_match_PGS002724.flip, 
                                2 - mt_match_PGS002724.DS.first(),
                                mt_match_PGS002724.DS.first()))


# In[38]:


mt_match_PGS002724 = mt_match_PGS002724.annotate_cols(prs=prs_PGS002724)


# PGS002725 についても同様のことを行ってみましょう。

# In[37]:


prs_PGS002725 = hl.agg.sum(hl.float64(mt_match_PGS002725.effect_weight) * 
                    hl.if_else( mt_match_PGS002725.flip, 
                                2 - mt_match_PGS002725.DS.first(),
                                mt_match_PGS002725.DS.first()))


# In[39]:


mt_match_PGS002725 = mt_match_PGS002725.annotate_cols(prs=prs_PGS002725)


# 下記のコードは、PRSの値を表示します。

# In[40]:


mt_match_PGS002724.cols().show(5)


# In[41]:


mt_match_PGS002725.cols().show(5)


# 下記のコードは、PRSの分布を表示します。

# In[42]:


p = hl.plot.histogram(mt_match_PGS002724.prs, title="PRS Histogram", legend="PGS002724", bins=20)
show(p)


# 下記のコードは、PRSの計算結果を chrAll.filtered.PGS002724.PRS.txt ファイルに保存します。

# In[43]:


mt_match_PGS002724.cols().export('chrAll.filtered.PGS002724.PRS.txt')


# PGS002725 についても同様のことを行ってみましょう。

# In[44]:


p = hl.plot.histogram(mt_match_PGS002725.prs, title="PRS Histogram", legend="PGS002725", bins=20)
show(p)


# In[45]:


mt_match_PGS002725.cols().export('chrAll.filtered.PGS002725.PRS.txt')


# ## Step 10 PRSのスコア値を比較します
# ここまでで 2 つの脳梗塞PRSモデルを用いて、PRSスコア値を計算しました。
# 
# このチュートリアルの最後に、そのPRSスコア値を比較してみます。
# 
# 下記のコードは、PRSモデル PGS002724 を用いて計算した PRS スコア値を読み込みます。

# In[46]:


prs_PGS002724 = hl.import_table('chrAll.filtered.PGS002724.PRS.txt', impute=True, force=True)


# In[47]:


prs_PGS002725 = hl.import_table('chrAll.filtered.PGS002725.PRS.txt', impute=True, force=True)


# 下記のコードは、PRS スコア値を subject ID (変数名 `s`) で検索できるようにします。

# In[48]:


prs_PGS002724 = prs_PGS002724.key_by('s')


# In[49]:


prs_PGS002725 = prs_PGS002725.key_by('s')


# 下記のコードは、2 種類の PRS スコア値を subject ID (変数名 s) で突合し、データマージを行います。

# In[50]:


prs_merge = prs_PGS002724.rename({'s':'subjectID', 'prs':'PGS002724'})


# In[51]:


prs_merge = prs_merge.annotate(PGS002725 = prs_PGS002725[prs_merge.subjectID].prs)


# 下記のコードは、データマージした結果を表示します。

# In[52]:


prs_merge.show(5)


# 下記のコードは、`PGS002724` と `PGS002725` のスコア値をプロットします。

# In[53]:


p = hl.plot.scatter(prs_merge.PGS002724, prs_merge.PGS002725, xlabel="PGS002724", ylabel="PGS002725")
show(p)


# 下記のコードは、PRS スコア値の相関係数を計算します。
# 
# `to_pandas()`関数を用いることで、`hail`特有のデータタイプから、`python` でよく使われる `pandas` ライブラリのデータフレームへと変換することができます。
# そうすることで、`hail`の関数だけでなく、`python` の様々な機能を使って分析することが可能になります。

# In[54]:


prs_merge_pandas = prs_merge.to_pandas()


# In[55]:


print(prs_merge_pandas.corr())


# この結果から、次のことが分かります。
# 
# - PGS002724のスコア値と PGS002725のスコア値の相関係数は　0.168966
# 
# ### 以上でこのチュートリアルは終了です

# In[ ]: