<jupyter_start><jupyter_text>One-click "Web of Life" builder <NAME>, 2017<jupyter_code>import re
import json
import tarfile
from os import (chdir, close, getcwd, listdir, makedirs,
                remove, symlink, wait4, write)
from os.path import join, isdir, isfile
from tempfile import mkdtemp, mkstemp
from copy import deepcopy
from subprocess import Popen, PIPE, STDOUT
import multiprocessing as mp
import pandas as pd
from skbio import io, TreeNode
class Configuration():
    """Configurations of this task.
    """
    def __init__(self):
        """Define attributes and their default values.
        """
        # working directory
        self.wkdir = '.'

        # compute resources
        self.cpus = 0  # number of CPU cores to use (0 for all)
        self.intermdir = ''  # directory to store intermediate files
                             # generated by external programs

        # intermediate files and directories
        self.genomes_dir = ''  # genome sequences in FASTA format
        self.proteins_dir = ''  # protein sequences per genome in FASTA format
        self.species_tree_fp = ''  # species tree in Newick format
        self.gene_families_dir = ''  # protein sequences per family in FASTA format
        self.gene_trees_dir = ''  # gene family trees in Newick format
        self.gene_networks_dir = ''  # gene family networks

        # external programs and databases
        self.py2_env = 'py2'  # Python 2 virtual environment
        self.prodigal = 'prodigal'
        self.phylophlan = 'phylophlan.py'
        self.phylophlan_data_dir = 'data'
        self.hmmsearch = 'hmmsearch'
        self.pfam_db = 'Pfam-A.hmm.gz'
        self.phylomizer = 'python phylomizer.py'
        self.muscle = 'muscle'
        self.mafft = 'mafft'
        self.kalign = 'kalign'
        self.t_coffee = 't_coffee'
        self.trimal = 'trimal'
        self.readal = 'readal'
        self.fasttree = 'FastTree'
        self.optroot = 'OptRoot'
        self.rangerdtl = 'Ranger-DTL'
        self.aggregateranger = 'AggregateRanger'

        # define datatypes of non-str attributes
        self.int_attrs = ['cpus']
        self.float_attrs = []
        self.bool_attrs = []
def read_config(config_fp):
    """Read configuration file.

    Parameters
    ----------
    config_fp : str
        path to configuration file

    Raises
    ------
    ValueError
        if configuration file is incorrectly formatted
        if configuration file contains invalid keys

    Notes
    -----
    allowed format of the configuration file:
        key = value
        key      =      value
          key=value
        key: value
        key:value  # comment
        key=
        # this is comment
    prohibited format of the configuration file:
        key = foo = bar
        key : foo : bar
    """
    global cfg
    with open(config_fp, 'r') as f:
        for line in f:
            line = line.strip().split('#')[0].rstrip()
            if line:
                if '=' in line:
                    key, value = [x.strip() for x in line.split('=')]
                elif ':' in line:
                    key, value = [x.strip() for x in line.split(':')]
                if key and hasattr(cfg, key):
                    for dtype in 'int', 'float', 'bool':
                        if key in getattr(cfg, '%s_attrs' % dtype):
                            value = eval('%s(value)' % dtype)
                    setattr(cfg, key, value)
def run(cmd, screen=False):
    """Run a Bash command and display screen output in real time.

    Parameters
    ----------
    cmd : str
        Bash command to benchmark
    screen : bool
        show screen output in real time

    Returns
    -------
    int
        return code
    """
    kwargs = {'stdout': PIPE, 'stderr': PIPE} if not screen else {}
    p = Popen(cmd, shell=True, **kwargs)
    rc = wait4(p.pid, 0)[1]
    return rc


def bench(cmd, screen=False):
    """Benchmark the time and memory usage of a Bash command.

    Parameters
    ----------
    cmd : str
        Bash command to benchmark
    screen : bool
        show screen output in real time

    Returns
    -------
    resource usage object
        resource usage information of the process

    Raises
    ------
    ValueError
        if return code is not 0

    Notes
    -----
    definition of resource usage object:
        https://docs.python.org/3/library/resource.html#resource.getrusage
    three most important attributes:
        ru_utime (0) : float
            user time in seconds
        ru_stime (1) : float
            system time in seconds
        ru_maxrss (2) : int
            maximum resident set size (peak memory usage) in kbytes
    """
    kwargs = {'stdout': PIPE, 'stderr': PIPE} if not screen else {}
    p = Popen(cmd, shell=True, **kwargs)
    rc, ru = wait4(p.pid, 0)[1:3]
    if rc != 0:
        raise ValueError('Benchmark failed.')
    return ru


def run0(cmd):
    """(obsolete) Run a Bash command and display screen output in real time.

    Parameters
    ----------
    cmd : str
        Bash command to benchmark

    Returns
    -------
    int
        return code
    """
    p = Popen(cmd, shell=True, stdout=PIPE, stderr=STDOUT)
    while True:
        line, rc = p.stdout.readline(), p.poll()
        if rc is not None:
            return rc
        if line:
            print(line.decode('utf-8').rstrip('\r\n'))

            
def bench0(cmd, screen=False):
    """(obsolete) Benchmark the time and memory usage of a Bash command.

    Parameters
    ----------
    cmd : str
        Bash command to benchmark
    screen : bool
        show screen output in real time

    Returns
    -------
    list of 4 floats
        wall (elapsed) time, user time and system time in seconds, maximum
        memory usage (resident set size) in kbytes

    Raises
    ------
    ValueError
        if return code is not 0
    """
    cmd = '$(which time) -f "%E,%U,%S,%M" bash -c "' + cmd + '"'
    p = Popen(cmd, shell=True, stdout=PIPE, stderr=STDOUT)
    last_line = ''
    while True:
        line, rc = p.stdout.readline(), p.poll()
        if rc is not None:
            if rc != 0:
                raise ValueError('Benchmark failed.')
            break
        if line:
            if screen == True:
                print(last_line)
            last_line = line.decode('utf-8').rstrip('\r\n')
    wall, usr, sys, ram = last_line.split(',')
    l = wall.split(':')
    wall, exp = 0.0, 0
    while l:
        wall += float(l.pop()) * (60 ** exp)
        exp += 1
    return [wall, float(usr), float(sys), float(ram) / 4]


def cat(fname):
    """Get the proper cat command by file extension name.

    Parameters
    ----------
    fname : str
        file path to read

    Returns
    -------
    str
        cat command
    """
    cats = {'gz': 'zcat', 'bz2': 'bzcat', 'xz': 'xzcat', 'lz': 'lzcat', 'lzma': 'lzcat'}
    ext = fname.split('.')[-1]
    return cats[ext] if ext in cats else 'cat'


def read(fname):
    """Read the content of a compressed or regular file.

    Parameters
    ----------
    fname : str
        file path to read

    Returns
    -------
    iterable of str
        content of the file
    """
    p = Popen(cat(fname) + ' ' + fname, shell=True, stdout=PIPE, stderr=STDOUT)
    return p.stdout
def check_file(fname):
    """Test if output file exists.

    Parameters
    ----------
    fname : str
        file name to be generated in this step
        e.g., 'species_tree.nwk' is the file name for the PhyloPhlAn step

    Returns
    -------
    bool
        if target file exists
    """
    global cfg
    stem = fname.split('.')[0]
    fp = getattr(cfg, '%s_fp' % stem)
    if not fp:
        fp = join(cfg.wkdir, fname)
        setattr(cfg, '%s_fp' % stem, fp)
    return isfile(fp)


def check_dir(dname):
    """Test if output directory exists, and create it if not.

    Parameters
    ----------
    dname : str
        directory name to be generated in this step
        e.g., 'faa' is the dname for the Prodigal step

    Returns
    -------
    bool
        if target directory exists
    """
    global cfg
    dp = getattr(cfg, '%s_dir' % dname)
    if not dp:
        dp = join(cfg.wkdir, dname)
        setattr(cfg, '%s_dir' % dname, dp)
    if isdir(dp):
        return True
    else:
        makedirs(dp)
        return False


def write_benchmarks(benchmarks, output_fp):
    """Write benchmark results to file.

    Parameters
    ----------
    benchmarks : dict of resource usage object(s)
        resource usage information per run
        each element can be a single object, or an iterable of
        multiple objects, in latter case times will be summed
        and memory will be the max
    output_fp : str
        output file path
    """
    with open(output_fp, 'w') as f:
        for key, bm in sorted(benchmarks.items()):
            utime, stime, maxrss = 0.0, 0.0, 0
            if hasattr(bm, 'ru_utime'):
                utime, stime, maxrss = bm.ru_utime, bm.ru_stime, bm.ru_maxrss
            else:
                for x in bm:
                    utime += x.ru_utime
                    stime += x.ru_stime
                    maxrss = max(maxrss, x.ru_maxrss)
            f.write('%s\t%.3f\t%.3f\t%i\n' % (key, utime, stime, maxrss))<jupyter_output><empty_output><jupyter_text>Preparations<jupyter_code># read configuration
global cfg
cfg = Configuration()
config_fp = 'config.txt'
if isfile(config_fp):
    read_config(config_fp)
# process configuration
if cfg.cpus == 0:
    cfg.cpus = mp.cpu_count()
if cfg.intermdir == '':
    cfg.intermdir = join(cfg.wkdir, 'interm')
if not isdir(cfg.intermdir):
    makedirs(cfg.intermdir)<jupyter_output><empty_output><jupyter_text>Step 1: Infer protein-coding genes from genomes, using ProdigalInput: Genome sequences (`G.fna`)Output: Protein sequences per genome (`G.faa`)Intermediate files: - Protein sequences per genome (`G.faa`) - Protein-coding sequences per genome (`G.ffn`) - Genome annotation tables (`G.gff`)<jupyter_code>def run_prodigal(fna_fp, faa_fp, ffn_fp, gff_fp, title=None):
    """Run prodigal to predict protein-coding genes of the genome.

    Parameters
    ----------
    fna_fp : str
        file path to input genome sequence in FASTA format
    faa_fp : str
        file path to output protein sequences in FASTA format
    ffn_fp : str
        file path to output nucleotide sequences in FASTA format
    gff_fp : str
        file path to output genome annotation table in GFF3 format
    title : str
        title of the task to display (optional)

    Returns
    -------
    resource usage object
        resource usage information of this prodigal run
    """
    if title:
        print('Running Prodigal on %s...' % title)
    cmd = ('%s %s | %s -c -q -f gff -a %s -d %s -o %s'
           % (cat(fna_fp), fna_fp, cfg.prodigal, faa_fp, ffn_fp, gff_fp))
    return bench(cmd)
if not check_dir('proteins'):
    makedirs(join(cfg.intermdir, 'prodigal'))
    for dname in 'faa', 'ffn', 'gff':
        makedirs(join(cfg.intermdir, 'prodigal', dname))
    bms = {}
    pool = mp.Pool(cfg.cpus)
    for fname in listdir(cfg.genomes_dir):
        l = fname.split('.')
        if 'fna' in l:
            gid = '.'.join(l[0:l.index('fna')])
            args = [join(cfg.genomes_dir, fname)]
            for dname in 'faa', 'ffn', 'gff':
                args.append(join(cfg.intermdir, 'prodigal', dname,
                                 '%s.%s' % (gid, dname)))
            args.append(gid)
            bms[gid] = pool.apply_async(run_prodigal, args=args)
    pool.close()
    pool.join()
    for gid in bms:
        bms[gid] = bms[gid].get()
        faa_fp = join(cfg.intermdir, 'prodigal', 'faa', '%s.faa' % gid)
        if isfile(faa_fp):
            symlink(faa_fp, join(cfg.proteins_dir, '%s.faa' % gid))
    write_benchmarks(bms, join(cfg.intermdir, 'prodigal', 'benchmark.txt'))<jupyter_output><empty_output><jupyter_text>Step 2: Extract marker genes and build species tree, using PhyloPhlAnInput: Protein sequences per genome (`G.faa`).Output: Species treeIntermediate files: - Protein sequences per marker (`p.faa`). - Alignments of protein sequences per marker (`p.aln`). - Protein-to-marker dictionaries (`G.b6o`). - Marker-to-proteins dictionary (`up2prots.txt`). - Genome-to-proteins dictionary (`orgs2prots.txt`).<jupyter_code>if not check_file('species_tree.nwk'):
    pwd = getcwd()
    cur_dir = join(cfg.intermdir, 'phylophlan')
    makedirs(cur_dir)
    for dname in 'input', 'output', 'temp':
        makedirs(join(cur_dir, dname))
    symlink(cfg.phylophlan_data_dir, join(cur_dir, 'data'))
    symlink(cfg.proteins_dir, join(cur_dir, 'input', 'all'))
    chdir(cur_dir)
    cmd_fp = join(cfg.wkdir, 'tmp.sh')
    with open(cmd_fp, 'w') as f:
        f.write('source activate %s\n' % cfg.py2_env)
        f.write('%s -u all --nproc %s --c_dat temp\n'
                % (join(phylophlan_dir, 'phylophlan.py'), cfg.cpus))
        f.write('source deactivate\n')
    print('Running PhyloPhlAn...')
    bms = {'all': bench('bash %s' % cmd_fp, True)}
    symlink(join(cur_dir, 'output', 'all', 'all.tree.nwk'), cfg.species_tree_fp)
    write_benchmarks(bms, join(cur_dir, 'benchmark.txt'))
    remove(cmd_fp)
    chdir(pwd)<jupyter_output><empty_output><jupyter_text>Step 3: Identify gene families using HMMER against PfamInput: Protein sequences per genome (`G.faa`).Output: - Protein sequences per family (`PF..faa`). - Statistics of families (`statistics.txt`).Intermediate files: HMMER search tabular reports (`G.hm3`).<jupyter_code>def run_hmmsearch(hmmfile, seqdb, outfile, title=None):
    """Run hmmsearch to search profiles within sequence database

    Parameters
    ----------
    hmmfile : str
        file path to database containing HMM profiles of reference protein
        families (e.g., Pfam)
    seqdb : str
        file path to query protein sequences in FASTA format
    outfile : str
        output file in standard HMMER3 tabular format, fields:
            target name, accession, tlen, query name, accession, qlen,
            full sequence: E-value, score, bias, #, of
            this domain: c-Evalue, i-Evalue, score, bias
            hmm_coord: from, to
            ali_coord: from, to
            env_coord: from, to
            acc, description of target
    title : str
        title of the task to display (optional)

    Returns
    -------
    resource usage object
        resource usage information of this hmmsearch run
    """
    if title:
        print('Running hmmsearch on %s...' % title)
    cmd = ('%s --notextw --noali --cut_ga --cpu 1 --domtblout %s %s %s > /dev/null'
           % (cfg.hmmsearch, outfile, hmmfile, seqdb))
    return bench(cmd)
if not check_dir('gene_families'):
    makedirs(join(cfg.intermdir, 'pfam'))
    bms = {}
    pool = mp.Pool(cfg.cpus)
    for fname in listdir(cfg.proteins_dir):
        if fname.endswith('.faa'):
            gid = fname[:-4]
            args = (cfg.pfam_db,
                    join(cfg.proteins_dir, '%s.faa' % gid),
                    join(cfg.intermdir, 'pfam', '%s.hm3' % gid),
                    gid)
            bms[gid] = pool.apply_async(run_hmmsearch, args=args)
    pool.close()
    pool.join()
    for gid in bms:
        bms[gid] = bms[gid].get()
    write_benchmarks(bms, join(cfg.intermdir, 'pfam', 'benchmark.txt'))
def parse_hm3(hm3_fp, multi_assign=True, evalue=None,
              plen=None, dlen=None, pcov=None, dcov=None):
    """Parse HMMER result to identify protein-to-domain matches.

    Parameters
    ----------
    hm3_fp : str
        file path to input HMMER tabular report
    multi_assign : bool
        whether one protein can be associated with multiple domains (default: True)
    evalue : float
        full sequence E-value cutoff (optional)
    plen : int
        protein sequence length cutoff (optional)
    dlen : int
        domain profile length cutoff (optional)
    pcov : float
        protein sequence percent coverage cutoff (optional)
    dcov : float
        domain profile percent coverage cutoff (optional)

    Returns
    -------
    tuple of 2 dicts of set of str
        domain-to-proteins dictionary
        protein-to-domains dictionary
    """
    dom2prots, prot2doms = {}, {}
    min_evalue = {}
    with open(hm3_fp, 'r') as f:
        for line in f:
            if line.startswith('#'):
                continue
            l = line.strip().split()
            if evalue and float(l[6]) > evalue:
                continue
            if plen and int(l[2]) < plen:
                continue
            if dlen and int(l[5]) < dlen:
                continue
            if pcov and (int(l[18]) - int(l[17]) + 1) / float(l[2]) * 100 < pcov:
                continue
            if dcov and (int(l[16]) - int(l[15]) + 1) / float(l[5]) * 100 < dcov:
                continue
            prot, dom = l[0], l[4]
            if prot not in min_evalue or min_evalue[prot][1] > float(l[6]):
                min_evalue[prot] = (dom, float(l[6]))
            if prot in prot2doms:
                prot2doms[prot].add(dom)
            else:
                prot2doms[prot] = set([dom])
            if dom in dom2prots:
                dom2prots[dom].add(prot)
            else:
                dom2prots[dom] = set([prot])
            if multi_assign == False:
                for prot in prot2doms:
                    prot2doms[prot] = set([min_evalue[prot][0]])
    return (dom2prots, prot2doms)
if not listdir(cfg.gene_families_dir) and isdir(join(cfg.intermdir, 'pfam')):
    fam2prots = {}
    fam2gs = {}
    for fname in listdir(join(cfg.intermdir, 'pfam')):
        if fname.endswith('.hm3'):
            gid = fname[:-4]
            print('Inferring gene families for %s...' % gid)
            f2p = parse_hm3(join(cfg.intermdir, 'pfam', fname), evalue=1e-50, dcov=90)[0]
            prots = {x.metadata['id']: x for x in
                     io.read(join(cfg.proteins_dir, '%s.faa' % gid), format='fasta')}
            for fam in f2p:
                # optional: only keep single-copy genes
                # if len(f2p[fam]) > 1:
                #     continue
                prots2add = [prots[x] for x in prots if x in f2p[fam]]
                if fam in fam2prots:
                    fam2prots[fam] += prots2add
                    fam2gs[fam].add(gid)
                else:
                    fam2prots[fam] = prots2add
                    fam2gs[fam] = set([gid])
    stats = []
    for fam, prots in fam2prots.items():
        # families represented in <10 genomes are excluded
        if len(fam2gs[fam]) < 10:
            continue
        print('Writing gene family %s...' % fam)
        def outprots():
            for prot in prots:
                prot.metadata.pop('description', None)
                yield prot
        io.write(outprots(), format='fasta', into=join(cfg.gene_families_dir, '%s.faa' % fam))
        stats.append([fam, len(fam2gs[fam]), len(prots)])
    df = pd.DataFrame(stats, columns=['family', 'genomes', 'proteins'])
    df.set_index('family', inplace=True)
    df.sort_values(by=['genomes', 'proteins'], ascending=[False, False], inplace=True)
    df.to_csv(join(cfg.gene_families_dir, 'statistics.txt'), sep='\t')<jupyter_output><empty_output><jupyter_text>Step 4: Build gene family trees, using PhylomizerInput: Protein sequences per family (`PF.faa`).Output: - Gene tree per family (`PF.nwk`).Intermediate files: - Clean alignment of protein sequences per family in Phylip (`PF.alg.clean`) and FASTA (`PF.alg.clean.fa`) formats. - Maximum likelihood gene family trees using WAG or JTT models in Newick format (`PF.tree.fasttree.ml.wag/jtt.nw`). - Ranking of per-model log likelihoods (`PF.tree.fasttree.rank.ml`).<jupyter_code># Phylomizer configuration file, inherited from the "KMM" protocol,
# except for that PhyML is replaced by FastTree
pcfg = """
verbose             parameter    1
residue_datatype    parameter    protein
force_seed_sequence parameter    True
alignment           mode         kalign muscle mafft
consensus           mode         m_coffee
trimming            mode         trimal
both_direction      parameter    True
min_seqs            parameter    10
in_letter           parameter    U:B
in_letter           parameter    O:Z
muscle              binary       %s
muscle_params       parameter    
mafft               binary       %s
mafft_params        parameter    --auto
kalign              binary       %s
kalign_params       parameter    -f fasta
m_coffee            binary       %s
m_coffee_params     parameter    -n_core 1 -output fasta -quiet
trimal              binary       %s
trimal_params       parameter    -phylip -gt 0.1
trimal_compare      parameter    -ct 0.1667
readal              binary       %s
tree                mode         fasttree
evol_models         parameter    wag jtt
numb_models         parameter    2
tree_approach       mode         ml
ml                  parameter    
fasttree            binary       %s
fasttree_params     parameter    -gamma
""" % (cfg.muscle, cfg.mafft, cfg.kalign, cfg.t_coffee, cfg.trimal,
       cfg.readal, cfg.fasttree)
def run_phylomizer(faa_fp, cfg_fp, out_dir, title=None):
    """Run Phylomizer to build phylogenetic tree of a gene family.

    Parameters
    ----------
    faa_fp : str
        file path to input protein sequences in FASTA format
    cfg_fp : str
        file path to Phylomizer configuration
    out_dir : str
        path to output directory
    title : str
        title of the task to display (optional)

    Returns
    -------
    resource usage object
        resource usage information of this Phylomizer run
    """
    if title:
        print('Running Phylomizer on %s...' % title)
    f, fp = mkstemp()
    close(f)
    with open(fp, 'w') as f:
        f.write('source activate %s\n' % cfg.py2_env)
        f.write('%s -i %s --steps alignments trees -c %s -o %s\n'
                % (cfg.phylomizer, faa_fp, cfg_fp, out_dir))
        f.write('source deactivate\n')
    bm = bench('bash %s' % fp)
    remove(fp)
    return bm
if not check_dir('gene_trees'):
    makedirs(join(cfg.intermdir, 'phylomizer'))
    with open(join(cfg.intermdir, 'phylomizer', 'config.txt'), 'w') as f:
        f.write(pcfg.strip())
    bms = {}
    pool = mp.Pool(cfg.cpus)
    for fname in listdir(cfg.gene_families_dir):
        if fname.endswith('.faa'):
            fid = fname[:-4]
            args = (join(cfg.gene_families_dir, fname),
                    join(cfg.intermdir, 'phylomizer', 'config.txt'),
                    join(cfg.intermdir, 'phylomizer', fid),
                    fid)
            bms[fid] = pool.apply_async(run_phylomizer, args=args)
    pool.close()
    pool.join()
    for fid in bms:
        bms[fid] = bms[fid].get()
        fid_stem = fid.split('.')[0]
        res_dir = join(cfg.intermdir, 'phylomizer', fid)
        ml_fp = join(res_dir, '%s.tree.fasttree.rank.ml' % fid_stem)
        if isfile(ml_fp):
            with open(ml_fp, 'r') as f:
                models = dict([x.split() for x in f.read().splitlines()])
            if models:
                # select the best model, which has the highest log likelihood
                best_model = sorted(models.items(),
                                    key=lambda x: float(x[1]), reverse=True)[0][0]
                best_tree_fp = join(res_dir, '%s.tree.fasttree.ml.%s.nw'
                                    % (fid_stem, best_model.lower()))
                symlink(best_tree_fp, join(cfg.gene_trees_dir, '%s.nwk' % fid))
    write_benchmarks(bms, join(cfg.intermdir, 'phylomizer', 'benchmark.txt'))<jupyter_output><empty_output><jupyter_text>Step 5: Inferring horizontal gene transfers, using Ranger-DTLInput: - Species tree (`species_tree.nwk`). - Gene family trees (`PF.nwk`).Output: - Reconciliation result per gene family (`PF/reconc.txt`).<jupyter_code>def run_rangerdtl(input_fp, result_dir, reps=100, cutoff=0.75, title=None):
    """Run RANGER-DTL to infer non-vertical evolutionary history of a gene family.

    Parameters
    ----------
    input_fp : str
        input file path (species tree followed by gene tree)
    result_dir : str
        directory to save intermediate files
    reps : int
        number of replicates to run (default: 100)
    cutoff : float
        probability cutoff for reporting HGT events
    title : str
        title of the task to display (optional)

    Returns
    -------
    list of resource usage objects
        resource usage information of these OptRoot, Ranger-DTL and
        AggregateRanger runs
    """
    if title:
        print('Running Ranger-DTL on %s...' % title)

    # read species tree
    with open(input_fp, 'r') as f:
        stree_str = f.readline().rstrip('\r\n')

    # run OptRoot
    # "-r --seed n" will print only one optimal rooting out of multiple
    cmd = ('%s -i %s -o %s -r --seed 1'
           % (cfg.optroot, input_fp, join(result_dir, 'optroot.txt')))
    bms = [bench(cmd)]
    
    # read (first) optimal rooting
    gtree_str = ''
    with open(join(result_dir, 'optroot.txt'), 'r') as f:
        for line in f:
            line = line.rstrip('\r\n')
            if line and not line.startswith(' '):
                gtree_str = line
                break

    # create new input file
    f, tmp_fp = mkstemp()
    close(f)
    with open(tmp_fp, 'w') as f:
        f.write('%s\n' % stree_str)
        f.write('%s\n' % gtree_str)

    # run Ranger-DTL for 100 times
    for i in range(reps):
        cmd = ('%s -i %s -o %s --seed %d'
               % (cfg.rangerdtl, tmp_fp, join(result_dir, 'reconc%d' % i), i))
        bms.append(bench(cmd))
    remove(tmp_fp)

    # run AggregateRanger
    cmd = ('%s %s > %s' % (cfg.aggregateranger, join(result_dir, 'reconc'),
                           join(result_dir, 'aggreconc.txt')))
    bms.append(bench(cmd))

    # generate report
    recipients = {}
    # in each reconciliation file, an HGT record line reads like:
    # m4 = LCA[G000008625, G000196115_1]: Transfer, Mapping --> G000008625, Recipient --> G000196115            
    p = re.compile(r'^(.+) = (.+): Transfer, Mapping --> (.+), Recipient --> (.+)$')
    for i in range(reps):
        with open(join(result_dir, 'reconc%d' % i), 'r') as f:
            for line in f:
                line = line.rstrip('\r\n')
                m = p.search(line)
                if m:
                    donor, recipient = m.group(1), m.group(4)
                    if donor not in recipients:
                        recipients[donor] = {recipient: 1}
                    elif recipient not in recipients[donor]:
                        recipients[donor][recipient] = 1
                    else:
                        recipients[donor][recipient] += 1
    top_recipient = {}
    for donor in recipients:
        for recipient in sorted(recipients[donor], key=recipients[donor].get, reverse=True):
            top_recipient[donor] = (recipient, recipients[donor][recipient])
            break
    # cutoffs = {'transfer': 75, 'donor': 90, 'recipient': 90}
    # in the aggregated reconciliation file, an HGT record line reads like:
    # m2 = LCA[G000008625, G000013045]: [Speciations = 35, Duplications = 0, Transfers = 65], [Most Frequent mapping --> G000013045, 48 times].
    dups, hgts = [], []
    p = re.compile(r'^(.+) = (.+): \[Speciations = (\d+), Duplications = (\d+), Transfers = (\d+)\], \[Most Frequent mapping --> (.+), (\d+) times\]')
    with open(join(result_dir, 'aggreconc.txt'), 'r') as f:
        for line in f:
            line = line.rstrip('\r\n')
            m = p.search(line)
            if m:
                donor, n_dup, n_hgt, mapping, n_mapping = m.group(1), int(m.group(4)), int(m.group(5)), m.group(6), int(m.group(7))
                if n_dup > 0:
                    dup_support = (1.0 * n_dup * n_mapping) / (reps ** 2)
                    if dup_support >= cutoff:
                        dups.append((mapping, dup_support))
                if n_hgt > 0:
                    hgt_support = (1.0 * n_hgt * n_mapping * top_recipient[donor][1]) / (reps ** 3)
                    if hgt_support >= cutoff:
                        hgts.append((mapping, top_recipient[donor][0], hgt_support))
    if dups:
        with open(join(result_dir, 'duplications.txt'), 'w') as f:
            for dup in dups:
                # node, support
                f.write('%s\t%.3f\n' % (dup[0], dup[1]))
    if hgts:
        with open(join(result_dir, 'transfers.txt'), 'w') as f:
            for hgt in hgts:
                # donor, recipient, support
                f.write('%s\t%s\t%.3f\n' % (hgt[0], hgt[1], hgt[2]))

    # clean up
    with tarfile.open(join(result_dir, 'reconcs.tar.gz'), 'w:gz') as tar:
        for i in range(reps):
            tar.add(join(result_dir, 'reconc%d' % i), arcname=('reconc%d' % i))
    for i in range(reps):
        remove(join(result_dir, 'reconc%d' % i))
    return bms
def get_protein_to_genome():
    """Generate a protein ID to genome ID dictionary.

    Returns
    -------
    dict of str
        protein ID to genome ID
    """
    prot2g = {}
    for fname in listdir(cfg.proteins_dir):
        if fname.endswith('.faa'):
            gid = fname[:-4]
            with open(join(cfg.proteins_dir, fname)) as f:
                for line in f:
                    if line.startswith('>'):
                        prot = line[1:].split()[0]
                        prot2g[prot] = gid
    return prot2g
if not check_dir('gene_networks'):
    makedirs(join(cfg.intermdir, 'ranger-dtl'))
    prot2g = get_protein_to_genome()

    # reformat species tree
    # replace nodal support values with incremental node indices
    stree = TreeNode.read(cfg.species_tree_fp)
    gids = set()
    idx, nodes, edges = 0, [], []
    for node in stree.levelorder():
        if node.is_tip():  # terminal tip
            gids.add(node.name)
        else:  # internal node
            support = node.name if node.name is not None else ''
            node.name = 'n%s' % idx
            idx += 1
            nodes.append((node.name, support))
        if node.parent is not None:  # non-root
            edges.append((node.parent.name, node.name, node.length))

    # export Cytoscape-compatible network files
    with open(join(cfg.wkdir, 'species_tree.nodes'), 'w') as f:
        f.write('%s\n' % '\t'.join(('name', 'support')))
        for node in nodes:
            f.write('%s\t%s\n' % (node[0], node[1]))
    with open(join(cfg.wkdir, 'species_tree.edges'), 'w') as f:
        f.write('%s\n' % '\t'.join(('source', 'target', 'interaction', 'directed', 'length')))
        for edge in edges:
            f.write('%s\t%s\tV\tTRUE\t%.5f\n' % (edge[0], edge[1], edge[2]))
            
    # Ranger-DTL cannot parse branch lengths as scientific notations
    # change them into zero
    sci = re.compile(r':-?[0-9\.]+e-[0-9]+([,;\)])')
    stree_str = str(stree).replace('\'', '')
    stree_str = re.sub(sci, r':0.0\1', stree_str)
    with open(join(cfg.intermdir, 'ranger-dtl', 'species_tree_labeled.nwk'), 'w') as f:
        f.write(stree_str)

    fid2gid2n = {}
    
    bms = {}
    pool = mp.Pool(cfg.cpus)
    for fname in listdir(cfg.gene_trees_dir):
        if fname.endswith('.nwk'):
            fid = fname[:-4]

            # reformat input file
            gtree = TreeNode.read(join(cfg.gene_trees_dir, fname),
                                  convert_underscores=False)
            gtree.bifurcate()
            for node in gtree.non_tips():
                node.name = ''
            gid2n = {}
            for node in gtree.tips():

                # replace protein ID with host genome ID
                gid = prot2g[node.name]

                # if multiple family members are present in one genome,
                # add suffix "_1", "_2",...
                if gid in gid2n:
                    node.name = '%s_%s' % (gid, gid2n[gid])
                    gid2n[gid] += 1
                else:
                    node.name = gid
                    gid2n[gid] = 1

            fid2gid2n[fid] = gid2n
            gtree_str = str(gtree).replace('\'', '')
            gtree_str = re.sub(sci, r':0.0\1', gtree_str)

            # prune species tree so that only tips present in a gene tree
            # are retained
            sub_stree = stree.copy()
            ntips = len(gid2n)
            while True:
                sub_stree.remove_deleted(lambda x: x.is_tip()
                                         and x.name not in gid2n)
                sub_stree.prune()
                if len([n.name for n in sub_stree.tips()]) == ntips:
                    break
            sub_stree_str = str(sub_stree).replace('\'', '')
            sub_stree_str = re.sub(sci, r':0.0\1', sub_stree_str)

            res_dir = join(cfg.intermdir, 'ranger-dtl', fid)
            makedirs(res_dir)
            input_fp = join(res_dir, 'input.txt')
            with open(input_fp, 'w') as f:
                f.write(sub_stree_str)
                f.write(gtree_str)
            args = (input_fp, res_dir, 100, 0.75, fid)
            bms[fid] = pool.apply_async(run_rangerdtl, args=args)
    pool.close()
    pool.join()
    for fid in bms:
        bms[fid] = bms[fid].get()
        with open(join(cfg.gene_networks_dir, '%s.edges' % fid), 'w') as fo:
            with open(join(cfg.wkdir, 'species_tree.edges'), 'r') as fi:
                line = fi.readline().rstrip('\r\n')
                fo.write('%s\tsupport\n' % line)
                for line in fi:
                    fo.write('%s\t\n' % line.rstrip('\r\n'))
            res_fp = join(cfg.intermdir, 'ranger-dtl', fid, 'transfers.txt')
            if isfile(res_fp):
                with open(res_fp, 'r') as fi:
                    for line in fi:
                        l = line.rstrip('\r\n').split('\t')
                        fo.write('%s\t%s\tH\tTRUE\t\t%s\n' % (l[0], l[1], l[2]))
        with open(join(cfg.gene_networks_dir, '%s.nodes' % fid), 'w') as f:
            f.write('%s\n' % '\t'.join(('name', 'support', 'copies', 'duplications')))
            dups = {}
            res_fp = join(cfg.intermdir, 'ranger-dtl', fid, 'duplications.txt')
            if isfile(res_fp):
                with open(res_fp, 'r') as fi:
                    for line in fi:
                        node = line.rstrip('\r\n').split('\t')[0]
                        if node in dups:
                            dups[node] += 1
                        else:
                            dups[node] = 1
            for gid in gids:
                n_copy = fid2gid2n[fid][gid] if gid in fid2gid2n[fid] else 0
                n_dup = dups[gid] if gid in dups else 0
                f.write('%s\t\t%d\t%d\n' % (gid, n_copy, n_dup))
            for node in nodes:
                n_dup = dups[node[0]] if node[0] in dups else 0
                f.write('%s\t%s\t\t%d\n' % (node[0], node[1], n_dup))
    write_benchmarks(bms, join(cfg.intermdir, 'ranger-dtl', 'benchmark.txt'))<jupyter_output><empty_output><jupyter_text>This step is manual<jupyter_code>taxonomy = pd.read_table(join(cfg.wkdir, 'taxonomy.txt'), index_col='name')
gid2species = {}
for gid in taxonomy.index.tolist():
    gid2species[gid] = taxonomy['species'][gid].replace('_', ' ')
families = {}
with open(join(cfg.wkdir, 'pfam.txt'), 'r') as f:
    for line in f:
        l = line.rstrip('\r\n').split('\t')
        families[l[0]] = l[2]
with open(join(cfg.wkdir, 'species_tree.cyjs'), 'r') as f:
    species_cyjs = json.loads(f.read())
node2suid = {}
for node in species_cyjs['elements']['nodes']:
    node2suid[node['data']['name']] = node['data']['SUID']
used_suids = set()
for e in 'edges', 'nodes':
    for element in species_cyjs['elements'][e]:
        used_suids.add(int(element['data']['SUID']))
start_idx = max(used_suids) + 1
for fid in listdir(join(cfg.intermdir, 'ranger-dtl')):
    if not isdir(join(cfg.intermdir, 'ranger-dtl', fid)):
        continue
    hgts = []
    res_fp = join(cfg.intermdir, 'ranger-dtl', fid, 'transfers.txt')
    if isfile(res_fp):
        with open(res_fp, 'r') as f:
            hgts = [x.split('\t') for x in f.read().splitlines()]
    tips, copies, dups = set(), {}, {}
    with open(join(cfg.gene_networks_dir, '%s.nodes' % fid), 'r') as f:
        next(f)
        for line in f:
            l = line.rstrip('\r\n').split('\t')
            if l[1] == '' and l[2] != '':
                tips.add(l[0])
            if l[2] and l[2] != '0':
                copies[l[0]] = int(l[2])
            if l[3] and l[3] != '0':
                dups[l[0]] = int(l[3])
    gene_cyjs = deepcopy(species_cyjs)
    title = '%s (%s)' % (fid, families[fid]) if fid in families else fid
    gene_cyjs['data']['name'] = title
    gene_cyjs['data']['shared_name'] = title
    for node in gene_cyjs['elements']['nodes']:
        nid = node['data']['name']
        if nid in tips:
            node['data']['species'] = gid2species[nid]
            node['data']['tip'] = 'TRUE'
        if nid in copies:
            node['data']['copies'] = copies[nid]
        if nid in dups:
            node['data']['duplications'] = dups[nid]
    i = start_idx
    for hgt in hgts:
        gene_cyjs['elements']['edges'].append({'data': {
            'SUID': i,
            'directed': True,
            'id': str(i),
            'interaction': 'H',
            'name': '%s (H) %s' % (hgt[0], hgt[1]),
            'selected': False,
            'shared_interaction': 'H',
            'shared_name': '%s (H) %s' % (hgt[0], hgt[1]),
            'source': str(node2suid[hgt[0]]),
            'target': str(node2suid[hgt[1]]),
            'support': float(hgt[2])}, 'selected': False})
        i += 1
    with open(join(cfg.gene_networks_dir, '%s.cyjs' % fid), 'w') as f:
        json.dump(gene_cyjs, f, indent=2)
for fid in listdir(join(cfg.intermdir, 'ranger-dtl')):
    res_dir = join(cfg.intermdir, 'ranger-dtl', fid)
    if isdir(res_dir):
        res_fp = join(res_dir, 'hgts.txt')
        if isfile(res_fp):
            with open(join(cfg.gene_networks_dir, '%s.edges' % fid), 'w') as fo:
                with open(join(cfg.wkdir, 'species_tree.edges'), 'r') as fi:
                    line = fi.readline.rstrip('\r\n')
                    fo.write('%s\tsupport\n' % line)
                    for line in fi:
                        fo.write('%s\t\n' % line.rstrip('\r\n'))
                with open(res_fp, 'r') as fi:
                    for line in fi:
                        l = line.rstrip('\r\n').split('\t')
                        fo.write('%s\t%s\tH\tTRUE\t\t%s\n' % (l[0], l[1], l[2]))
        gtree = TreeNode.read()
        tips = set([x.name for x in tree.tips()])
                with open(join(cfg.gene_networks_dir, '%s.nodes' % fid), 'w') as f:
                    for tip in [x.name for x in stree.tips()]:
                        isin = True if tip in tips else False
                        f.write('%s\t%s\n' % (tip, str(isin)))
def parse_rangerdtl_result(fname):
    with open(fname, 'r') as f:
        for line in f:
            line = line.rstrip('\r\n')
            if line.startswith('Species Tree:'):
                line = next(f).rstrip('\r\n')
                yield line
            elif 'Transfer, Mapping' in line:
                # m4 = LCA[G000008625, G000196115_1]: Transfer, Mapping --> G000008625, Recipient --> G000196115
                l = line.split(': ')[-1].split(', ')
                if l[-2].startswith('Mapping') and l[-1].startswith('Recipient'):
                    donor, recipient = l[-2].split(' --> ')[1], l[-1].split(' --> ')[1]
                    yield donor, recipient
                # m2 = LCA[G000008625, G000013045]: [Speciations = 35, Duplications = 0, Transfers = 65], [Most Frequent mapping --> G000013045, 48 times].
vert_edges = []
for node in stree.postorder(include_self=False):
    vert_edges.append((node.parent.name, node.name))
for fid in listdir(join(cfg.intermdir, 'ranger-dtl')):
    if isdir(join(cfg.intermdir, 'ranger-dtl', fid)):
        gen = parse_rangerdtl_result(join(cfg.intermdir, 'ranger-dtl',
                                          fid, 'aggreconc.txt'))
        tree = TreeNode.read([next(gen)])
        nodes = list(gen)

        hori_edges = []
        for node in nodes:
            d, r = node
            sd, sr = tree.find(d), tree.find(r)
            if sd.is_tip():
                dlca = stree.find(sd.name).name
            else:
                sd_tips = list(sd.subset())
                dlca = stree.lowest_common_ancestor(sd_tips).name
            if sr.is_tip():
                rlca = stree.find(sr.name).name
            else:
                sr_tips = list(sr.subset())
                rlca = stree.lowest_common_ancestor(sr_tips).name
            hori_edges.append((dlca, rlca))

        edges = pd.DataFrame(hori_edges + vert_edges, columns=['source', 'target'])
        edges['hgt'] = len(hori_edges) * [True] + len(vert_edges) * [False]
        edges = edges.set_index('source')
        edges.to_csv(join(cfg.gene_networks_dir, '%s.edges' % fid), sep='\t')

        edges = pd.DataFrame(hori_edges, columns=['source', 'target'])
        edges['interaction'] = len(hori_edges) * ['H']
        edges['directed'] = len(hori_edges) * ['TRUE']
        edges['support'] = len(hori_edges) * ['H']

        tips = set([x.name for x in tree.tips()])
        with open(join(cfg.gene_networks_dir, '%s.nodes' % fid), 'w') as f:
            for tip in [x.name for x in stree.tips()]:
                isin = True if tip in tips else False
                f.write('%s\t%s\n' % (tip, str(isin)))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13533129851714773, 'fraction_numerical': 0.011079333797976399, 'mean_word_length': 1.9529486906536087, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 55, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '596508', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 14189, 'n_tokens_neox': 13469, 'n_words': 3570, 'path': 'code/notebooks/ocWoL.ipynb', 'repo_name': 'YimengYang/wol', 'size': '41612'}