Borzoi prediction from personal genome - jupyter notebook

try:
    import kipoiseq
except ImportError:
    %pip install kipoiseq
from kipoiseq import Interval
try:
    import cyvcf2
except ImportError:
    %pip install cyvcf2
import os
import time
import io
import gzip

#os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

import h5py
import numpy as np
import pandas as pd
import tensorflow as tf

import baskerville
from baskerville import seqnn
from baskerville import gene as bgene
from baskerville import dna

import json

import pysam

import pyfaidx

import matplotlib.pyplot as plt
import matplotlib.patches as patches

print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
gpu_device = tf.config.list_physical_devices('GPU')[0]
tf.config.experimental.set_memory_growth(gpu_device, True)

PRE = '/Users/haekyungim/Library/CloudStorage/Box-Box/LargeFiles/imlab-data/data-Github/web-data/web-GENE-46100'
CONTENT_DIR = PRE + '/borzoi'
# Create borzoi/data directory if it doesn't exist
if not os.path.exists(os.path.join(CONTENT_DIR)):
    !mkdir {CONTENT_DIR}
if not os.path.exists(os.path.join(CONTENT_DIR, '/data')):
    !mkdir {CONTENT_DIR}/data

# Create model file structure
saved_models_path = os.path.join(CONTENT_DIR, 'data/saved_models')
if not os.path.exists(saved_models_path):
    !mkdir {saved_models_path}
if not os.path.exists(os.path.join(saved_models_path, 'f0')):
    !mkdir {saved_models_path}/f0
if not os.path.exists(os.path.join(saved_models_path, 'f1')):
    !mkdir {saved_models_path}/f1
if not os.path.exists(os.path.join(saved_models_path, 'f2')):
    !mkdir {saved_models_path}/f2
if not os.path.exists(os.path.join(saved_models_path, 'f3')):
    !mkdir {saved_models_path}/f3
    
#%cd {CONTENT_DIR}/data

#DELET
os.path.join(CONTENT_DIR, 'data/hg38.fa')

#Download model weights
if not os.path.exists(os.path.join(saved_models_path, 'f0', 'model0_best.h5')):
    !wget https://storage.googleapis.com/seqnn-share/borzoi/f0/model0_best.h5 -O {saved_models_path}/f0/model0_best.h5
if not os.path.exists(os.path.join(saved_models_path, 'f1', 'model0_best.h5')): 
    !wget https://storage.googleapis.com/seqnn-share/borzoi/f1/model0_best.h5 -O {saved_models_path}/f1/model0_best.h5
if not os.path.exists(os.path.join(saved_models_path, 'f2', 'model0_best.h5')):
    !wget https://storage.googleapis.com/seqnn-share/borzoi/f2/model0_best.h5 -O {saved_models_path}/f2/model0_best.h5
if not os.path.exists(os.path.join(saved_models_path, 'f3', 'model0_best.h5')):
    !wget https://storage.googleapis.com/seqnn-share/borzoi/f3/model0_best.h5 -O {saved_models_path}/f3/model0_best.h5

#Download and uncompress annotation files
if not os.path.exists(os.path.join(CONTENT_DIR, 'data', 'gencode41_basic_nort.gtf')):
    !wget  -O - https://storage.googleapis.com/seqnn-share/helper/gencode41_basic_nort.gtf.gz | gunzip -c > {os.path.join(CONTENT_DIR, 'data', 'gencode41_basic_nort.gtf')}
if not os.path.exists(os.path.join(CONTENT_DIR, 'data', 'gencode41_basic_protein_splice.csv.gz')):
    !wget https://storage.googleapis.com/seqnn-share/helper/gencode41_basic_protein_splice.csv.gz -O {os.path.join(CONTENT_DIR, 'data', 'gencode41_basic_protein_splice.csv.gz')}
if not os.path.exists(os.path.join(CONTENT_DIR, 'data', 'polyadb_human_v3.csv.gz')):
    !wget https://storage.googleapis.com/seqnn-share/helper/polyadb_human_v3.csv.gz -O {os. path.join(CONTENT_DIR, 'data', 'polyadb_human_v3.csv.gz')}

#fasta_file = '/content/data/genome.fa'
fasta_file = os.path.join(CONTENT_DIR, 'data/hg38.fa')
if not os.path.exists(fasta_file):
    !wget -O - http://hgdownload.cse.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz | gunzip -c > {fasta_file}

### download vcf file and index for chromosome 22
VCF_FILE = os.path.join(CONTENT_DIR, 'data/ALL.chr22.shapeit2_integrated_SNPs_v2a_27022019.GRCh38.phased.gz')
if not os.path.exists(VCF_FILE):
    !wget https://uchicago.box.com/shared/static/u526pjh29w93ufn65yomlch4z8lkb1sp.gz -O {VCF_FILE}

#!ls /content/data
!ls {CONTENT_DIR}/data

%cd {CONTENT_DIR}

models_path = CONTENT_DIR + '/data/saved_models/'
params_file = GITHUB_DIR + '/borzoi/examples/params_pred.json'
targets_file = GITHUB_DIR + '/borzoi/examples/targets_human.txt'

splice_df = pd.read_csv(CONTENT_DIR+'/data/gencode41_basic_protein_splice.csv.gz', sep='\t', compression='gzip')
print("len(splice_df) = " + str(len(splice_df)))

transcriptome = bgene.Transcriptome(CONTENT_DIR + '/data/gencode41_basic_nort.gtf')

class FastaStringExtractor:

    def __init__(self, fasta_file):
        self.fasta = pyfaidx.Fasta(fasta_file)
        self._chromosome_sizes = {k: len(v) for k, v in self.fasta.items()}
    #import pd.Interval as Interval

    def extract(self, interval: Interval, **kwargs) -> str:
        # Truncate interval if it extends beyond the chromosome lengths.
        chromosome_length = self._chromosome_sizes[interval.chrom]
        trimmed_interval = Interval(interval.chrom,
                                    max(interval.start, 0),
                                    min(interval.end, chromosome_length),
                                    )
        # pyfaidx wants a 1-based interval
        sequence = str(self.fasta.get_seq(trimmed_interval.chrom,
                                          trimmed_interval.start + 1,
                                          trimmed_interval.stop).seq).upper()
        # Fill truncated values with N's.
        pad_upstream = 'N' * max(-interval.start, 0)
        pad_downstream = 'N' * max(interval.end - chromosome_length, 0)
        return pad_upstream + sequence + pad_downstream # returns padded sequence

    def close(self):
        return self.fasta.close()

# Make one-hot coded sequence (modified from borzoi)
def make_seq_1hot(sequence):
    return dna.dna_1hot(sequence).astype("float32")

def predict_tracks(models, sequence_one_hot):

  predicted_tracks = []
  for fold_ix in range(len(models)):

    yh = models[fold_ix](sequence_one_hot[None, ...])[:, None, ...].astype("float16")

    predicted_tracks.append(yh)

  predicted_tracks = np.concatenate(predicted_tracks, axis=1)

  return predicted_tracks

def vcf_to_seq_faster(target_interval, individual, vcf_file, fasta_extractor):
  # target_inverval is a kipoiseq interval, e.g. kipoiseq.Interval("chr22", 18118779, 18145669)
  # individual is the id of the individual in the vcf file

  target_fa = fasta_extractor.extract(target_interval.resize(SEQUENCE_LENGTH))
  window_coords = target_interval.resize(SEQUENCE_LENGTH)
      # two haplotypes per individual
  haplo_1 = list(target_fa[:])
  haplo_2 = list(target_fa[:])

  # Open the VCF file
  vcf_reader = cyvcf2.VCF(vcf_file)

  # Specific genomic region
  CHROM = window_coords.chrom
  START = max(1,window_coords.start)
  END = min(window_coords.end, fasta_extractor._chromosome_sizes[CHROM]-1)

  count1 = 0
  count2 = 0

  # Iterate through variants in the specified region
  for variant in vcf_reader(CHROM + ':' + str(START) + '-' + str(END)):
      # Access the individual's genotype using index (0-based) of the sample
      individual_index = vcf_reader.samples.index(individual)
      genotype = variant.genotypes[individual_index]
      ALT=variant.ALT[0]
      REF=variant.REF
      POS=variant.POS
      if REF == target_fa[POS - window_coords.start - 1]:
        if genotype[0] == 1:
          haplo_1[POS - window_coords.start - 1] = ALT
          count1 = count1 + 1
        if genotype[1] == 1:
          haplo_2[POS - window_coords.start - 1] = ALT
          count2 = count2 + 1
      else:
        print("ERROR: REF in vcf "+ REF + "!= REF from the build" + target_fa[POS - window_coords.start - 1])
  print("number of changes haplo1:")
  print(count1)
  print("number of changes haplo2:")
  print(count2)
  return haplo_1, haplo_2

def make_prediction(gene, interval, haplo, sequence_one_hot, window = 131072): # snp_pos, alt_allele
    with tf.device('/GPU:0'):
        search_gene = gene # gene to predict

        resized_interval = interval.resize(SEQUENCE_LENGTH)
        start = resized_interval.start
        end = resized_interval.end
        center_pos = start + SEQUENCE_LENGTH//2 # figure center position

        chrom = resized_interval.chr # chromosome

        #Get exon bin range
        gene_keys = [gene_key for gene_key in transcriptome.genes.keys() if search_gene in gene_key]

        gene = transcriptome.genes[gene_keys[0]]

        #Determine output sequence start
        seq_out_start = start + seqnn_model.model_strides[0]*seqnn_model.target_crops[0]
        seq_out_len = seqnn_model.model_strides[0]*seqnn_model.target_lengths[0]

        #Determine output positions of gene exons
        gene_slice = gene.output_slice(seq_out_start, seq_out_len, seqnn_model.model_strides[0], False)

        #Make predictions
        #y_wt
        return predict_tracks(models, sequence_one_hot)

# Get indexes of track
def get_track_idx(tracks):
  track_idx = []
  targets_df['local_index'] = np.arange(len(targets_df))
  for track in tracks:
    track_idx.append(targets_df.loc[targets_df['description'] == track]['local_index'].tolist())
  return track_idx

def plot_tracks(prediction, interval, haplo, anno_df=None, tracks=[],
                log = [False, False, False],
                sqrt_scale = [False, False, False]):
  # Gene slice not included
  bin_size = 32
  pad = 16

  resized_interval = interval.resize(SEQUENCE_LENGTH)
  start = resized_interval.start
  end = resized_interval.end
  center_pos = start + SEQUENCE_LENGTH//2 # figure center position
  window=131072

  plot_start = center_pos - window // 2
  plot_end = center_pos + window // 2

  plot_start_bin = (plot_start - start) // bin_size - pad
  plot_end_bin = (plot_end - start) // bin_size - pad

  track_idx = get_track_idx(tracks)
  track_scales = [0.01, 0.01, 0.01]
  track_transforms = [3./4., 3./4., 3./4.]
  soft_clips = [384., 384., 384.]

  # Get annotation positions
  anno_poses = []
  if anno_df is not None:
      anno_poses = anno_df.query(
          "chrom == '"
          + interval.chr
          + "' and position_hg38 >= "
          + str(plot_start)
          + " and position_hg38 < "
          + str(plot_end)
        )["position_hg38"].values.tolist()
  # BIG plot
  fig, axes = plt.subplots(len(tracks), figsize=(20, 1.5 * len(tracks)), sharex = True)

  for ax, track, track_index, track_scale, track_transform, clip_soft in zip(axes,
        tracks, track_idx, track_scales, track_transforms, soft_clips):
    # Plot track densities
    y = np.array(np.copy(prediction), dtype = np.float32)
    y = np.mean(y[..., track_index], axis=(0, 1, 3))
    y = y[plot_start_bin:plot_end_bin]
    signal = np.max(y)
    if log:
      y = np.log2(y + 1.0)
    elif sqrt_scale:
      y = np.sqrt(y + 1.0)

    max_y = np.max(y)
    print("Max signal for "+ str(track) + " = " + str(round(signal, 4)))
    print("Max transformed signal for "+ str(track) + " = " + str(round(max_y, 4)))
    print("---")
    ax.bar(np.arange(plot_end_bin - plot_start_bin) + plot_start_bin,
            y,
            width=1.0,
            color="green",
            alpha=0.5)
            # label="Prediction",)
    # X axis tick annotations
    xtick_vals = []
    for _, anno_pos in enumerate(anno_poses):
      pas_bin = int((anno_pos - start) // 32) - 16
      xtick_vals.append(pas_bin)
      bin_end = pas_bin + 3 - 0.5
      bin_start = bin_end - 5

      ax.axvline(x=pas_bin,
                color="cyan",
                linewidth=2,
                alpha=0.5,
                linestyle="-",
                zorder=-1,)

    plt.xlim(plot_start_bin, plot_end_bin - 1)
    plt.xticks([], [])
    plt.yticks([], [])

    ax.set(ylabel = "Signal(log)" if log else "Signal")# , fontsize = 8)
    ax.set_title("Track: " + str(track)+ " in "+ str(haplo), fontsize=10)
    # plt.legend(fontsize=8)
    plt.tight_layout()
  ax.set_xlabel(
            interval.chr
            + ":"
            + str(plot_start)
            + "-"
            + str(plot_end)
            + " ("
            + str(window)
            + "bp window)")
            # fontsize=8,)
  plt.show()

fasta_extractor = FastaStringExtractor(fasta_file)