partis-bcr 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

bin/gcdyn-simu-run.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+from __future__ import absolute_import, division, unicode_literals
+from __future__ import print_function
+import glob
+import sys
+import csv
+from io import open
+csv.field_size_limit(sys.maxsize)  # make sure we can write very large csv fields
+import os
+import copy
+import argparse
+import colored_traceback.always
+import json
+
+# if you move this script, you'll need to change this method of getting the imports
+from pathlib import Path
+partis_dir = str(Path(__file__).parent.parent)
+sys.path.insert(1, partis_dir) # + '/python')
+
+import python.utils as utils
+import python.paircluster as paircluster
+import python.glutils as glutils
+from python.clusterpath import ClusterPath
+import python.treeutils as treeutils
+import python.indelutils as indelutils
+from python.event import RecombinationEvent
+
+# ----------------------------------------------------------------------------------------
+def get_replay_naive_antn(glfos, ltmp, add_empty_mature_keys=False, debug=False):
+    antn, naive_seq = {}, []
+    for region in utils.getregions(ltmp):
+        gene, seq = utils.get_single_entry(list(glfos[ltmp]['seqs'][region].items()))
+        antn['%s_gene'%region] = gene
+        if ltmp == 'igh' and region == 'j' and args.rm_last_ighj_base:
+            if debug:
+                print('    removing last ighj base: %s --> %s' % (seq, seq[:-1]))
+            seq = seq[:-1]
+        naive_seq.append(seq)
+    for dstr in utils.all_erosions:
+        antn['%s_del'%dstr] = 0
+    if not utils.has_d_gene(ltmp):
+        antn['d_gene'] = glutils.dummy_d_genes[ltmp]
+        antn['d_5p_del'] = 1
+    for bstr in utils.all_boundaries:
+        antn['%s_insertion'%bstr] = ''
+    antn['naive_seq'] = ''.join(naive_seq)
+    antn['invalid'] = False
+    if add_empty_mature_keys:  # keys that need to be replaced in mature annotation
+        antn['unique_ids'] = []
+        antn['seqs'] = []
+        antn['input_seqs'] = antn['seqs']
+        antn['indelfos'] = []
+        antn['paired-uids'] = []
+    else:
+        antn['unique_ids'] = ['replay_naive']
+        antn['seqs'] = [''.join(naive_seq)]
+        antn['input_seqs'] = antn['seqs']
+        antn['indelfos'] = [indelutils.get_empty_indel()]
+        antn['paired-uids'] = [[]]
+    if debug:
+        utils.add_implicit_info(glfos[ltmp], antn)
+        utils.print_reco_event(antn)
+    return antn
+
+# ----------------------------------------------------------------------------------------
+def get_uid(name, ltmp):
+    return '%s-%s' % (name, ltmp)
+
+# ----------------------------------------------------------------------------------------
+def process_tree(glfos, treelist, tree_sfos, leaf_meta, itree, lp_infos, lpair):
+    antns = {}
+    for ltmp in lpair:
+        antns[ltmp] = get_replay_naive_antn(glfos, ltmp, add_empty_mature_keys=True, debug=args.debug)
+        for tkey in args.meta_columns:
+            antns[ltmp][utils.input_metafile_keys[tkey]] = []
+
+    hloc = utils.heavy_locus(args.ig_or_tr)
+    lloc = utils.get_single_entry([l for l in glfos if l!=hloc])
+    h_seq_len = len(antns[hloc]['naive_seq'])
+    for sfo in tree_sfos[itree]:
+        joint_naive_seq = ''.join(antns[l]['naive_seq'] for l in lpair)
+        if len(sfo['seq']) == len(joint_naive_seq) - 1:
+            raise Exception('seq read from gcdyn file has len %d, one less than replay naive seq %d, you probably need to set --rm-last-ighj-base' % (len(sfo['seq']), len(joint_naive_seq)))
+        assert len(sfo['seq']) == len(joint_naive_seq)
+        if 'naive' in sfo['name']:
+            assert sfo['seq'] == joint_naive_seq
+            continue
+        dumb_offset = 0
+        if args.rm_last_ighj_base:
+            dumb_offset = 1
+        antns[hloc]['seqs'].append(sfo['seq'][:h_seq_len + dumb_offset])
+        antns[lloc]['seqs'].append(sfo['seq'][h_seq_len:])
+    
+        for ltmp in lpair:
+            assert antns[ltmp]['input_seqs'][-1] == antns[ltmp]['seqs'][-1]
+            antns[ltmp]['unique_ids'].append(get_uid(sfo['name'], ltmp))
+            antns[ltmp]['indelfos'].append(indelutils.get_empty_indel())
+            other_locus = utils.get_single_entry([l for l in lpair if l!=ltmp])
+            antns[ltmp]['paired-uids'].append([get_uid(sfo['name'], other_locus)])
+            for tkey in args.meta_columns:
+                antns[ltmp][utils.input_metafile_keys[tkey]].append(leaf_meta[sfo['name']][tkey])
+
+    for ltmp in lpair:
+        dtree = treeutils.get_dendro_tree(treestr=treelist[itree])
+        treeutils.translate_labels(dtree, [(s['name'], get_uid(s['name'], ltmp)) for s in tree_sfos[itree] if 'naive' not in s['name']], expect_missing=True)
+        antns[ltmp]['tree'] = dtree.as_string(schema='newick').strip()
+        tmp_event = RecombinationEvent(glfos[ltmp])  # I don't want to move the function out of event.py right now
+        tmp_event.set_reco_id(antns[ltmp], irandom=itree)  # not sure that setting <irandom> here actually does anything
+        utils.add_implicit_info(glfos[ltmp], antns[ltmp])  # easiest way to add codon_positions, which we want to write to file
+
+    for ltmp in lpair:
+        lp_infos[lpair]['antn_lists'][ltmp].append(antns[ltmp])
+
+    if args.debug:
+        for ltmp in sorted(glfos):
+            utils.print_reco_event(antns[ltmp], extra_str='        ')
+    
+# ----------------------------------------------------------------------------------------
+def mfname():  # have to look for old meta file name for backwards compatibility
+    mfn = '%s/meta.csv'%gcd_dir
+    if os.path.exists(mfn):  # if new name exists, return that
+        return mfn
+    return '%s/leaf-meta.csv'%gcd_dir  # otherwise return old name
+
+# ----------------------------------------------------------------------------------------
+def run_gcdyn():
+    if os.path.exists('%s/encoded-trees.npy'%gcd_dir):
+        print('    gcdyn output exists in %s' % gcd_dir)
+        return
+    cmds = ['#!/bin/bash']
+    cmds += utils.mamba_cmds('gcdyn')
+    gcmd = 'gcd-simulate --sample-internal-nodes --label-leaf-internal-nodes --outdir %s --xshift-values 2.5 --xscale-values 5 --yscale-values 1000000' % gcd_dir
+    if args.n_sub_procs is not None:
+        gcmd += ' --n-sub-procs %d' % args.n_sub_procs
+    if args.seed is not None:
+        gcmd += ' --seed %d' % args.seed
+    if args.n_sim_events is not None:
+        gcmd += ' --n-trials %d' % args.n_sim_events
+    if args.obs_times is not None:
+        gcmd += ' --time-to-sampling-values %d' % args.obs_times
+    cmds += [gcmd]
+    cmdfname = '%s/run.sh' % gcd_dir
+    utils.simplerun('\n'.join(cmds) + '\n', cmdfname=cmdfname)
+
+# ----------------------------------------------------------------------------------------
+def process_output():
+    glfos = {}
+    for ltmp in utils.loci:
+        if os.path.exists('%s/%s' % (args.replay_germline_dir, ltmp)):
+            glfos[ltmp] = glutils.read_glfo(args.replay_germline_dir, ltmp)
+
+    seqfos = utils.read_fastx('%s/seqs.fasta'%gcd_dir)
+    treelist = treeutils.get_treestrs_from_file('%s/trees.nwk'%gcd_dir)
+    lmetalines = utils.csvlines(mfname())
+    leaf_meta = {l['name'] : {'affinity' : float(l['affinity'])} for l in lmetalines}
+    print('  read %d trees, %d seqs  (plus leaf metafo) from %s' % (len(treelist), len(seqfos), gcd_dir))
+    tree_sfos, all_uids = {}, set()  # collect up the seqfos for each tree
+    for sfo in seqfos:
+        if sfo['name'].count('-') == 1:  # old-style, and naive seq
+            itree, sname = sfo['name'].split('-')
+        elif sfo['name'].count('-') == 2:  # new-style
+            itree, listr, sname = sfo['name'].split('-')  # listr is either 'leaf' or 'mrca'
+        else:
+            assert False
+        itree = int(itree)
+        if itree not in tree_sfos:
+            tree_sfos[itree] = []
+        tree_sfos[itree].append(sfo)
+        if sfo['name'] in all_uids:
+            raise Exception('found uid %s twice' % sfo['name'])
+        all_uids.add(sfo['name'])
+    print('      %d tree seqfos with lengths: %s' % (len(tree_sfos), ' '.join(str(len(slist)) for slist in sorted(tree_sfos.values(), key=len, reverse=True))))
+    if sorted(tree_sfos) != list(range(len(treelist))):
+        raise Exception('tree indices from sequence names didn\'t match number of trees')
+
+    lpairs = [tuple(sorted(glfos))]  # could use utils.locus_pairs() if i update to more than one lpair
+    lp_infos = {lp : {'antn_lists' : {l : [] for l in glfos}, 'glfos' : {l : g for l, g in glfos.items()}, 'cpaths' : {}} for lp in lpairs}
+    for itree in range(len(treelist)):
+        process_tree(glfos, treelist, tree_sfos, leaf_meta, itree, lp_infos, utils.get_single_entry(lpairs))
+
+    print('  writing annotations to %s' % args.outdir)
+    headers = utils.simulation_headers + args.meta_columns
+    def ofn_fn(locus, lpair=None, joint=None):
+        return paircluster.paired_fn(args.outdir, locus, lpair=lpair, suffix='.yaml')
+    paircluster.write_lpair_output_files(lpairs, lp_infos, ofn_fn, headers=headers)
+    glfos, antn_lists, _ = paircluster.concat_heavy_chain(lpairs, lp_infos)  # per-locus glfos with concat'd heavy chain
+    paircluster.write_concatd_output_files(glfos, antn_lists, ofn_fn, headers)
+    outfos, metafos = paircluster.get_combined_outmetafos(antn_lists, extra_meta_headers=[utils.input_metafile_keys[k] for k in args.meta_columns])
+    paircluster.write_combined_fasta_and_meta(args.outdir+'/all-seqs.fa', args.outdir+'/meta.yaml', outfos, metafos)
+
+# ----------------------------------------------------------------------------------------
+helpstr = """
+run gcdyn simulation, then process results into partis-format paired output dir, for example:
+    ./bin/gcdyn-simu-run.py --gcd-dir <gcd> --n-sub-procs 10 --seed 0 --n-trials 10 --rm-last-ighj-base
+"""
+class MultiplyInheritedFormatter(argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter):
+    pass
+formatter_class = MultiplyInheritedFormatter
+parser = argparse.ArgumentParser(formatter_class=MultiplyInheritedFormatter, description=helpstr)
+parser.add_argument('--outdir', required=True)
+parser.add_argument('--actions', default='run:process')
+parser.add_argument('--input-simu-dir', help='if set, only run \'process\' action, reading gcdyn simu from this dir and writing partis files to --outdir')
+parser.add_argument('--replay-germline-dir', default='datascripts/meta/taraki-gctree-2021-10/germlines', help='dir with gcreplay germline sequences')
+parser.add_argument('--rm-last-ighj-base', action='store_true', help='sometimes the ighj gene has an extra G at the end, sometimes not, this says to remove it from the seqs read from --replay-germline-dir')
+parser.add_argument('--n-sub-procs', type=int)
+parser.add_argument('--seed', type=int)
+parser.add_argument('--n-sim-events', type=int)
+parser.add_argument('--obs-times', type=int)
+parser.add_argument('--meta-columns', default='affinity')
+parser.add_argument('--ig-or-tr', default='ig')
+parser.add_argument('--debug', action='store_true')
+args = parser.parse_args()
+args.meta_columns = utils.get_arg_list(args.meta_columns, choices=utils.input_metafile_keys.keys())
+if args.input_simu_dir is None:
+    gcd_dir = '%s/gcdyn' % args.outdir
+    args.actions = utils.get_arg_list(args.actions)
+else:
+    gcd_dir = args.input_simu_dir
+    args.actions = ['process']
+if 'run' in args.actions:
+    run_gcdyn()
+if 'process' in args.actions:
+    process_output()

bin/gctree-run.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python3
+from __future__ import absolute_import, division, unicode_literals
+from __future__ import print_function
+import numpy
+import csv
+import yaml
+import time
+import colored_traceback.always
+import argparse
+import subprocess
+import sys
+import os
+import dendropy
+import json
+from io import open
+import random
+from pathlib import Path
+
+partis_dir = str(Path(__file__).parent.parent)
+sys.path.insert(1, partis_dir) #'./python')
+import python.utils as utils
+import python.glutils as glutils
+import python.treeutils as treeutils
+
+# ----------------------------------------------------------------------------------------
+def get_inf_int_name(gname):  # <gname> is just an integer, which won't be unique and will break things
+    return '%s-%s' % (args.inf_int_label, gname)
+
+# ----------------------------------------------------------------------------------------
+def gctofn(ft):
+    ftstrs = {
+        'tree' : 'gctree.out.inference.1.nk',
+        'seqs' : 'gctree.out.inference.1.fasta',
+        'dnapars' : 'outfile',
+    }
+    return '%s/%s' % (args.outdir, ftstrs[ft])
+
+# ----------------------------------------------------------------------------------------
+def fofn(ft):
+    assert ft in ['tree', 'seqs']
+    return '%s/%s%s' % (args.outdir, ft if ft=='tree' else 'inferred-%s'%ft, '.nwk' if ft=='tree' else '.fa')
+
+# ----------------------------------------------------------------------------------------
+def idfn():
+    return 'idmap.txt'
+
+# ----------------------------------------------------------------------------------------
+def install():
+    cmds = ['#!/bin/bash']
+    cmds += utils.mamba_cmds(args.env_label, only_prep=True)
+    cmds += ['micromamba create -y -n %s -c conda-forge python=3.9' % args.env_label]  # 3.10 currently has problems with ete
+    cmds += ['micromamba activate %s' % args.env_label]
+    cmds += ['micromamba install -y -c bioconda -c conda-forge phylip']
+    cmds += ['micromamba install -y -c conda-forge%s click' % ('' if args.no_dag else ' gctree')]
+    if args.no_dag:
+        cmds += ['pip install gctree==3.3.0']  # I think having --user makes it install in ~/.local (outside mamba env)
+    # micromamba remove -n gctree --all  # to nuke it and start over
+    utils.simplerun('\n'.join(cmds) + '\n', cmdfname='/tmp/tmprun.sh', debug=True)
+
+# ----------------------------------------------------------------------------------------
+def update():
+    cmds = ['#!/bin/bash']
+    cmds += utils.mamba_cmds(args.env_label)
+    cmds += ['micromamba update phylip gctree click']
+    utils.simplerun('\n'.join(cmds) + '\n', cmdfname='/tmp/tmprun.sh', debug=True)
+
+# ----------------------------------------------------------------------------------------
+def add_mfo(tcmd, mfn):
+    kdict = {'frame' : 'frame', 'h_frame' : 'frame', 'l_frame' : 'frame2', 'l_offset' : 'chain_split'}  # translates from metafo dict to gctree command line args
+    with open(args.metafname) as mfile:
+        metafo = json.load(mfile)
+    for tk, tc in kdict.items():
+        if tk in metafo:
+            tcmd += ' --%s %d' % (tc, metafo[tk])
+    return tcmd
+
+# ----------------------------------------------------------------------------------------
+def run_gctree():
+    # ----------------------------------------------------------------------------------------
+    def get_gctree_cmd():
+        tcmd = '%s/bin/xvfb-run -a gctree infer outfile abundances.csv --root %s --verbose --idlabel' % (utils.get_partis_dir(), args.root_label)  # --idlabel writes the output fasta file
+        if not args.base_model and not args.no_dag:
+            tcmd += ' --mutability %s/HS5F_Mutability.csv --substitution %s/HS5F_Substitution.csv' % (args.data_dir, args.data_dir)
+        if args.ranking_coeffs is not None:
+            tcmd += ' --ranking_coeffs %s' % (' '.join(c for c in args.ranking_coeffs))
+        if args.branching_process_ranking_coeff is not None:
+            tcmd += ' --branching_process_ranking_coeff %d' % args.branching_process_ranking_coeff
+        if os.path.exists(args.metafname):
+            tcmd = add_mfo(tcmd, args.metafname)
+        return tcmd
+    # ----------------------------------------------------------------------------------------
+    def get_cmds():
+        cmds = ['#!/bin/bash']
+        cmds += utils.mamba_cmds(args.env_label)
+        if args.run_help:
+            cmds += ['gctree infer -h']
+            return cmds
+        if not os.path.exists(args.infname):
+            raise Exception('--infname %s doesn\'t exist' % args.infname)
+        cmds += ['cd %s' % args.outdir]
+        if args.input_forest_dir is None:
+            ofn = '%s/outfile' % args.outdir  # dnapars output file (this is what takes the longest to make
+            if os.path.exists(ofn) and os.stat(ofn).st_size > 0:
+                print('    dnapars output already exists, not rerunning: %s' % ofn)
+            else:
+                if os.path.exists(ofn) and os.stat(ofn).st_size == 0:
+                    print('    removing zero length dnapars output %s' % ofn)
+                    utils.prep_dir(args.outdir, wildlings=['outfile', 'outtree'], allow_other_files=True)  # phylip barfs like a mfer if its outputs exist (probably you'll get a KeyError 'naive')
+                cmds += ['deduplicate %s --root %s --abundance_file abundances.csv --idmapfile %s > deduplicated.phylip' % (args.infname, args.root_label, idfn())]
+                cmds += ['mkconfig deduplicated.phylip dnapars > dnapars.cfg']
+                cmds += ['dnapars < dnapars.cfg > dnapars.log']  # NOTE if things fail, look in dnaparse.log (but it's super verbose so we can't print it to std out by default)
+        else:
+            print('    --input-forest-dir: copying abundance, idmap, and forest files from %s' % args.input_forest_dir)
+            cmds += ['cp %s/{abundances.csv,%s,outfile} %s/' % (args.input_forest_dir, idfn(), args.outdir)]
+        if not args.only_write_forest:
+            cmds.append(get_gctree_cmd())
+        return cmds
+    # ----------------------------------------------------------------------------------------
+    if not args.run_help and utils.output_exists(args, gctofn('dnapars' if args.only_write_forest else 'tree')):
+        return
+
+    cmds = get_cmds()  # also preps dir + other stuff
+
+    utils.simplerun('\n'.join(cmds) + '\n', cmdfname=args.outdir + '/run.sh', print_time='gctree', debug=True, dryrun=args.dry_run)
+    if args.run_help:
+        sys.exit()
+
+# ----------------------------------------------------------------------------------------
+def parse_output():
+    if utils.output_exists(args, fofn('seqs')):
+        return
+
+    # read translations (this only includes input sequences, not inferred intermediates)
+    idm_trns = {}
+    with open('%s/idmap.txt' % args.outdir) as idfile:
+        reader = csv.DictReader(idfile, fieldnames=('name', 'orig_names'))
+        for line in reader:
+            if line['orig_names'] == '':
+                continue
+            idm_trns[line['name']] = line['orig_names'].split(':')
+
+    # read fasta (mostly for inferred intermediate seqs)
+    seqfos = utils.read_fastx(gctofn('seqs'), look_for_tuples=True)
+    print('    read %d seqs from gctree output fasta' % len(seqfos))
+    if any(s['name']=='' for s in seqfos):
+        n_removed = len([s for s in seqfos if s['name']==''])
+        seqfos = [s for s in seqfos if s['name']!='']
+        print('  %s removed %d seqs with zero-length names \'\' (I\'m *not* sure this is the right thing to do, but it just kicked this error when I was doing the python 3 conversion)' % (utils.wrnstr(), n_removed))
+    nfos = [s for s in seqfos if s['name']==args.root_label]
+    if len(nfos) != 1:
+        print('  %s expected 1 naive seq with label \'%s\' but found %d: %s  (in %s)' % (utils.wrnstr(), args.root_label, len(nfos), ' '.join(n['name'] for n in nfos), gctofn('seqs')))
+    seqfos = [s for s in seqfos if s['name'] != args.root_label]  # don't want naive seq in final fasta
+    seq_len = numpy.mean([len(s['seq']) for s in seqfos])
+    if not args.expand_all_nodes:  # also remove input seqs (well, gctree's new names for input seqs), unless we're expanding all nodes, in which case we need the gctree-named-nodes as fake new internal nodes
+        seqfos = [s for s in seqfos if s['name'] not in idm_trns]
+    if len(seqfos) == 0:
+        print('  %s no inferred sequences (all seqs read from gctree output were input seqs' % utils.wrnstr())
+    inf_int_trns = []
+    for sfo in seqfos:
+        inf_int_trns.append((sfo['name'], get_inf_int_name(sfo['name'])))
+        sfo['name'] = get_inf_int_name(sfo['name'])
+
+    # read tree
+    dtree = treeutils.get_dendro_tree(treefname=gctofn('tree'), debug=args.debug)
+    dtree.scale_edges(1. / seq_len)
+    dtree.seed_node.taxon.label = args.root_label
+    ndict = {n.taxon.label : n for n in dtree.preorder_node_iter()}
+    for gname, onames in idm_trns.items():
+        node = ndict[gname]
+        if node is None:
+            raise Exception('couldn\'t find node with name \'%s\' in tree from gctree in %s' % (gname, gctofn('tree')))
+        if args.debug and len(onames) > 1:
+            print('    abundance > 1 for %s: %d (%s)' % (gname, len(onames), ' '.join(onames)))
+        for onm in onames:
+            if node.taxon.label == gname and not args.expand_all_nodes:
+                node.taxon.label = onm
+                if args.debug and len(onames) > 1:
+                    print('        setting node to %s' % onm)
+                continue
+            treeutils.add_zero_length_child(node, dtree, child_name=onm)  # add duplicates as children with zero-length edges
+            if args.debug and len(onames) > 1:
+                print('        adding child node %s' % onm)
+    treeutils.translate_labels(dtree, inf_int_trns, expect_missing=True, debug=args.debug)
+
+    if args.fix_multifurcations:
+        input_seqfos = utils.read_fastx(args.infname)
+        dtree, new_seqfos = treeutils.get_binary_tree(dtree, nfos + input_seqfos + seqfos, debug=args.debug)
+        seqfos += new_seqfos
+    if args.debug:
+        print('    final tree:')
+        print(treeutils.get_ascii_tree(dendro_tree=dtree, extra_str='      ', width=350))
+    with open(fofn('tree'), 'w') as ofile:
+        ofile.write('%s\n' % treeutils.as_str(dtree))
+    utils.write_fasta(fofn('seqs'), nfos + seqfos)
+
+# ----------------------------------------------------------------------------------------
+ustr = """
+Run gctree tree inference on sequences from fasta input file <--infname>.
+Output trees and sequences are written to <--outdir> as inferred-seqs.fa and tree.nwk (gctree output files are also there, but they don't have any postprocessing e.g. fixing names and/or multifurcations.
+  gctree-run.py --infname <fasta> --outdir <outdir>
+"""
+parser = argparse.ArgumentParser(usage=ustr)
+parser.add_argument('--actions', default='run:parse')
+parser.add_argument('--infname')
+parser.add_argument('--metafname', help='if you need --frame (v region doesn\'t start at first position) or --chain_split and --frame2 (heavy/light chain smooshed together), pass the info in json format with this arg (see code above for format).')
+parser.add_argument('--outdir')
+parser.add_argument('--only-write-forest', action='store_true', help='only run preparatory steps for gctree, i.e. up through dnapars, to write parsimony forest')
+parser.add_argument('--input-forest-dir', help='If set, skips preparatory steps (see --only-write-forest), and looks for \'abundance.csv\' and parsimony forest file (\'outfile\') in the specified dir')
+parser.add_argument('--overwrite', action='store_true')
+parser.add_argument('--base-model', action='store_true', help='By default, we pass gctree info for the s5f mutation model; if this is set, we don\'t, and it instead use the base model.')
+parser.add_argument('--no-dag', action='store_true', help='If set, use old v1 non-DAG gctree version (v3.3.0). Note that this uses a different env (see --env-label)')
+parser.add_argument('--ranking-coeffs', nargs='+', help='see gctree help')
+parser.add_argument('--branching-process-ranking-coeff', type=int, help='see gctree help')
+parser.add_argument('--env-label', default='gctree')
+parser.add_argument('--root-label', default='naive')
+parser.add_argument('--data-dir', default='%s/data/s5f'%utils.get_partis_dir())
+parser.add_argument('--inf-int-label', default='inf', help='base name for inferred intermediate seqs (numerical name is appended with -')
+parser.add_argument('--expand-all-nodes', action='store_true', help='Gctree collapses duplicate observed seqs into nodes with new names and abundance N > 1. By default, we expand these such that the node is named for one of the observed seqs, and add N-1 (zero-length) children. If this arg is set, however, we leave the node and add N (zero-length) children.')
+parser.add_argument('--run-help', action='store_true', help='run gctree help')
+parser.add_argument('--debug', action='store_true')
+parser.add_argument('--dry-run', action='store_true')
+parser.add_argument('--random-seed', type=int, default=0)
+parser.add_argument('--fix-multifurcations', action='store_true', help='resolves multifurcations (by adding zero length intermediates) and move input seqs that have been extend unifurcations onto zero length branches')
+
+args = parser.parse_args()
+random.seed(args.random_seed)
+numpy.random.seed(args.random_seed)
+if args.only_write_forest and args.input_forest_dir:
+    raise Exception('doesn\'t make sense to specify both')
+args.actions = utils.get_arg_list(args.actions, choices=['install', 'update', 'run', 'parse'])
+args.infname = utils.fpath(args.infname)
+args.outdir = utils.fpath(args.outdir)
+if args.no_dag:
+    assert not args.base_model and args.branching_process_ranking_coeff is None and args.ranking_coeffs is None
+    args.env_label = 'gctree-no-dag'
+
+if 'install' in args.actions:
+    install()
+if 'update' in args.actions:
+    update()
+if 'run' in args.actions:
+    run_gctree()
+if 'parse' in args.actions:
+    parse_output()

bin/get-naive-probabilities.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+from __future__ import absolute_import, division, unicode_literals
+from __future__ import print_function
+import os
+import sys
+import csv
+import argparse
+import operator
+import argparse
+import yaml
+import colored_traceback.always
+from io import open
+
+# if you move this script, you'll need to change this method of getting the imports
+from pathlib import Path
+partis_dir = str(Path(__file__).parent.parent)
+sys.path.insert(1, partis_dir) # + '/python')
+
+import python.utils as utils
+
+# ----------------------------------------------------------------------------------------
+def is_acceptable(scol, acceptable_values, lval):
+    if lval in acceptable_values:
+        return True
+    if args.any_allele and '_gene' in scol and any(utils.are_alleles(g, lval) for g in acceptable_values):
+        return True
+    return False
+
+class MultiplyInheritedFormatter(argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter):
+    pass
+formatter_class = MultiplyInheritedFormatter
+parser = argparse.ArgumentParser(formatter_class=MultiplyInheritedFormatter)
+parser.add_argument('--infname', default='test/ref-results/test/parameters/data/hmm/all-probs.csv', help='input all-probs.csv file from a previously-inferred partis parameter directory, for instance: test/reference-results/test/parameters/data/hmm/all-probs.csv')
+parser.add_argument('--config-fname', help='yaml file with info on columns for which we want to specify particular values (and skip others). See default/example set below. To create a yaml config file to start from, uncomment the yaml.dump() line below and rerun with no arguments.')
+parser.add_argument('--outfname')
+parser.add_argument('--any-allele', action='store_true', help='if set, also include any other alleles of any of the genes specified in \'skip_column_vals\' (note: can also set it in the cfg file).')
+parser.add_argument('--debug', action='store_true', default=True)  # it's kind of confusing without the debug printout
+args = parser.parse_args()
+
+non_summed_column = None
+if args.config_fname is None:
+    non_summed_column = 'v_gene'
+    skip_column_vals = {  # to input your own dict on the command line, just convert with str() and quote it
+        # 'cdr3_length' : ['33', '36', '39', '42', '45', '48'],  # <value> is list of acceptable values NOTE need to all be strings, otherwise you have to worry about converting the values in the csv file
+        'v_gene' : ['IGHV1-2*02+G35A', 'IGHV1-2*02+T147C', 'IGHV1-2*02'],
+        # 'd_gene' : ['IGHD3-22*01'],
+        'j_gene' : ['IGHJ4*02'],
+        'cdr3_length' : ['66',],
+    }
+    print('%s using default skip column/non-summed column values (which probably don\'t correspond to what you\'re actually interested in)' % utils.color('red', 'note'))
+    # # uncomment to create a yaml file to start from:
+    # with open('tmp.yaml', 'w') as tfile:
+    #     yaml.dump({'non_summed_column' : non_summed_column, 'skip_column_vals' : skip_column_vals}, tfile)
+else:
+    with open(args.config_fname) as yamlfile:
+        yamlfo = yaml.load(yamlfile, Loader=yaml.Loader)
+    if 'non_summed_column' in yamlfo:
+        non_summed_column = yamlfo['non_summed_column']
+    skip_column_vals = yamlfo['skip_column_vals']
+    for scol in skip_column_vals:
+        skip_column_vals[scol] = [str(v) for v in skip_column_vals[scol]]  # yaml.load() converts to integers, which is usually nice, but here we don't want it to since we're not converting when reading all-probs.csv (I think there's options to yaml.load to change this, I just don't want to figure it out now)
+    if 'any_allele' in yamlfo:
+        if args.any_allele and not yamlfo['any_allele']:  # if it's set to true on the command line, but false in the file
+            print(' %s overwriting --any-allele with value from cfg file %s' % (utils.color('red', 'warning'), args.config_fname))
+        args.any_allele = yamlfo['any_allele']
+
+info = {}
+lines_skipped, lines_used = 0, 0
+counts_skipped, counts_used = 0, 0
+print('  reading probs from %s' % args.infname)
+with open(args.infname) as csvfile:
+    reader = csv.DictReader(csvfile)
+    # if args.debug:
+    #     print '  all columns in file: %s' % ' '.join(reader.fieldnames)
+    if len(set(skip_column_vals) - set(reader.fieldnames)) > 0:
+        raise Exception('keys in --skip-column-fname not in file: %s' % ' '.join(set(skip_column_vals) - set(reader.fieldnames)))
+    for line in reader:
+        skip_this_line = False
+        for scol, acceptable_values in skip_column_vals.items():
+            if not is_acceptable(scol, acceptable_values, line[scol]):
+                skip_this_line = True
+                lines_skipped += 1
+                counts_skipped += int(line['count'])
+                break
+        if skip_this_line:
+            continue
+
+        if non_summed_column is not None:
+            if line[non_summed_column] not in info:
+                info[line[non_summed_column]] = 0
+            info[line[non_summed_column]] += int(line['count'])
+
+        lines_used += 1
+        counts_used += int(line['count'])
+
+# ----------------------------------------------------------------------------------------
+import python.fraction_uncertainty as fraction_uncertainty
+def frac_err(obs, total):
+    lo, hi = fraction_uncertainty.err(obs, total)
+    return 0.5 * (hi - lo)
+count_fraction = counts_used / float(counts_used + counts_skipped)
+
+if args.debug:
+    print('  applied restrictions:%s' % ('     (including all alleles of these genes)' if args.any_allele else ''))
+    for scol, acceptable_values in skip_column_vals.items():
+        print('      %15s in %s' % (scol, acceptable_values))
+    print('   used:')
+    print('     %6d / %-6d = %.3f  lines'  % (lines_used, lines_used + lines_skipped, lines_used / float(lines_used + lines_skipped)))
+    print('     %6d / %-6d = %.3f +/- %.3f counts'  % (counts_used, counts_used + counts_skipped, count_fraction, frac_err(counts_used, counts_used + counts_skipped)))
+
+    if non_summed_column is not None:
+        print('    %18s      count      / %d = fraction' % (non_summed_column, counts_used))
+        for val, count in sorted(list(info.items()), key=operator.itemgetter(1), reverse=True):  # sort by counts
+        # for val, count in sorted(info.items()):  # sort by column value (e.g. cdr3 length)
+            print('   %18s   %6d          %.3f +/- %.3f' % (val, count, count / float(counts_used), frac_err(count, counts_used)))
+
+if args.outfname is not None:
+    if args.debug:
+        print('  writing total counts (plus %d info entries) to %s' % (len(info), args.outfname))
+    with open(args.outfname, 'w') as outfile:
+        yamlfo = {'counts' : counts_used,
+                  'total' : counts_used + counts_skipped,
+                  'fraction' : count_fraction,
+                  'frac_err' : frac_err(counts_used, counts_used + counts_skipped),
+                  'info' : info}
+        yaml.dump(yamlfo, outfile, width=150)