<commit_before>#!/usr/bin/env python
# Copyright 2017 Calico LLC

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =========================================================================

from optparse import OptionParser
from collections import OrderedDict
import math
import os
import random
import subprocess
import sys
import tempfile
import time

import h5py
import joblib
import numpy as np
import pyBigWig
import pysam
import tensorflow as tf

import slurm

import basenji

'''
basenji_hdf5.py

Tile the genome and project the full functional profile to latent space
using a given model. Save the result in HDF5 for Basenji learning.

To Do:
 -If unmappable regions are cutting my data, I could squeeze a little more out
   by allowing them to sit at the edges of sequences where I'm not making
   predictions anyway.
'''

################################################################################
def main():
    usage = 'usage: %prog [options] <fasta_file> <sample_wigs_file> <hdf5_file>'
    parser = OptionParser(usage)
    parser.add_option('-b', dest='limit_bed', help='Limit to segments that overlap regions in a BED file')
    parser.add_option('-c', dest='clip', default=None, type='float', help='Clip target values to have minimum [Default: %default]')
    parser.add_option('--cluster_dir', dest='cluster_dir', default='basenji_hdf5')
    parser.add_option('-d', dest='sample_pct', default=1.0, type='float', help='Down-sample the segments')
    parser.add_option('-f', dest='fourier_dim', default=None, type='int', help='Fourier transform dimension [Default: %default]')
    parser.add_option('-g', dest='gaps_file', help='Genome assembly gaps BED [Default: %default]')
    parser.add_option('-l', dest='seq_length', default=1024, type='int', help='Sequence length [Default: %default]')
    parser.add_option('--log2', dest='log10to2', default=False, action='store_true', help='Transform values from log10 to log2 [Default: %default]')
    parser.add_option('--mult_cov', dest='cov_multiplier', default=1, type='float', help='Coverage multiplier, useful when the read extension and pool width do not match [Default: %default]')
    parser.add_option('-n', dest='na_t', default=0.25, type='float', help='Remove sequences with an NA% greater than this threshold [Default: %default]')
    parser.add_option('-o', dest='out_bed_file', help='Output the train/valid/test sequences as a BED file')
    parser.add_option('-p', dest='processes', default=1, type='int', help='Number parallel processes to load data [Default: %default]')
    parser.add_option('-s', dest='stride', type='int', help='Stride to advance segments [Default: seq_length]')
    parser.add_option('-t', dest='test_pct_or_chr', type='str', default=0.05, help='Proportion of the data for testing [Default: %default]')
    parser.add_option('-u', dest='unmap_bed', help='Unmappable segments to set to NA')
    parser.add_option('-w', dest='pool_width', type='int', default=1, help='Average pooling width [Default: %default]')
    parser.add_option('-v', dest='valid_pct_or_chr', type='str', default=0.05, help='Proportion of the data for validation [Default: %default]')
    parser.add_option('-z', dest='compression', help='h5py compression [Default: %default]')
    (options,args) = parser.parse_args()

    if len(args) != 3:
        parser.error('Must provide genome file, sample Wig/BigWig labels and paths, and model output file')
    else:
        fasta_file = args[0]
        sample_wigs_file = args[1]
        hdf5_file = args[2]

    random.seed(1)

    if options.stride is None:
        options.stride = options.seq_length

    ################################################################
    # assess bigwigs
    ################################################################
    # get wig files and labels
    target_wigs = OrderedDict()
    for line in open(sample_wigs_file, encoding='UTF-8'):
        a = line.split()
        target_wigs[a[0]] = a[1]

    if options.fourier_dim is not None and 2*options.fourier_dim >= options.seq_length/options.pool_width:
        print("Fourier transform to %d dims won't compress %d length sequences with %d pooling" % (options.fourier_dim, options.seq_length, options.pool_width), file=sys.stderr)
        exit(1)


    ################################################################
    # prepare genomic segments
    ################################################################
    chrom_segments = basenji.genome.load_chromosomes(fasta_file)

    # remove gaps
    if options.gaps_file:
        chrom_segments = basenji.genome.split_contigs(chrom_segments, options.gaps_file)

    # ditch the chromosomes
    segments = []
    for chrom in chrom_segments:
        segments += [(chrom, seg_start, seg_end) for seg_start, seg_end in chrom_segments[chrom]]

    # standardize order
    segments.sort()

    # filter for large enough
    segments = [cse for cse in segments if cse[2]-cse[1] >= options.seq_length]

    # down-sample
    if options.sample_pct < 1.0:
        segments = random.sample(segments, int(options.sample_pct*len(segments)))

    # limit to a BED file
    if options.limit_bed is not None:
        segments = limit_segments(segments, options.limit_bed)

    if not os.path.isdir(options.cluster_dir):
        os.mkdir(options.cluster_dir)

    # print segments to BED file
    seg_bed_file = '%s/segments.bed' % options.cluster_dir
    seg_bed_out = open(seg_bed_file, 'w')
    for chrom, seg_start, seg_end in segments:
        print('%s\t%d\t%d' % (chrom,seg_start,seg_end), file=seg_bed_out)
    seg_bed_out.close()


    ################################################################
    # bigwig read and process
    ################################################################
    print('Reading and pre-processing bigwigs for %d segments' % len(segments), flush=True)

    targets_real = []
    targets_imag = []

    # generate numpy arrays on cluster
    jobs = []
    for target_label in target_wigs.keys():
        wig_file = target_wigs[target_label]
        npy_file = '%s/%s' % (options.cluster_dir, target_label)
        if not os.path.isfile(npy_file) and not os.path.isfile('%s.npy'%npy_file):
            print(npy_file)

            if os.path.splitext(wig_file)[1] == 'h5':
                script = 'seqs_hdf5.py'
            else:
                script = 'bigwig_hdf5.py'

            cmd = 'echo $HOSTNAME; %s -l %d -s %d -w %d %s %s %s' % (script, options.seq_length, options.stride, options.pool_width, wig_file, seg_bed_file, npy_file)
            name = 'hdf5_%s'%target_label
            outf = '%s/%s.out' % (options.cluster_dir, target_label)
            errf = '%s/%s.err' % (options.cluster_dir, target_label)
            j = slurm.Job(cmd, name, outf, errf, queue='flash', mem=16000, time='4:0:0')
            jobs.append(j)

    slurm.multi_run(jobs)

    # load into targets_real, targets_imag
    for target_label in target_wigs.keys():
        npy_file = '%s/%s.npy' % (options.cluster_dir, target_label)
        wig_targets = np.load(npy_file)
        targets_real.append(wig_targets)

    # transpose from TxSxL to SxLxT<commit_msg>path split bug<commit_after>#!/usr/bin/env python
# Copyright 2017 Calico LLC

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =========================================================================

from optparse import OptionParser
from collections import OrderedDict
import math
import os
import random
import subprocess
import sys
import tempfile
import time

import h5py
import joblib
import numpy as np
import pyBigWig
import pysam
import tensorflow as tf

import slurm

import basenji

'''
basenji_hdf5.py

Tile the genome and project the full functional profile to latent space
using a given model. Save the result in HDF5 for Basenji learning.

To Do:
 -If unmappable regions are cutting my data, I could squeeze a little more out
   by allowing them to sit at the edges of sequences where I'm not making
   predictions anyway.
'''

################################################################################
def main():
    usage = 'usage: %prog [options] <fasta_file> <sample_wigs_file> <hdf5_file>'
    parser = OptionParser(usage)
    parser.add_option('-b', dest='limit_bed', help='Limit to segments that overlap regions in a BED file')
    parser.add_option('-c', dest='clip', default=None, type='float', help='Clip target values to have minimum [Default: %default]')
    parser.add_option('--cluster_dir', dest='cluster_dir', default='basenji_hdf5')
    parser.add_option('-d', dest='sample_pct', default=1.0, type='float', help='Down-sample the segments')
    parser.add_option('-f', dest='fourier_dim', default=None, type='int', help='Fourier transform dimension [Default: %default]')
    parser.add_option('-g', dest='gaps_file', help='Genome assembly gaps BED [Default: %default]')
    parser.add_option('-l', dest='seq_length', default=1024, type='int', help='Sequence length [Default: %default]')
    parser.add_option('--log2', dest='log10to2', default=False, action='store_true', help='Transform values from log10 to log2 [Default: %default]')
    parser.add_option('--mult_cov', dest='cov_multiplier', default=1, type='float', help='Coverage multiplier, useful when the read extension and pool width do not match [Default: %default]')
    parser.add_option('-n', dest='na_t', default=0.25, type='float', help='Remove sequences with an NA% greater than this threshold [Default: %default]')
    parser.add_option('-o', dest='out_bed_file', help='Output the train/valid/test sequences as a BED file')
    parser.add_option('-p', dest='processes', default=1, type='int', help='Number parallel processes to load data [Default: %default]')
    parser.add_option('-s', dest='stride', type='int', help='Stride to advance segments [Default: seq_length]')
    parser.add_option('-t', dest='test_pct_or_chr', type='str', default=0.05, help='Proportion of the data for testing [Default: %default]')
    parser.add_option('-u', dest='unmap_bed', help='Unmappable segments to set to NA')
    parser.add_option('-w', dest='pool_width', type='int', default=1, help='Average pooling width [Default: %default]')
    parser.add_option('-v', dest='valid_pct_or_chr', type='str', default=0.05, help='Proportion of the data for validation [Default: %default]')
    parser.add_option('-z', dest='compression', help='h5py compression [Default: %default]')
    (options,args) = parser.parse_args()

    if len(args) != 3:
        parser.error('Must provide genome file, sample Wig/BigWig labels and paths, and model output file')
    else:
        fasta_file = args[0]
        sample_wigs_file = args[1]
        hdf5_file = args[2]

    random.seed(1)

    if options.stride is None:
        options.stride = options.seq_length

    ################################################################
    # assess bigwigs
    ################################################################
    # get wig files and labels
    target_wigs = OrderedDict()
    for line in open(sample_wigs_file, encoding='UTF-8'):
        a = line.split()
        target_wigs[a[0]] = a[1]

    if options.fourier_dim is not None and 2*options.fourier_dim >= options.seq_length/options.pool_width:
        print("Fourier transform to %d dims won't compress %d length sequences with %d pooling" % (options.fourier_dim, options.seq_length, options.pool_width), file=sys.stderr)
        exit(1)


    ################################################################
    # prepare genomic segments
    ################################################################
    chrom_segments = basenji.genome.load_chromosomes(fasta_file)

    # remove gaps
    if options.gaps_file:
        chrom_segments = basenji.genome.split_contigs(chrom_segments, options.gaps_file)

    # ditch the chromosomes
    segments = []
    for chrom in chrom_segments:
        segments += [(chrom, seg_start, seg_end) for seg_start, seg_end in chrom_segments[chrom]]

    # standardize order
    segments.sort()

    # filter for large enough
    segments = [cse for cse in segments if cse[2]-cse[1] >= options.seq_length]

    # down-sample
    if options.sample_pct < 1.0:
        segments = random.sample(segments, int(options.sample_pct*len(segments)))

    # limit to a BED file
    if options.limit_bed is not None:
        segments = limit_segments(segments, options.limit_bed)

    if not os.path.isdir(options.cluster_dir):
        os.mkdir(options.cluster_dir)

    # print segments to BED file
    seg_bed_file = '%s/segments.bed' % options.cluster_dir
    seg_bed_out = open(seg_bed_file, 'w')
    for chrom, seg_start, seg_end in segments:
        print('%s\t%d\t%d' % (chrom,seg_start,seg_end), file=seg_bed_out)
    seg_bed_out.close()


    ################################################################
    # bigwig read and process
    ################################################################
    print('Reading and pre-processing bigwigs for %d segments' % len(segments), flush=True)

    targets_real = []
    targets_imag = []

    # generate numpy arrays on cluster
    jobs = []
    for target_label in target_wigs.keys():
        wig_file = target_wigs[target_label]
        npy_file = '%s/%s' % (options.cluster_dir, target_label)
        if not os.path.isfile(npy_file) and not os.path.isfile('%s.npy'%npy_file):
            print(npy_file)

            if os.path.splitext(wig_file)[1] == '.h5':
                script = 'seqs_hdf5.py'
            else:
                script = 'bigwig_hdf5.py'

            cmd = 'echo $HOSTNAME; %s -l %d -s %d -w %d %s %s %s' % (script, options.seq_length, options.stride, options.pool_width, wig_file, seg_bed_file, npy_file)
            name = 'hdf5_%s'%target_label
            outf = '%s/%s.out' % (options.cluster_dir, target_label)
            errf = '%s/%s.err' % (options.cluster_dir, target_label)
            j = slurm.Job(cmd, name, outf, errf, queue='flash', mem=16000, time='4:0:0')
            jobs.append(j)

    slurm.multi_run(jobs)

    # load into targets_real, targets_imag
    for target_label in target_wigs.keys():
        npy_file = '%s/%s.npy' % (options.cluster_dir, target_label)
        wig_targets = np.load(npy_file)
        targets_real.append(wig_targets)

    # transpose from TxSxL to SxLxT
{'fraction_non_alphanumeric': 0.1821330545026794, 'fraction_numerical': 0.008757025225460725, 'mean_word_length': 4.077305905773059, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 13, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6116629', 'n_tokens_mistral': 4785, 'n_tokens_neox': 4558, 'n_words': 1299}