local.py

"""
what local.py does
"""
import argparse
import os
import sys
import json
from datetime import datetime, date
from csv import DictReader
import subprocess
import requests
from Bio import Phylo

import covizu
from covizu.utils import seq_utils, gisaid_utils
from covizu.utils.progress_utils import Callback
from covizu.utils.batch_utils import (unpack_records, build_timetree, parse_alias,
                                      make_beadplots, get_mutations)
from covizu.utils.seq_utils import SC2Locator


def parse_args():
    """parse arguments"""
    parser = argparse.ArgumentParser(
        description="CoVizu analysis pipeline automation for execution on local files")

    parser.add_argument(
        "infile",
        type=str,
        help="input, path to xz-compressed JSON; if not specified, "
        "download xz file from GISAID provision feed.")
    parser.add_argument("pangolineages", type=argparse.FileType('r'),
                        help="input, CSV output generated by Pangolin")
    parser.add_argument("--outdir", type=str, default='data/',
                        help="option, path to write output files")

    parser.add_argument(
        "--bylineage",
        type=str,
        default='data/by_lineage.json',
        help="path to write JSON of features by lineage")

    parser.add_argument(
        '--lineages',
        type=str,
        default=os.path.join(
            covizu.__path__[0],
            "data/pango-designation/lineages.csv"),
        help="optional, path to CSV file containing Pango lineage designations.")

    parser.add_argument(
        '--minlen',
        type=int,
        default=29000,
        help='minimum genome length (nt)')
    parser.add_argument(
        '--mindate',
        type=str,
        default='2019-12-01',
        help='earliest possible sample collection date (ISO format, default '
        '2019-12-01')
    parser.add_argument(
        '--poisson-cutoff',
        type=float,
        default=0.001,
        help='filtering outlying genomes whose distance exceeds the upper '
        'quantile of Poisson distribution (molecular clock).  Default 0.001 '
        'corresponds to 99.9%% cutoff.')

    parser.add_argument(
        '--batchsize',
        type=int,
        default=500,
        help='number of records to batch process with minimap2')
    parser.add_argument(
        '--max-variants',
        type=int,
        default=5000,
        help='option, limit number of variants per lineage (default 5000)')

    parser.add_argument(
        "--ref",
        type=str,
        default=os.path.join(
            covizu.__path__[0],
            "data/NC_045512.fa"),
        help="path to FASTA file with reference genome")
    parser.add_argument('--mmbin', type=str, default='minimap2',
                        help="path to minimap2 binary executable")
    parser.add_argument('-mmt', "--mmthreads", type=int, default=8,
                        help="number of threads for minimap2.")

    parser.add_argument('--misstol', type=int, default=300,
                        help="maximum tolerated number of missing bases per "
                             "genome (default 300).")
    parser.add_argument(
        "--vcf",
        type=str,
        default=os.path.join(
            covizu.__path__[0],
            "data/ProblematicSites_SARS-CoV2/problematic_sites_sarsCov2.vcf"),
        help="Path to VCF file of problematic sites in SARS-COV-2 genome. "
        "Source: https://github.com/W-L/ProblematicSites_SARS-CoV2")

    parser.add_argument('--ft2bin', default='fasttree2',
                        help='path to fasttree2 binary executable')

    parser.add_argument(
        '--alias',
        type=str,
        default=os.path.join(
            covizu.__path__[0],
            "data/pango-designation/pango_designation/alias_key.json"),
        help="optional, path to JSON file containing alias.")

    parser.add_argument('--ttbin', default='treetime',
                        help='path to treetime binary executable')
    parser.add_argument('--clock', type=float, default=8e-4,
                        help='specify molecular clock rate for '
                             'constraining Treetime analysis (default 8e-4).')

    parser.add_argument('--datetol', type=float, default=0.1,
                        help='exclude tips from time-scaled tree '
                             'with high discordance between estimated and '
                             'known sample collection dates (year units,'
                             'default: 0.1)')

    parser.add_argument('--binpath', type=str, default='rapidnj',
                        help='path to RapidNJ binary executable')
    parser.add_argument('--mincount', type=int, default=500,
                        help='minimum number of variants in lineage '
                             'above which MPI processing will be used.')
    parser.add_argument('--machine_file', type=str, default='mfile',
                        help='path to machine file for MPI.')
    parser.add_argument("-n", "--nboot", type=int, default=100,
                        help="Number of bootstrap samples, default 100.")

    parser.add_argument(
        "--boot-cutoff",
        type=float,
        default=0.5,
        help="Bootstrap cutoff for consensus tree (default 0.5). "
        "Only used if --cons is specified.")

    return parser.parse_args()

def id_local_lineage(input_local_lineage, given_callback, input_header):
    """helps stream_local avoid too many branches"""
    if input_local_lineage is None:
        if given_callback:
            given_callback(
                f"Failed to retrieve lineage assignment for {input_header}",
                level='ERROR'
            )
        sys.exit()

def stream_local(
        path,
        lineage_file,
        local_regions,
        mindata,
        callback=None):
    """ Convert local FASTA file to feed-like object - replaces load_gisaid() """
    if not mindata:
        mindate = '2019-12-01'
        minlen = 29000
    else:
        minlen, mindate = mindata
    mindate = seq_utils.fromisoformat(mindate)

    # parse CSV output from Pangolin
    reader = DictReader(lineage_file)
    if ['taxon', 'lineage'] in reader.fieldnames:
        if callback:
            callback(
                "Lineage CSV header does not match expected.",
                level='ERROR')
        sys.exit()

    lineages = {}
    for row in reader:
        lineages.update({row['taxon']: row['lineage']})

    with open(path, 'r', encoding='utf-8') as local_handle:
        rejects = {'short': 0, 'baddate': 0, 'nonhuman': 0}
        for header, seq in seq_utils.iter_fasta(local_handle):
            if len(seq) < minlen:
                rejects['short'] += 1
                continue  # sequence is too short

            # hCoV-19/Canada/Qc-L00240569/2020|EPI_ISL_465679|2020-03-27
            label, accn, coldate = header.split('|')
            country = label.split('/')[1]
            if country == '' or country[0].islower():
                rejects['nonhuman'] += 1
                continue

            if coldate.count('-') != 2:
                rejects['baddate'] += 1
                continue  # incomplete collection date
            date_time = seq_utils.fromisoformat(coldate)
            if date_time < mindate or date_time > date.today():
                rejects['baddate'] += 1
                continue  # reject records with non-sensical collection date

            local_lineage = lineages.get(header, None)
            id_local_lineage(local_lineage, callback, header)

            region = local_regions.get(country.lower().replace('_', ' '), None)

            record = {
                'covv_virus_name': label,
                'covv_accession_id': accn,
                'sequence': seq,
                'covv_collection_date': coldate,
                'covv_lineage': local_lineage,
                'covv_location': country if region is None else f"{region} / {country}"}
            yield record

        if callback:
            callback("Rejected {short} short genomes\n         {baddate} records with bad "
                    "dates\n         {nonhuman} non-human genomes".format(**rejects))


def process_local(local_args, local_regions, callback=None):
    """ Analyze genome sequences from local FASTA file """
    loader = stream_local(
        local_args.infile,
        local_args.pangolineages,
        local_regions,
        (local_args.minlen,
        local_args.mindate),
        callback=callback)
    batcher = gisaid_utils.batch_fasta(loader, size=local_args.batchsize)
    aligned = gisaid_utils.extract_features(
        batcher,
        ref_file=local_args.ref,
        binpath=local_args.mmbin,
        nthread=local_args.mmthreads,
        minlen=local_args.minlen)
    filtered = gisaid_utils.filter_problematic(
        aligned,
        vcf_file=local_args.vcf,
        cutoff=local_args.poisson_cutoff,
        callback=callback)
    return gisaid_utils.sort_by_lineage(filtered, callback=callback)


if __name__ == "__main__":
    args = parse_args()
    cb = Callback()

    # check that user has loaded openmpi module
    try:
        subprocess.check_call(['mpirun', '-np', '2', 'ls'],
                              stdout=subprocess.DEVNULL)
    except FileNotFoundError:
        cb.callback(
            "mpirun not loaded - run `module load openmpi/gnu`",
            level='ERROR')
        sys.exit()

    # check that the user has included submodules
    if (
        not os.path.exists(
            os.path.join(
            covizu.__path__[0],
            "data/pango-designation/lineages.csv")) or not os.path.exists(
                os.path.join(
                    covizu.__path__[0],
                    "data/ProblematicSites_SARS-CoV2/problematic_sites_sarsCov2.vcf"))):
        try:
            subprocess.check_call(
                "git submodule init; git submodule update", shell=True)
        except subprocess.CalledProcessError as e:
            cb.callback(f"Error adding the required submodules, exit {e}")
            sys.exit()

    # update submodules
    try:
        subprocess.check_call(
            "git submodule foreach git pull origin master",
            shell=True)
    except subprocess.CalledProcessError as e:
        cb.callback(f"Error updating submodules, exit {e}")
        sys.exit()

    response = requests.get(
        "https://geoenrich.arcgis.com/arcgis/rest/services/World/geoenrichmentserver/" +
        "Geoenrichment/countries?f=pjson", timeout=300)
    res = response.json()
    regions = {c['name'].lower(): c['continent'] for c in res['countries']} | {
        c['abbr3'].lower(): c['continent'] for c in res['countries']}

    by_lineage = process_local(args, regions, cb.callback)
    with open(args.bylineage, 'w', encoding='utf-8') as handle:
        # export to file to process large lineages with MPI
        json.dump(by_lineage, handle)

    # reconstruct time-scaled tree
    timetree, residuals = build_timetree(by_lineage, args, cb.callback)
    timestamp = datetime.now().isoformat().split('.')[0]
    nwk_file = os.path.join(args.outdir, f'timetree.{timestamp}.nwk')
    with open(nwk_file, 'w', encoding='utf-8') as handle:
        Phylo.write(timetree, file=handle, format='newick')

    # generate beadplots and serialize to file
    result, infection_prediction = make_beadplots(
        by_lineage, args, cb.callback, initial_time=cb.initial_time.timestamp())
    outfile = os.path.join(args.outdir, f'clusters.{timestamp}.json')
    with open(outfile, 'w', encoding='utf-8') as handle:  # serialize results to JSON
        json.dump(result, fp=handle)

    # get mutation info
    locator = SC2Locator()
    mutations = {}
    for lineage, features in get_mutations(by_lineage).items():
        annots = {
            locator.parse_mutation(f): freq for f,
            freq in features.items()}
        mutations.update(
            {lineage: {a: freq for a, freq in annots.items() if a is not None}})

    # write data stats
    dbstat_file = os.path.join(args.outdir, f'dbstats.{timestamp}.json')

    alias = parse_alias(args.alias)

    with open(dbstat_file, 'w', encoding='utf-8') as handle:
        nseqs = sum((len(rows) for rows in by_lineage.values()))
        val = {
            'lastupdate': timestamp.split('T')[0],
            'noseqs': nseqs,
            'lineages': {}
        }
        for lineage, samples in by_lineage.items():
            prefix = lineage.split('.')[0]
            lname = lineage.replace(
                prefix,
                alias[prefix]) if lineage.lower() not in [
                'unclassifiable',
                'unassigned'] and not prefix.startswith('X') and alias[prefix] != '' else lineage
            samples = unpack_records(samples)
            ndiffs = [len(x['diffs']) for x in samples]
            val['lineages'][lineage] = {
                'nsamples': len(samples),
                'lastcoldate': max(x['covv_collection_date'] for x in samples),
                'residual': residuals.get(lineage, 0),
                'max_ndiffs': max(ndiffs),
                'mean_ndiffs': sum(ndiffs) / len(ndiffs),
                'mutations': mutations[lineage],
                'infections': infection_prediction[lineage],
                'raw_lineage': lname
            }
        json.dump(val, handle)

    cb.callback("All done!")