<commit_before>﻿#!/usr/bin/env python2.7
from __future__ import print_function

import argparse
import os
import multiprocessing
import subprocess
import sys
import textwrap
import tarfile
from urlparse import urlparse
import math
import shutil
import glob
import time
import datetime
import logging
from contextlib import closing

import yaml
from bd2k.util.files import mkdir_p
from bd2k.util.processes import which
from toil.job import Job
from toil.lib.docker import dockerCall, dockerCheckOutput, _fixPermissions
from toil_lib import require, UserError
from toil_lib.files import tarball_files, copy_files
from toil_lib.jobs import map_job
from toil_lib.urls import download_url, s3am_upload

# formats
SCHEMES = ('http', 'file', 's3', 'ftp')

# filenames
DEFAULT_CONFIG_NAME = 'config-toil-marginphase.yaml'
DEFAULT_MANIFEST_NAME = 'manifest-toil-marginphase.tsv'

# docker images
DOCKER_SAMTOOLS = "quay.io/ucsc_cgl/samtools:1.3--256539928ea162949d8a65ca5c79a72ef557ce7c"
DOCKER_MARGIN_PHASE = "quay.io/ucsc_cgl/margin_phase:latest"

# resource
MP_CPU = 2
MP_MEM_BAM_FACTOR = 2048
MP_MEM_REF_FACTOR = 2
MP_DSK_BAM_FACTOR = 2.5
MP_DSK_REF_FACTOR = 1.5


def parse_samples(path_to_manifest):
    """
    Parses samples, specified in either a manifest or listed with --samples

    :param str path_to_manifest: Path to configuration file
    :return: Samples and their attributes as defined in the manifest
    :rtype: list[list]
    """

    samples = []
    with open(path_to_manifest, 'r') as f:
        for line in f.readlines():
            if line.isspace() or line.startswith('#'):
                continue
            sample = line.strip().split('\t')

            # validate structure
            if len(sample) != 2:
                raise UserError('Bad manifest format! Expected 2 tab-separated columns, got: {}'.format(sample))

            # extract sample parts
            uuid, url = sample

            # validation?

            sample = [uuid, url]
            samples.append(sample)
    return samples


def prepare_input(job, sample, config):

    # job prep
    config = argparse.Namespace(**vars(config))
    uuid, url = sample
    config.uuid = uuid
    work_dir = job.fileStore.getLocalTempDir()
    start = time.time()
    job.fileStore.logToMaster("START:{}:{}".format(config.uuid, datetime.datetime.now()))

    # todo global resource estimation
    #config.cores = min(config.maxCores, multiprocessing.cpu_count())
    #config.memory
    #config.disk

    # download references
    #ref genome
    download_url(job, url=config.reference_contig, work_dir=work_dir)
    ref_genome_filename = os.path.basename(config.reference_contig)
    ref_genome_fileid = job.fileStore.writeGlobalFile(os.path.join(work_dir, ref_genome_filename))
    config.reference_genome_fileid = ref_genome_fileid
    ref_genome_size = os.stat(os.path.join(work_dir, ref_genome_filename)).st_size
    #ref vcf
    download_url(job, url=config.reference_vcf, work_dir=work_dir)
    ref_vcf_filename = os.path.basename(config.reference_vcf)
    ref_vcf_fileid = job.fileStore.writeGlobalFile(os.path.join(work_dir, ref_vcf_filename))
    config.reference_vcf_fileid = ref_vcf_fileid
    #params
    download_url(job, url=config.params, work_dir=work_dir)
    params_filename = os.path.basename(config.params)
    params_fileid = job.fileStore.writeGlobalFile(os.path.join(work_dir, params_filename))
    config.params_fileid = params_fileid


    # download bam
    download_url(job, url=url, work_dir=work_dir)
    bam_filename = os.path.basename(url)
    data_bam_location = os.path.join("/data", bam_filename)

    # index the bam
    docker_params = ["index", data_bam_location]
    job.fileStore.logToMaster("Running {} with parameters: {}".format(DOCKER_SAMTOOLS, docker_params))
    dockerCall(job, tool=DOCKER_SAMTOOLS, workDir=work_dir, parameters=docker_params)
    _fixPermissions(DOCKER_SAMTOOLS, work_dir)  # todo tmpfix

    # sanity check
    workdir_bai_location = os.path.join(work_dir, bam_filename + ".bai")
    if not os.path.isfile(workdir_bai_location):
        raise UserError("BAM index file not created for {}: {}".format(bam_filename, workdir_bai_location))

    # get start and end location
    get_idx_cmd = [
        ["samtools", "view", data_bam_location],
        ["head", "-n", "1"],
        [os.path.join("/data", write_select_column_script(work_dir))]
    ]
    start_idx_str = dockerCheckOutputExcept141(job, tool=DOCKER_SAMTOOLS, work_dir=work_dir, parameters=get_idx_cmd).strip()
    get_idx_cmd[1][0] = "tail"
    end_idx_str = dockerCheckOutputExcept141(job, tool=DOCKER_SAMTOOLS, work_dir=work_dir, parameters=get_idx_cmd).strip()
    job.fileStore.logToMaster("{}: start_pos:{}, end_pos:{}".format(config.uuid, start_idx_str, end_idx_str))
    start_idx = int(start_idx_str) - 1
    end_idx = int(end_idx_str) + 1

    # get reads from positions
    positions = list()
    idx = start_idx
    while idx < end_idx:
        chunk_start = idx - (config.partition_margin if idx != start_idx else 0)
        idx += config.partition_size
        chunk_end = idx + config.partition_margin
        positions.append([chunk_start, chunk_end])

    # enqueue jobs
    idx = 0
    return_values = list()
    for position in positions:
        #prep
        chunk_position_description = "{}:{}-{}".format(config.contig_name, position[0], position[1])
        bam_split_command = ["view", "-b", data_bam_location, chunk_position_description]
        chunk_name = "{}.{}.bam".format(config.uuid, idx)
        #write chunk
        chunk_location = os.path.join(work_dir, chunk_name)
        with open(chunk_location, 'w') as out:
            job.fileStore.logToMaster("Running {} with parameters: {}".format(DOCKER_SAMTOOLS, bam_split_command))
            dockerCall(job, tool=DOCKER_SAMTOOLS, workDir=work_dir, parameters=bam_split_command, outfile=out)
        #document read count
        read_count= get_bam_read_count(job, work_dir, chunk_name)
        job.fileStore.logToMaster("{}: chunk from {} for idx {} has {} reads"
                                  .format(config.uuid, chunk_position_description, idx, read_count))
        # enqueue marginPhase job
        if read_count > 0:
            chunk_size = os.stat(chunk_location).st_size
            chunk_fileid = job.fileStore.writeGlobalFile(chunk_location)
            mp_cores = int(min(MP_CPU, config.maxCores))
            mp_mem = int(min(chunk_size * MP_MEM_BAM_FACTOR + ref_genome_size * MP_MEM_REF_FACTOR, config.maxMemory))
            mp_disk = int(min(chunk_size * MP_DSK_BAM_FACTOR + ref_genome_size * MP_DSK_REF_FACTOR, config.maxDisk))
            margin_phase_job = job.addChildJobFn(run_margin_phase, config, chunk_fileid, idx,
                                                 memory=mp_mem, cores=mp_cores, disk=mp_disk)
            return_values.append(margin_phase_job.rv())
        idx += 1

    # enqueue consolidation job
    job.addFollowOnJobFn(consolidate_output, config, return_values)

    # log
    _log_time(job, "prepare_input", start, config.uuid)


def write_select_column_script(work_dir, column=4):
    # I feel bad for doing this, but I can't send single quotes into a toil command without them becoming escaped
    # so I'm just creating a script which does that
    filename = "select_column_{}.sh".format(column)
    file_location = os.path.join(work_dir, filename)
    with open(file_location, 'w') as out:
        print("#!/usr/bin/awk -f", file=out)
        print("{print $%d}" % column, file=out)
    os.chmod(file_location, 1023) #rwxrwxrwx
    return filename
<commit_msg>fiddling with some numbers<commit_after>﻿#!/usr/bin/env python2.7
from __future__ import print_function

import argparse
import os
import multiprocessing
import subprocess
import sys
import textwrap
import tarfile
from urlparse import urlparse
import math
import shutil
import glob
import time
import datetime
import logging
from contextlib import closing

import yaml
from bd2k.util.files import mkdir_p
from bd2k.util.processes import which
from toil.job import Job
from toil.lib.docker import dockerCall, dockerCheckOutput, _fixPermissions
from toil_lib import require, UserError
from toil_lib.files import tarball_files, copy_files
from toil_lib.jobs import map_job
from toil_lib.urls import download_url, s3am_upload

# formats
SCHEMES = ('http', 'file', 's3', 'ftp')

# filenames
DEFAULT_CONFIG_NAME = 'config-toil-marginphase.yaml'
DEFAULT_MANIFEST_NAME = 'manifest-toil-marginphase.tsv'

# docker images
DOCKER_SAMTOOLS = "quay.io/ucsc_cgl/samtools:1.3--256539928ea162949d8a65ca5c79a72ef557ce7c"
DOCKER_MARGIN_PHASE = "quay.io/ucsc_cgl/margin_phase:latest"

# resource
MP_CPU = 2
MP_MEM_BAM_FACTOR = 1024 #todo account for learning iterations
MP_MEM_REF_FACTOR = 2
MP_DSK_BAM_FACTOR = 2.5
MP_DSK_REF_FACTOR = 1.1


def parse_samples(path_to_manifest):
    """
    Parses samples, specified in either a manifest or listed with --samples

    :param str path_to_manifest: Path to configuration file
    :return: Samples and their attributes as defined in the manifest
    :rtype: list[list]
    """

    samples = []
    with open(path_to_manifest, 'r') as f:
        for line in f.readlines():
            if line.isspace() or line.startswith('#'):
                continue
            sample = line.strip().split('\t')

            # validate structure
            if len(sample) != 2:
                raise UserError('Bad manifest format! Expected 2 tab-separated columns, got: {}'.format(sample))

            # extract sample parts
            uuid, url = sample

            # validation?

            sample = [uuid, url]
            samples.append(sample)
    return samples


def prepare_input(job, sample, config):

    # job prep
    config = argparse.Namespace(**vars(config))
    uuid, url = sample
    config.uuid = uuid
    work_dir = job.fileStore.getLocalTempDir()
    start = time.time()
    job.fileStore.logToMaster("START:{}:{}".format(config.uuid, datetime.datetime.now()))

    # todo global resource estimation
    #config.cores = min(config.maxCores, multiprocessing.cpu_count())
    #config.memory
    #config.disk

    # download references
    #ref genome
    download_url(job, url=config.reference_contig, work_dir=work_dir)
    ref_genome_filename = os.path.basename(config.reference_contig)
    ref_genome_fileid = job.fileStore.writeGlobalFile(os.path.join(work_dir, ref_genome_filename))
    config.reference_genome_fileid = ref_genome_fileid
    ref_genome_size = os.stat(os.path.join(work_dir, ref_genome_filename)).st_size
    #ref vcf
    download_url(job, url=config.reference_vcf, work_dir=work_dir)
    ref_vcf_filename = os.path.basename(config.reference_vcf)
    ref_vcf_fileid = job.fileStore.writeGlobalFile(os.path.join(work_dir, ref_vcf_filename))
    config.reference_vcf_fileid = ref_vcf_fileid
    #params
    download_url(job, url=config.params, work_dir=work_dir)
    params_filename = os.path.basename(config.params)
    params_fileid = job.fileStore.writeGlobalFile(os.path.join(work_dir, params_filename))
    config.params_fileid = params_fileid


    # download bam
    download_url(job, url=url, work_dir=work_dir)
    bam_filename = os.path.basename(url)
    data_bam_location = os.path.join("/data", bam_filename)

    # index the bam
    docker_params = ["index", data_bam_location]
    job.fileStore.logToMaster("Running {} with parameters: {}".format(DOCKER_SAMTOOLS, docker_params))
    dockerCall(job, tool=DOCKER_SAMTOOLS, workDir=work_dir, parameters=docker_params)
    _fixPermissions(DOCKER_SAMTOOLS, work_dir)  # todo tmpfix

    # sanity check
    workdir_bai_location = os.path.join(work_dir, bam_filename + ".bai")
    if not os.path.isfile(workdir_bai_location):
        raise UserError("BAM index file not created for {}: {}".format(bam_filename, workdir_bai_location))

    # get start and end location
    get_idx_cmd = [
        ["samtools", "view", data_bam_location],
        ["head", "-n", "1"],
        [os.path.join("/data", write_select_column_script(work_dir))]
    ]
    start_idx_str = dockerCheckOutputExcept141(job, tool=DOCKER_SAMTOOLS, work_dir=work_dir, parameters=get_idx_cmd).strip()
    get_idx_cmd[1][0] = "tail"
    end_idx_str = dockerCheckOutputExcept141(job, tool=DOCKER_SAMTOOLS, work_dir=work_dir, parameters=get_idx_cmd).strip()
    job.fileStore.logToMaster("{}: start_pos:{}, end_pos:{}".format(config.uuid, start_idx_str, end_idx_str))
    start_idx = int(start_idx_str) - 1
    end_idx = int(end_idx_str) + 1

    # get reads from positions
    positions = list()
    idx = start_idx
    while idx < end_idx:
        chunk_start = idx - (config.partition_margin if idx != start_idx else 0)
        idx += config.partition_size
        chunk_end = idx + config.partition_margin
        positions.append([chunk_start, chunk_end])

    # enqueue jobs
    idx = 0
    return_values = list()
    for position in positions:
        #prep
        chunk_position_description = "{}:{}-{}".format(config.contig_name, position[0], position[1])
        bam_split_command = ["view", "-b", data_bam_location, chunk_position_description]
        chunk_name = "{}.{}.bam".format(config.uuid, idx)
        #write chunk
        chunk_location = os.path.join(work_dir, chunk_name)
        with open(chunk_location, 'w') as out:
            job.fileStore.logToMaster("Running {} with parameters: {}".format(DOCKER_SAMTOOLS, bam_split_command))
            dockerCall(job, tool=DOCKER_SAMTOOLS, workDir=work_dir, parameters=bam_split_command, outfile=out)
        #document read count
        chunk_size = os.stat(chunk_location).st_size
        read_count= get_bam_read_count(job, work_dir, chunk_name)
        job.fileStore.logToMaster("{}: chunk from {} for idx {} is {}b ({}mb) and has {} reads"
                                  .format(config.uuid, chunk_position_description, idx, chunk_size,
                                          int(chunk_size / 1024 / 1024), read_count))
        # enqueue marginPhase job
        if read_count > 0:
            chunk_fileid = job.fileStore.writeGlobalFile(chunk_location)
            mp_cores = int(min(MP_CPU, config.maxCores))
            mp_mem = int(min(chunk_size * MP_MEM_BAM_FACTOR + ref_genome_size * MP_MEM_REF_FACTOR, config.maxMemory))
            mp_disk = int(min(chunk_size * MP_DSK_BAM_FACTOR + ref_genome_size * MP_DSK_REF_FACTOR, config.maxDisk))
            margin_phase_job = job.addChildJobFn(run_margin_phase, config, chunk_fileid, idx,
                                                 memory=mp_mem, cores=mp_cores, disk=mp_disk)
            return_values.append(margin_phase_job.rv())
        idx += 1

    # enqueue consolidation job
    job.addFollowOnJobFn(consolidate_output, config, return_values)

    # log
    _log_time(job, "prepare_input", start, config.uuid)


def write_select_column_script(work_dir, column=4):
    # I feel bad for doing this, but I can't send single quotes into a toil command without them becoming escaped
    # so I'm just creating a script which does that
    filename = "select_column_{}.sh".format(column)
    file_location = os.path.join(work_dir, filename)
    with open(file_location, 'w') as out:
        print("#!/usr/bin/awk -f", file=out)
        print("{print $%d}" % column, file=out)
    os.chmod(file_location, 1023) #rwxrwxrwx
    return filename

{'fraction_non_alphanumeric': 0.10199166883624056, 'fraction_numerical': 0.009502733663108566, 'mean_word_length': 3.568837347606304, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '248489', 'n_tokens_mistral': 5213, 'n_tokens_neox': 4968, 'n_words': 1064}