description = """

Write seqchrom dataset to disk

"""

import os

import numpy as np

import functools

import math

import logging

from collections import defaultdict

from multiprocessing import Pool

import pyfaidx

import pysam

import webdataset as wds

from seqchromloader import utils, config

from seqchromloader.loader import _SeqChromDatasetByWds

logger = logging.getLogger(__name__)

def convert_data_webdataset(wds_in, wds_out, transforms=None, compress=False):

"""

Transform the provided webdataset

:param wds_in: input webdataset file

:type wds_in: string

:param wds_out: output webdataset file

:type wds_out: string

:param transforms: A dictionary of functions to transform the output data, accepted keys are *["seq", "chrom", "target", "label"]*

:type transforms: dict of functions

:param compress: whether to compress the output file

:type compress: boolean

"""

ds = _SeqChromDatasetByWds(wds_in, transforms=transforms, keep_key=True)

sink = wds.TarWriter(wds_out, compress=compress)

for (key, seq, chrom, target, label) in ds:

feature_dict = defaultdict()

feature_dict["__key__"] = key

feature_dict["seq.npy"] = seq

feature_dict["chrom.npy"] = chrom

feature_dict["target.npy"] = target

feature_dict["label.npy"] = label

sink.write(feature_dict)

sink.close()

def dump_data_webdataset(coords, genome_fasta, bigwig_filelist,

target_bam=None,

target_bw=None,

patch_left=0, patch_right=0,

outdir="dataset/", outprefix="seqchrom",

compress=True,

numProcessors=1,

transforms=None,

braceexpand=False,

samples_per_tar=10000,

batch_size=None):

"""

Given coordinates dataframe, extract the sequence and chromatin signal, save in webdataset format

:param coords: pandas DataFrame containing genomic coordinates with columns **[chrom, start, end, label]**

:type coords: pandas DataFrame

:param genome_fasta: Genome fasta file.

:type genome_fasta: str

:param bigwig_filelist: A list of bigwig files containing track information (e.g., histone modifications)

:type bigwig_filelist: list of str or None

:param target_bam: bam file to get # reads in each region, mutually exclusive with `target_bw`

:type target_bam: str or None

:param target_bw: bigwig file to get # reads in each region, mutually exclusive with `target_bam`

:type target_bw: str or None

:param patch_left: extend the seq and chrom inputs on the left by `patch_left`bp

:type patch_left: int

:param patch_right: extend the seq and chrom inputs on the right by `patch_right`bp

:type patch_right: int

:param transforms: A dictionary of functions to transform the output data, accepted keys are *["seq", "chrom", "target", "label"]*

:type transforms: dict of functions

:param outdir: output directory to save files in

:type outdir: str

:param outprefix: prefix of output files

:type outprefix: str

:param compress: whether to compress the output files

:type compress: boolean

:param numProcessors: number of processors

:type numProcessors: int

:param braceexpand: if use brace to simplify the wds file list into a string

:param braceexpand: boolean

:param DALI: Set to True if you want to use the dataset for NVIDIA DALI, it would save all arrays in bytes, which results in losing the array shape info

:param DALI: boolean

:param samples_per_tar: Number of samples included per tar file

:param samples_per_tar: int

"""

# check parameters

if (target_bam is not None and target_bw is not None):

raise Exception("Only one of target_bam and target_bw should be provided!")

# split coordinates and assign chunks to workers

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

chunks = [coords.iloc[idx] for idx in np.array_split(np.arange(len(coords)), num_chunks)]

# freeze the common parameters

## create a scaler to get statistics for normalizing chromatin marks input

## also create a multiprocessing lock

dump_data_worker_freeze = functools.partial(dump_data_webdataset_worker,

fasta=genome_fasta,

bigwig_files=bigwig_filelist,

target_bam=target_bam,

target_bw=target_bw,

patch_left=patch_left,

patch_right=patch_right,

compress=compress,

outdir=outdir,

transforms=transforms,

batch_size=batch_size)

count_of_digits = 0

nc = num_chunks

while nc > 0:

nc = int(nc/10)

count_of_digits += 1

pool = Pool(numProcessors)

res = pool.starmap_async(dump_data_worker_freeze, zip(chunks, [outprefix + "_" + format(i, f'0{count_of_digits}d') for i in range(num_chunks)]))

files = res.get()

if braceexpand:

begin = format(0, f'0{count_of_digits}d')

end = format(range(num_chunks)[-1], f'0{count_of_digits}d')

return os.path.join(outdir, f"{outprefix}_{{{begin}..{end}}}.tar.gz" if compress else f"{outprefix}_{{{begin}..{end}}}.tar")

else:

return files

def dump_data_webdataset_worker(coords,

outprefix,

fasta,

bigwig_files=None,

target_bam=None,

target_bw=None,

patch_left=0, patch_right=0,

outdir="dataset/",

compress=True,

transforms=None,

batch_size=None,

):

#get handlers

genome_pyfaidx = pyfaidx.Fasta(fasta)

bigwigs = [utils.BigWig(bw, backend=config.BIGWIG_BACKEND) for bw in bigwig_files] if bigwig_files is not None else None

if target_bam is not None:

if isinstance(target_bam, list):

target = [pysam.AlignmentFile(b) for b in target_bam]

else:

target = pysam.AlignmentFile(target_bam)

elif target_bw is not None:

if isinstance(target_bw, list):

target = [utils.BigWig(b, backend=config.BIGWIG_BACKEND) for b in target_bw]

else:

target = utils.BigWig(target_bw, backend=config.BIGWIG_BACKEND)

else:

target = None

# iterate all records

filename = os.path.join(outdir, f"{outprefix}.tar.gz" if compress else f"{outprefix}.tar")

sink = wds.TarWriter(filename, compress=compress)

counter = 0; keys = []; seq_arr = []; chrom_arr = []; target_arr = []; label_arr = []

for rindex, item in enumerate(coords.itertuples()):

try:

feature = utils.extract_info(

item.chrom,

item.start,

item.end,

item.label,

genome_pyfaidx=genome_pyfaidx,

bigwigs=bigwigs,

target=target,

strand=item.strand,

transforms=transforms,

patch_left=patch_left,

patch_right=patch_right

)

except utils.BigWigInaccessible as e:

logger.warning(f"Skip the region {item.chrom}:{item.start}-{item.end} due to BigWigInaccessible Exception")

continue

except pyfaidx.FetchError as e:

logger.warning(f"Skip the region {item.chrom}:{item.start}-{item.end} due to pyfaidx FetchError")

continue

except AssertionError as e:

logger.warning(f"Skip the region {item.chrom}:{item.start}-{item.end} due to AssertionError: {e}")

continue

feature_dict = defaultdict()

if batch_size is None:

feature_dict["__key__"] = f"{rindex}_{item.chrom}:{item.start-patch_left}-{item.end+patch_right}_{item.strand}"

feature_dict["seq.npy"] = feature['seq']

feature_dict["chrom.npy"] = feature['chrom']

feature_dict["target.npy"] = feature['target']

feature_dict["label.npy"] = feature['label']

sink.write(feature_dict)

feature_dict = defaultdict()

else:

counter += 1

keys.append(f"{rindex}_{item.chrom}:{item.start}-{item.end}_{item.strand}")

seq_arr.append(feature['seq']); chrom_arr.append(feature['chrom']); target_arr.append(feature['target']); label_arr.append(feature['label'])

if counter>=batch_size:

feature_dict["__key__"] = ','.join(keys)

feature_dict["seq.npy"] = np.array(seq_arr)

feature_dict["chrom.npy"] = np.array(chrom_arr)

feature_dict["target.npy"] = np.array(target_arr)

feature_dict["label.npy"] = np.array(label_arr)

sink.write(feature_dict)

keys, seq_arr, chrom_arr, target_arr, label_arr = [], [], [], [] ,[]

counter = 0

feature_dict = defaultdict()

if batch_size is not None:

feature_dict["__key__"] = ','.join(keys)

feature_dict["seq.npy"] = np.array(seq_arr)

feature_dict["chrom.npy"] = np.array(chrom_arr)

feature_dict["target.npy"] = np.array(target_arr)

feature_dict["label.npy"] = np.array(label_arr)

sink.write(feature_dict)

sink.close()

if bigwigs is not None:

for bw in bigwigs: bw.close()

return filename