<commit_before>#!/usr/bin/env python

# native Python imports
import os.path
import shutil
import sys
import sqlite3

import annotations
import subprocess
from cluster_helper.cluster import cluster_view
import database as gemini_db
from gemini_load_chunk import GeminiLoader
import uuid
import time
import datetime


def load(parser, args):
    if (args.db is None or args.vcf is None):
        parser.print_help()
        exit("ERROR: load needs both a VCF file and a database file\n")

    annos = annotations.get_anno_files(args)
    # force skipping CADD and GERP if the data files have not been installed
    if args.skip_cadd is False:
        if 'cadd_score' not in annos:
            sys.stderr.write("\nCADD scores are not being loaded because the"
            " annotation file could not be found.\n"
            "`Run gemini update --dataonly --extra cadd_score`"
            " to install the annotation file.\n\n")
            args.skip_cadd = True
        else:
            sys.stderr.write("CADD scores are being loaded (to skip use:--skip-cadd).\n")
    if args.skip_gerp_bp is False:
        if 'gerp_bp' not in annos:
            sys.stderr.write("\nGERP per bp is not being loaded because the annotation file"
                        " could not be found.\n    Run `gemini update --dataonly --extra gerp_bp`"
                        " to install the annotation file.\n\n")
            args.skip_gerp_bp = True
        else:
            sys.stderr.write("GERP per bp is being loaded (to skip use:--skip-gerp-bp).\n")
    # collect of the the add'l annotation files
    annotations.load_annos(args)

    if args.scheduler:
        load_ipython(args)
    elif args.cores > 1:
        load_multicore(args)
    else:
        load_singlecore(args)

    if not args.no_bcolz:
        from gemini.gemini_bcolz import create
        create(args.db)


def load_singlecore(args):
    # create a new gemini loader and populate
    # the gemini db and files from the VCF
    gemini_loader = GeminiLoader(args)
    gemini_loader.store_resources()
    gemini_loader.store_version()
    gemini_loader.store_vcf_header()
    gemini_loader.populate_from_vcf()
    gemini_db.add_max_aaf(gemini_loader.c)

    if not args.skip_gene_tables and not args.test_mode:
        gemini_loader.update_gene_table()
    if not args.test_mode:
        gemini_loader.build_indices_and_disconnect()

    if not args.no_genotypes and not args.no_load_genotypes:
        gemini_loader.store_sample_gt_counts()


def load_multicore(args):
    grabix_file = bgzip(args.vcf)
    chunks = load_chunks_multicore(grabix_file, args)
    merge_chunks_multicore(chunks, args)

def load_ipython(args):
    grabix_file = bgzip(args.vcf)
    with cluster_view(*get_ipython_args(args)) as view:
        chunks = load_chunks_ipython(grabix_file, args, view)
        merge_chunks_ipython(chunks, args, view)

def merge_chunks(chunks, db, kwargs):
    cmd = get_merge_chunks_cmd(chunks, db, tempdir=kwargs.get("tempdir"))
    print "Merging chunks."
    subprocess.check_call(cmd, shell=True)
    cleanup_temp_db_files(chunks)
    return db

def get_merge_chunks_cmd(chunks, db, tempdir=None):
    chunk_names = ""
    for chunk in chunks:
        chunk_names += " --chunkdb  " + chunk

    tempdir_string = ""
    if tempdir is not None:
        tempdir_string = " --tempdir " + tempdir

    return "gemini merge_chunks {chunk_names} {tempdir_string} --db {db}".format(**locals())


def finalize_merged_db(tmp_db, db):
    ts = time.time()
    st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
    print st, "indexing final database."

    main_conn = sqlite3.connect(tmp_db)
    main_conn.isolation_level = None
    main_curr = main_conn.cursor()
    main_curr.execute('PRAGMA synchronous = OFF')
    main_curr.execute('PRAGMA journal_mode=MEMORY')

    gemini_db.add_max_aaf(main_curr)
    gemini_db.create_indices(main_curr)

    main_conn.commit()
    main_curr.close()

    shutil.move(tmp_db, db)

def merge_chunks_ipython(chunks, args, view):
    if len(chunks) == 1:
        finalize_merged_db(chunks[0], args.db)
        return args.db
    else:
        ts = time.time()
        st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
        print st, "merging", len(chunks), "chunks."
        sub_merges = get_chunks_to_merge(chunks)
        tmp_dbs = get_temp_dbs(len(sub_merges), os.path.dirname(sub_merges[0][0]))
        merge_args = {"tempdir": args.tempdir}
        view.map(merge_chunks, sub_merges, tmp_dbs, [merge_args] * len(sub_merges))
        merge_chunks_ipython(tmp_dbs, args, view)

def merge_chunks_multicore(chunks, args):
    if len(chunks) <= 1:
        finalize_merged_db(chunks[0], args.db)
        return args.db
    else:
        ts = time.time()
        st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
        print st, "merging", len(chunks), "chunks."
        procs = []
        sub_merges = get_chunks_to_merge(chunks)<commit_msg>Avoid use of tempdir for chunks when using IPython<commit_after>#!/usr/bin/env python

# native Python imports
import os.path
import shutil
import sys
import sqlite3

import annotations
import subprocess
from cluster_helper.cluster import cluster_view
import database as gemini_db
from gemini_load_chunk import GeminiLoader
import uuid
import time
import datetime


def load(parser, args):
    if (args.db is None or args.vcf is None):
        parser.print_help()
        exit("ERROR: load needs both a VCF file and a database file\n")

    annos = annotations.get_anno_files(args)
    # force skipping CADD and GERP if the data files have not been installed
    if args.skip_cadd is False:
        if 'cadd_score' not in annos:
            sys.stderr.write("\nCADD scores are not being loaded because the"
            " annotation file could not be found.\n"
            "`Run gemini update --dataonly --extra cadd_score`"
            " to install the annotation file.\n\n")
            args.skip_cadd = True
        else:
            sys.stderr.write("CADD scores are being loaded (to skip use:--skip-cadd).\n")
    if args.skip_gerp_bp is False:
        if 'gerp_bp' not in annos:
            sys.stderr.write("\nGERP per bp is not being loaded because the annotation file"
                        " could not be found.\n    Run `gemini update --dataonly --extra gerp_bp`"
                        " to install the annotation file.\n\n")
            args.skip_gerp_bp = True
        else:
            sys.stderr.write("GERP per bp is being loaded (to skip use:--skip-gerp-bp).\n")
    # collect of the the add'l annotation files
    annotations.load_annos(args)

    if args.scheduler:
        load_ipython(args)
    elif args.cores > 1:
        load_multicore(args)
    else:
        load_singlecore(args)

    if not args.no_bcolz:
        from gemini.gemini_bcolz import create
        create(args.db)


def load_singlecore(args):
    # create a new gemini loader and populate
    # the gemini db and files from the VCF
    gemini_loader = GeminiLoader(args)
    gemini_loader.store_resources()
    gemini_loader.store_version()
    gemini_loader.store_vcf_header()
    gemini_loader.populate_from_vcf()
    gemini_db.add_max_aaf(gemini_loader.c)

    if not args.skip_gene_tables and not args.test_mode:
        gemini_loader.update_gene_table()
    if not args.test_mode:
        gemini_loader.build_indices_and_disconnect()

    if not args.no_genotypes and not args.no_load_genotypes:
        gemini_loader.store_sample_gt_counts()


def load_multicore(args):
    grabix_file = bgzip(args.vcf)
    chunks = load_chunks_multicore(grabix_file, args)
    merge_chunks_multicore(chunks, args)

def load_ipython(args):
    grabix_file = bgzip(args.vcf)
    with cluster_view(*get_ipython_args(args)) as view:
        chunks = load_chunks_ipython(grabix_file, args, view)
        merge_chunks_ipython(chunks, args, view)

def merge_chunks(chunks, db, kwargs):
    cmd = get_merge_chunks_cmd(chunks, db, tempdir=kwargs.get("tempdir"))
    print "Merging chunks."
    subprocess.check_call(cmd, shell=True)
    cleanup_temp_db_files(chunks)
    return db

def get_merge_chunks_cmd(chunks, db, tempdir=None):
    chunk_names = ""
    for chunk in chunks:
        chunk_names += " --chunkdb  " + chunk

    tempdir_string = ""
    if tempdir is not None:
        tempdir_string = " --tempdir " + tempdir

    return "gemini merge_chunks {chunk_names} {tempdir_string} --db {db}".format(**locals())


def finalize_merged_db(tmp_db, db):
    ts = time.time()
    st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
    print st, "indexing final database."

    main_conn = sqlite3.connect(tmp_db)
    main_conn.isolation_level = None
    main_curr = main_conn.cursor()
    main_curr.execute('PRAGMA synchronous = OFF')
    main_curr.execute('PRAGMA journal_mode=MEMORY')

    gemini_db.add_max_aaf(main_curr)
    gemini_db.create_indices(main_curr)

    main_conn.commit()
    main_curr.close()

    shutil.move(tmp_db, db)

def merge_chunks_ipython(chunks, args, view):
    if len(chunks) == 1:
        finalize_merged_db(chunks[0], args.db)
        return args.db
    else:
        ts = time.time()
        st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
        print st, "merging", len(chunks), "chunks."
        sub_merges = get_chunks_to_merge(chunks)
        tmp_dbs = get_temp_dbs(len(sub_merges), os.getcwd())
        merge_args = {"tempdir": args.tempdir}
        view.map(merge_chunks, sub_merges, tmp_dbs, [merge_args] * len(sub_merges))
        merge_chunks_ipython(tmp_dbs, args, view)

def merge_chunks_multicore(chunks, args):
    if len(chunks) <= 1:
        finalize_merged_db(chunks[0], args.db)
        return args.db
    else:
        ts = time.time()
        st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
        print st, "merging", len(chunks), "chunks."
        procs = []
        sub_merges = get_chunks_to_merge(chunks)
{'fraction_non_alphanumeric': 0.10204897551224387, 'fraction_numerical': 0.0015992003998001, 'mean_word_length': 2.9992006394884094, 'pattern_counts': {'":': 2, '<': 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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1837421', 'n_tokens_mistral': 3460, 'n_tokens_neox': 3366, 'n_words': 791}