code: #!/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
  import os
  import sys
  
  import h5py
  import numpy as np
  import pdb
  import pysam
  
  from basenji_data import ModelSeq
  from basenji.dna_io import dna_1hot
  
  import tensorflow as tf
  
  """
  basenji_data_write.py
  
  Write TF Records for batches of model sequences.
  """
  
  ################################################################################
  # main
  ################################################################################
  def main():
    usage = 'usage: %prog [options] <fasta_file> <seqs_bed_file> <seqs_cov_dir> <tfr_file>'
    parser = OptionParser(usage)
    parser.add_option('-g', dest='genome_index',
        default=None, type='int', help='Genome index')
    parser.add_option('-s', dest='start_i',
        default=0, type='int',
        help='Sequence start index [Default: %default]')
    parser.add_option('-e', dest='end_i',
        default=None, type='int',
        help='Sequence end index [Default: %default]')
    parser.add_option('--te', dest='target_extend',
        default=None, type='int', help='Extend targets vector [Default: %default]')
    parser.add_option('--ts', dest='target_start',
        default=0, type='int', help='Write targets into vector starting at index [Default: %default')
    parser.add_option('-u', dest='umap_npy',
        help='Unmappable array numpy file')
    parser.add_option('--umap_set', dest='umap_set',
        default=None, type='float',
        help='Sequence distribution value to set unmappable positions to, eg 0.25.')
    (options, args) = parser.parse_args()
  
    if len(args) != 4:
      parser.error('Must provide input arguments.')
    else:
      fasta_file = args[0]
      seqs_bed_file = args[1]
      seqs_cov_dir = args[2]
      tfr_file = args[3]
  
    ################################################################
    # read model sequences
  
    model_seqs = []
    for line in open(seqs_bed_file):
      a = line.split()
      model_seqs.append(ModelSeq(a[0],int(a[1]),int(a[2]),None))
  
    if options.end_i is None:
      options.end_i = len(model_seqs)
  
    num_seqs = options.end_i - options.start_i
  
    ################################################################
    # determine sequence coverage files
  
    seqs_cov_files = []
    ti = 0
    if options.genome_index is None:
      seqs_cov_file = '%s/%d.h5' % (seqs_cov_dir, ti)
    else:
      seqs_cov_file = '%s/%d-%d.h5' % (seqs_cov_dir, options.genome_index, ti)
    while os.path.isfile(seqs_cov_file):
      seqs_cov_files.append(seqs_cov_file)
      ti += 1
      if options.genome_index is None:
        seqs_cov_file = '%s/%d.h5' % (seqs_cov_dir, ti)
      else:
        seqs_cov_file = '%s/%d-%d.h5' % (seqs_cov_dir, options.genome_index, ti)
  
    if len(seqs_cov_files) == 0:
      print('Sequence coverage files not found, e.g. %s' % seqs_cov_file, file=sys.stderr)
      exit(1)
  
    seq_pool_len = h5py.File(seqs_cov_files[0], 'r')['seqs_hic'].shape[1]
    num_targets = len(seqs_cov_files)
  
    ################################################################
    # read targets
  
    # extend targets
    num_targets_tfr = num_targets
    if options.target_extend is not None:
      assert(options.target_extend >= num_targets_tfr)
      num_targets_tfr = options.target_extend
  
    # initialize targets
    targets = np.zeros((num_seqs, seq_pool_len, seq_pool_len, num_targets_tfr), dtype='float16')
  
    # read each target
    for ti in range(num_targets):
      seqs_cov_open = h5py.File(seqs_cov_files[ti], 'r')
      tii = options.target_start + ti
      targets[:,:,:,tii] = seqs_cov_open['seqs_hic'][options.start_i:options.end_i,:,:]
      seqs_cov_open.close()
  
    ################################################################
    # modify unmappable
  
    if options.umap_npy is not None and options.umap_set is not None:
      unmap_mask = np.load(options.umap_npy)
  
      for si in range(num_seqs):
        msi = options.start_i + si
  
        # determine unmappable null value
        seq_target_null = np.percentile(targets[si], q=[100*options.umap_set], axis=0)[0]
  
        # set unmappable positions to null
        targets[si,unmap_mask[msi,:],unmap_mask[msi,:],:] = np.minimum(targets[si,unmap_mask[msi,:],unmap_mask[msi,:],:], seq_target_null)
  
    ################################################################
    # write TFRecords
  
    # open FASTA
    fasta_open = pysam.Fastafile(fasta_file)
  
    # define options
    tf_opts = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB)
  
    with tf.python_io.TFRecordWriter(tfr_file, tf_opts) as writer:
      for si in range(num_seqs):
        msi = options.start_i + si
        mseq = model_seqs[msi]
  
        # read FASTA
        seq_dna = fasta_open.fetch(mseq.chr, mseq.start, mseq.end)
  
        # one hot code
        seq_1hot = dna_1hot(seq_dna)
  
        if options.genome_index is None:
          example = tf.train.Example(features=tf.train.Features(feature={
              'genome': _int_feature(0),
              'sequence': _bytes_feature(seq_1hot.flatten().tostring()),
              'target': _bytes_feature(targets[si,:,:,:].flatten().tostring())}))
        else:
          example = tf.train.Example(features=tf.train.Features(feature={
              'genome': _int_feature(options.genome_index),
              'sequence': _bytes_feature(seq_1hot.flatten().tostring()),
              'target': _bytes_feature(targets[si,:,:,:].flatten().tostring())}))
  
        writer.write(example.SerializeToString())
  
      fasta_open.close()
  
  
  def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
  
  def _int_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
  
  ################################################################################
  # __main__
  ################################################################################
  if __name__ == '__main__':
    main()
   #!/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
  import glob
  import os
  import pickle
  import shutil
  import subprocess
  import sys
  
  import h5py
  import numpy as np
  
  import slurm
  
  """
  sonnet_sad_multi.py
  
  Compute SNP expression difference scores for variants in a VCF file,
  using multiple processes.
  """
  
  ################################################################################
  # main
  ################################################################################
  def main():
    usage = 'usage: %prog [options] <model> <vcf_file>'
    parser = OptionParser(usage)
  
    # sad
    parser.add_option('-b', dest='batch_size',
        default=4, type='int',
        help='Batch size [Default: %default]')
    parser.add_option('-c', dest='slice_center',
        default=None, type='int',
        help='Slice center positions [Default: %default]')
    parser.add_option('-f', dest='genome_fasta',
        default='%s/data/hg19.fa' % os.environ['BASENJIDIR'],
        help='Genome FASTA for sequences [Default: %default]')
    parser.add_option('-o',dest='out_dir',
        default='sad',
        help='Output directory for tables and plots [Default: %default]')
    parser.add_option('--pseudo', dest='log_pseudo',
        default=1, type='float',
        help='Log2 pseudocount [Default: %default]')
    parser.add_option('--rc', dest='rc',
        default=False, action='store_true',
        help='Average forward and reverse complement predictions [Default: %default]')
    parser.add_option('--shifts', dest='shifts',
        default='0', type='str',
        help='Ensemble prediction shifts [Default: %default]')
    parser.add_option('--species', dest='species',
        default='human')
    parser.add_option('--stats', dest='sad_stats',
        default='SAD',
        help='Comma-separated list of stats to save. [Default: %default]')
    parser.add_option('-t', dest='targets_file',
        default=None, type='str',
        help='File specifying target indexes and labels in table format')
  
    # multi
    parser.add_option('-e', dest='conda_env',
        default='tf2.6',
        help='Anaconda environment [Default: %default]')
    parser.add_option('--name', dest='name',
        default='sad', help='SLURM name prefix [Default: %default]')
    parser.add_option('--max_proc', dest='max_proc',
        default=None, type='int',
        help='Maximum concurrent processes [Default: %default]')
    parser.add_option('-p', dest='processes',
        default=None, type='int',
        help='Number of processes, passed by multi script')
    parser.add_option('-q', dest='queue',
        default='gtx1080ti',
        help='SLURM queue on which to run the jobs [Default: %default]')
    parser.add_option('-r', dest='restart',
        default=False, action='store_true',
        help='Restart a partially completed job [Default: %default]')
    (options, args) = parser.parse_args()
  
    if len(args) != 2:
      parser.error('Must provide model and VCF file')
    else:
      model_file = args[0]
      vcf_file = args[1]
  
    #######################################################
    # prep work
  
    # output directory
    if not options.restart:
      if os.path.isdir(options.out_dir):
        print('Please remove %s' % options.out_dir, file=sys.stderr)
        exit(1)
      os.mkdir(options.out_dir)
  
    # pickle options
    options_pkl_file = '%s/options.pkl' % options.out_dir
    options_pkl = open(options_pkl_file, 'wb')
    pickle.dump(options, options_pkl)
    options_pkl.close()
  
    #######################################################
    # launch worker threads
    jobs = []
    for pi in range(options.processes):
      if not options.restart or not job_completed(options, pi):
        cmd = '. /home/drk/anaconda3/etc/profile.d/conda.sh;'
        cmd += ' conda activate %s;' % options.conda_env
  
        cmd += ' sonnet_sad.py %s %s %d' % (
            options_pkl_file, ' '.join(args), pi)
  
        name = '%s_p%d' % (options.name, pi)
        outf = '%s/job%d.out' % (options.out_dir, pi)
        errf = '%s/job%d.err' % (options.out_dir, pi)
  
        j = slurm.Job(cmd, name,
                      outf, errf,
                      queue=options.queue, gpu=1,
                      mem=22000, time='14-0:0:0')
        jobs.append(j)
  
    slurm.multi_run(jobs, max_proc=options.max_proc, verbose=True,
                    launch_sleep=10, update_sleep=60)
  
    #######################################################
    # collect output
  
    collect_h5('sad.h5', options.out_dir, options.processes)
  
    # for pi in range(options.processes):
    #     shutil.rmtree('%s/job%d' % (options.out_dir,pi))
  
  
  def collect_h5(file_name, out_dir, num_procs):
    # count variants
    num_variants = 0
    for pi in range(num_procs):
      # open job
      job_h5_file = '%s/job%d/%s' % (out_dir, pi, file_name)
      job_h5_open = h5py.File(job_h5_file, 'r')
      num_variants += len(job_h5_open['snp'])
      job_h5_open.close()
  
    # initialize final h5
    final_h5_file = '%s/%s' % (out_dir, file_name)
    final_h5_open = h5py.File(final_h5_file, 'w')
  
    # keep dict for string values
    final_strings = {}
  
    job0_h5_file = '%s/job0/%s' % (out_dir, file_name)
    job0_h5_open = h5py.File(job0_h5_file, 'r')
    for key in job0_h5_open.keys():
      if key in ['percentiles', 'target_ids', 'target_labels']:
        # copy
        final_h5_open.create_dataset(key, data=job0_h5_open[key])
  
      elif key[-4:] == '_pct':
        values = np.zeros(job0_h5_open[key].shape)
        final_h5_open.create_dataset(key, data=values)
  
      elif job0_h5_open[key].dtype.char == 'S':
          final_strings[key] = []
  
      elif job0_h5_open[key].ndim == 1:
        final_h5_open.create_dataset(key, shape=(num_variants,), dtype=job0_h5_open[key].dtype)
  
      else:
        num_targets = job0_h5_open[key].shape[1]
        final_h5_open.create_dataset(key, shape=(num_variants, num_targets), dtype=job0_h5_open[key].dtype)
  
    job0_h5_open.close()
  
    # set values
    vi = 0
    for pi in range(num_procs):
      # open job
      job_h5_file = '%s/job%d/%s' % (out_dir, pi, file_name)
      job_h5_open = h5py.File(job_h5_file, 'r')
  
      # append to final
      for key in job_h5_open.keys():
        if key in ['percentiles', 'target_ids', 'target_labels']:
          # once is enough
          pass
  
        elif key[-4:] == '_pct':
          # average
          u_k1 = np.array(final_h5_open[key])
          x_k = np.array(job_h5_open[key])
          final_h5_open[key][:] = u_k1 + (x_k - u_k1) / (pi+1)
  
        else:
          if job_h5_open[key].dtype.char == 'S':
            final_strings[key] += list(job_h5_open[key])
          else:
            job_variants = job_h5_open[key].shape[0]
            final_h5_open[key][vi:vi+job_variants] = job_h5_open[key]
  
      vi += job_variants
      job_h5_open.close()
  
    # create final string datasets
    for key in final_strings:
      final_h5_open.create_dataset(key,
        data=np.array(final_strings[key], dtype='S'))
  
    final_h5_open.close()
  
  
  def job_completed(options, pi):
    """Check whether a specific job has generated its
       output file."""
    out_file = '%s/job%d/sad.h5' % (options.out_dir, pi)
    return os.path.isfile(out_file) or os.path.isdir(out_file)
  
  
  ################################################################################
  # __main__
  ################################################################################
  if __name__ == '__main__':
    main()
  
apis: numpy.load numpy.zeros tensorflow.python_io.TFRecordOptions tensorflow.train.Int64List tensorflow.python_io.TFRecordWriter tensorflow.train.BytesList numpy.percentile numpy.minimum
apis: numpy.array numpy.zeros