"""

Utilities for iterating constructing data sets and iterating over

DNA sequence data.

"""

import multiprocessing

import pandas as pd

import numpy as np

import functools

import math

from collections import defaultdict

from multiprocessing import Pool

import pyfasta

import pyBigWig

from pybedtools import Interval, BedTool

import logging

import tensorflow as tf

def filter_chromosomes(input_df, to_filter=None, to_keep=None):

"""

This function takes as input a pandas DataFrame

Parameters:

input_df (dataFrame): A pandas dataFrame, the first column is expected to

be a chromosome. Example: chr1.

to_filter (list): Default None (bool = False), will iterate over list

objects and filter the listed chromosomes.

( Default: None, i.e. this condition will not be triggered unless a list

is supplied)

to_keep (list): Default None, will iterate over list objects and only

retain the listed chromosomes.

Returns:

output_df (dataFrame): The filtered pandas dataFrame

"""

if to_filter:

output_df = input_df.copy()

for chromosome in to_filter:

# note: using the str.contains method to remove all

# contigs; for example: chrUn_JH584304

bool_filter = ~(output_df['chrom'].str.contains(chromosome))

output_df = output_df[bool_filter]

elif to_keep:

# keep only the to_keep chromosomes:

# note: this is slightly different from to_filter, because

# at a time, if only one chromosome is retained, it can be used

# sequentially.

filtered_chromosomes = []

for chromosome in to_keep:

filtered_record = input_df[(input_df['chrom'] == chromosome)]

filtered_chromosomes.append(filtered_record)

# merge the retained chromosomes

output_df = pd.concat(filtered_chromosomes)

else:

output_df = input_df

return output_df

def get_genome_sizes(genome_sizes_file, to_filter=None, to_keep=None):

"""

Loads the genome sizes file which should look like this:

chr1 45900011

chr2 10001401

...

chrX 9981013

This function parses this file, and saves the resulting intervals file

as a BedTools object.

"Random" contigs, chrUns and chrMs are filtered out.

Parameters:

genome_sizes_file (str): (Is in an input to the class,

can be downloaded from UCSC genome browser)

to_filter (list): Default None (bool = False), will iterate over list

objects and filter the listed chromosomes.

( Default: None, i.e. this condition will not be triggered unless a list

is supplied)

to_keep (list): Default None, will iterate over list objects and only

retain the listed chromosomes.

Returns:

A BedTools (from pybedtools) object containing all the chromosomes,

start (0) and stop (chromosome size) positions

"""

genome_sizes = pd.read_csv(genome_sizes_file, sep='\t',

header=None, names=['chrom', 'length'])

genome_sizes_filt = filter_chromosomes(genome_sizes, to_filter=to_filter,

to_keep=to_keep)

genome_bed_data = []

# Note: Modifying this to deal with unexpected (incorrect) edge case \

# BedTools shuffle behavior.

# While shuffling data, BedTools shuffle is placing certain windows at the \

# edge of a chromosome

# Why it's doing that is unclear; will open an issue on GitHub.

# It's probably placing the "start" co-ordinate within limits of the genome,

# with the end coordinate not fitting.

# This leads to the fasta file returning an incomplete sequence \

# (< 500 base pairs)

# This breaks the generator feeding into Model.fit.

# Therefore, in the genome sizes file, buffering 550 from the edges

# to allow for BedTools shuffle to place window without running of the

# chromosome.

for chrom, sizes in genome_sizes_filt.values:

genome_bed_data.append(Interval(chrom, 0 + 550, sizes - 550))

genome_bed_data = BedTool(genome_bed_data)

return genome_bed_data

def load_chipseq_data(chip_peaks_file, genome_sizes_file, to_filter=None,

to_keep=None):

"""

Loads the ChIP-seq peaks data.

The chip peaks file is an events bed file:

chr1:451350

chr2:91024

...

chrX:870000

This file can be constructed using a any peak-caller. We use multiGPS.

Also constructs a 1 bp long bedfile for each coordinate and a

BedTools object which can be later used to generate

negative sets.

"""

chip_seq_data = pd.read_csv(chip_peaks_file, delimiter=':', header=None,

names=['chrom', 'start'])

chip_seq_data['end'] = chip_seq_data['start'] + 1

chip_seq_data = filter_chromosomes(chip_seq_data, to_filter=to_filter,

to_keep=to_keep)

sizes = pd.read_csv(genome_sizes_file, names=['chrom', 'chrsize'],

sep='\t')

# filtering out any regions that are close enough to the edges to

# result in out-of-range windows when applying data augmentation.

chrom_sizes_dict = (dict(zip(sizes.chrom, sizes.chrsize)))

chip_seq_data['window_max'] = chip_seq_data['end'] + 500

chip_seq_data['window_min'] = chip_seq_data['start'] - 500

chip_seq_data['chr_limits_upper'] = chip_seq_data['chrom'].map(

chrom_sizes_dict)

chip_seq_data = chip_seq_data[chip_seq_data['window_max'] <=

chip_seq_data['chr_limits_upper']]

chip_seq_data = chip_seq_data[chip_seq_data['window_min'] >= 0]

chip_seq_data = chip_seq_data[['chrom', 'start', 'end']]

return chip_seq_data

def exclusion_regions(blacklist_file, chip_seq_data):

"""

This function takes as input a bound bed file (from multiGPS).

The assumption is that the bed file reports the peak center

For example: chr2 45 46

It converts these peak centers into 501 base pair windows, and adds them to

the exclusion list which will be used when constructing negative sets.

It also adds the mm10 blacklisted windows to the exclusion list.

Parameters:

blacklist_file (str): Path to the blacklist file.

chip_seq_data (dataFrame): The pandas chip-seq data loaded by load_chipseq_data

Returns:

exclusion_windows (BedTool): A bedtools object containing all exclusion windows.

bound_exclusion_windows (BedTool): A bedtool object containing only

those exclusion windows where there exists a binding site.

"""

temp_chip_file = chip_seq_data.copy() # Doesn't modify OG array.

temp_chip_file['start'] = temp_chip_file['start'] - 250

temp_chip_file['end'] = temp_chip_file['end'] + 250

if blacklist_file is None:

print('No blacklist file specified ...')

exclusion_windows = BedTool.from_dataframe(temp_chip_file[['chrom', 'start','end']])

else:

bound_exclusion_windows = BedTool.from_dataframe(temp_chip_file[['chrom', 'start','end']])

blacklist_exclusion_windows = BedTool(blacklist_file)

exclusion_windows = BedTool.cat(

*[blacklist_exclusion_windows, bound_exclusion_windows])

return exclusion_windows

def make_random_shift(coords, L, buffer=25):

"""

This function takes as input a set of bed coordinates dataframe

It finds the mid-point for each record or Interval in the bed file,

shifts the mid-point, and generates a windows of length L.

If training window length is L, then we must ensure that the

peak center is still within the training window.

Therefore: -L/2 < shift < L/2

To add in a buffer: -L/2 + 25 <= shift <= L/2 + 25

# Note: The 50 here is a tunable hyper-parameter.

Parameters:

coords(pandas dataFrame): This is an input bedfile (first 3 column names: "chr", "start", "end")

Returns:

shifted_coords(pandas dataFrame): The output bedfile with shifted coords

"""

low = int(-L/2 + buffer)

high = int(L/2 - buffer)

return (coords.assign(midpoint=lambda x: (x["start"]+x["end"])/2)

.astype({"midpoint": int})

.assign(midpoint=lambda x: x["midpoint"] + np.random.randint(low=low, high=high, size=len(coords)))

.apply(lambda s: pd.Series([s["chrom"], int(s["midpoint"]-L/2), int(s["midpoint"]+L/2)],

index=["chrom", "start", "end"]), axis=1))

def make_flank(coords, L, d):

"""

Make flanking regions by:

1. Shift midpoint by d

2. Expand midpoint to upstream/downstream by L/2

"""

return (coords.assign(midpoint=lambda x: (x["start"]+x["end"])/2)

.astype({"midpoint": int})

.assign(midpoint=lambda x: x["midpoint"] + d)

.apply(lambda s: pd.Series([s["chrom"], int(s["midpoint"]-L/2), int(s["midpoint"]+L/2)],

index=["chrom", "start", "end"]), axis=1))

def random_coords(gs, incl, excl, l=500, n=1000):

"""

Randomly sample n intervals of length l from the genome,

shuffle to make all intervals inside the desired regions

and outside exclusion regions

"""

return (BedTool()

.random(l=l, n=n, g=gs)

.shuffle(g=gs, incl=incl.fn, excl=excl.fn)

.to_dataframe()[["chrom", "start", "end"]])

def chop_genome(gs, chroms, excl, stride=500, l=500):

"""

Given a genome size file and chromosome list,

chop these chromosomes into intervals of length l,

with include/exclude regions specified

"""

def intervals_loop(chrom, start, stride, l, size):

intervals = []

while True:

if (start + l) < size:

intervals.append((chrom, start, start+l))

else:

break

start += stride

return pd.DataFrame(intervals, columns=["chrom", "start", "end"])

genome_sizes = (pd.read_csv(gs, sep="\t", names=["chrom", "len"])

.set_index("chrom")

.loc[chroms])

genome_chops = pd.concat([intervals_loop(i.Index, 0, stride, l, i.len)

for i in genome_sizes.itertuples()])

genome_chops_bdt = BedTool.from_dataframe(genome_chops)

return (genome_chops_bdt.intersect(excl, v=True)

.to_dataframe()[["chrom", "start", "end"]])

def clean_bed(coords):

"""

Clean the bed file:

1. Remove intervals with start < 0

"""

return coords.loc[coords["start"]>=0]

def get_data(coords, genome_fasta, chromatin_tracks, nbins, reverse=False, numProcessors=1):

"""

Given coordinates dataframe, extract the sequence and chromatin signal

"""

y = coords["label"]

# get pointer

genome_pyfasta = pyfasta.Fasta(genome_fasta)

# split coordinates and assign chunks to workers

chunks = np.array_split(coords, numProcessors)

get_coverage_worker_freeze = functools.partial(get_coverage_worker, nbins=nbins,

bigwig_files=chromatin_tracks, reverse=reverse)

pool = Pool(numProcessors)

res = pool.map_async(get_coverage_worker_freeze, chunks)

# let's take care of sequence

X_seq = get_sequence_worker(coords, genome_pyfasta, reverse=reverse)

# gather the results

chromatin_out_lists = res.get()

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

return X_seq, chromatin_out_lists, y

def get_data_TFRecord(coords, genome_fasta, chromatin_tracks, nbins, outprefix, reverse=False, numProcessors=1):

"""

Given coordinates dataframe, extract the sequence and chromatin signal,

Then save in **TFReocrd** format

"""

# split coordinates and assign chunks to workers

num_chunks = math.ceil(len(coords) / 7000)

chunks = np.array_split(coords, num_chunks)

get_data_TFRecord_worker_freeze = functools.partial(get_data_TFRecord_worker,

fasta=genome_fasta, nbins=nbins,

bigwig_files=chromatin_tracks, reverse=reverse)

pool = Pool(numProcessors)

res = pool.starmap_async(get_data_TFRecord_worker_freeze, zip(chunks, [outprefix + "_" + str(i) for i in range(num_chunks)]))

res = res.get()

return res

def get_data_TFRecord_worker(coords, outprefix, fasta, bigwig_files, nbins, reverse=False):

genome_pyfasta = pyfasta.Fasta(fasta)

bigwigs = [pyBigWig.open(bw) for bw in bigwig_files]

# Reference: https://stackoverflow.com/questions/47861084/how-to-store-numpy-arrays-as-tfrecord

def serialize_array(array):

array = tf.io.serialize_tensor(array)

return array

def _bytes_feature(value):

"""Returns a bytes_list from a string / byte."""

if isinstance(value, type(tf.constant(0))): # if value ist tensor

value = value.numpy() # get value of tensor

return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

TFRecord_file = outprefix + ".TFRecord"

with tf.io.TFRecordWriter(TFRecord_file) as writer:

for item in coords.itertuples():

feature_dict = defaultdict()

# seq

seq = genome_pyfasta[item.chrom][int(item.start):int(item.end)]

if reverse:

seq = rev_comp(seq)

seq_serialized = serialize_array(dna2onehot(seq))

feature_dict["seq"] = _bytes_feature(seq_serialized)

# chromatin track

try:

for idx, bigwig in enumerate(bigwigs):

m = (np.nan_to_num(bigwig.values(item.chrom, item.start, item.end))

.reshape((nbins, -1))

.mean(axis=1, dtype=float))

if reverse:

m = m[::-1]

m_serialized = serialize_array(m)

feature_dict[bigwig_files[idx]] = _bytes_feature(m_serialized)

except RuntimeError as e:

logging.warning(e)

logging.warning(f"Chromatin track {bigwig_files[idx]} doesn't have information in {item} Skip this region...")

continue

# label

feature_dict["label"] = tf.train.Feature(int64_list=tf.train.Int64List(value=[item.label]))

example = tf.train.Example(features=tf.train.Features(feature=feature_dict))

writer.write(example.SerializeToString())

for bw in bigwigs: bw.close()

return TFRecord_file

def dna2onehot(dnaSeq):

DNA2index = {

"A": 0,

"T": 1,

"G": 2,

"C": 3

}

seqLen = len(dnaSeq)

# initialize the matrix to seqlen x 4

seqMatrixs = np.zeros((seqLen,4), dtype=int)

# change the value to matrix

dnaSeq = dnaSeq.upper()

for j in range(0,seqLen):

try:

seqMatrixs[j, DNA2index[dnaSeq[j]]] = 1

except KeyError as e:

continue

return seqMatrixs

def rev_comp(inp_str):

rc_dict = {'A': 'T', 'G': 'C', 'T': 'A', 'C': 'G', 'c': 'g',

'g': 'c', 't': 'a', 'a': 't', 'n': 'n', 'N': 'N'}

outp_str = list()

for nucl in inp_str:

outp_str.append(rc_dict[nucl])

return ''.join(outp_str)[::-1]

def get_sequence_worker(coords, fasta, reverse=False):

"""

Get the sequence in provided regions

"""

seqs = []

for item in coords.itertuples():

seq = fasta[item.chrom][int(item.start):int(item.end)]

if reverse:

seq = rev_comp(seq)

seqs.append(seq)

return seqs

def get_coverage_worker(coords, bigwig_files, nbins, reverse=False):

"""

Get the signal coverage in provided regions, summarize mean in each bin

"""

bigwigs = [pyBigWig.open(bw) for bw in bigwig_files]

ms = [[] for x in bigwigs]

for idx, bigwig in enumerate(bigwigs):

for item in coords.itertuples():

try:

m = (np.nan_to_num(bigwig.values(item.chrom, item.start, item.end))

.reshape((nbins, -1))

.mean(axis=1))

except RuntimeError as e:

logging.warning(e)

logging.warning(f"Skip region: {item}")

continue

if reverse:

m = m[::-1]

ms[idx].append(m)

return np.array(ms)

if __name__ == "__main__":

chip_seq_coordinates = load_chipseq_data("Bichrom/sample_data/Ascl1.peaks", genome_sizes_file="Bichrom/sample_data/mm10.info",

to_keep=None, to_filter=['chr17', 'chr11', 'chrM', 'chrUn'])

# pickle.dump(chip_seq_coordinates, open("chip_coords.pickle", "wb"))

# pickle.dump(make_random_shift(chip_seq_coordinates, 500), open("shifted.pickle", "wb"))