diff --git a/inStrain/argumentParser.py b/inStrain/argumentParser.py index 4ca4366..c86244f 100644 --- a/inStrain/argumentParser.py +++ b/inStrain/argumentParser.py @@ -183,13 +183,29 @@ def parse_args(args): compare_parent = argparse.ArgumentParser(add_help=False) # Required positional arguments Rflags = compare_parent.add_argument_group('REQUIRED') - Rflags.add_argument('-i', '--input', help="A list of inStrain objects, all mapped to the same .fasta file", + Rflags.add_argument('-i', '--input', + help="A list of inStrain objects, all mapped to the same .fasta file", nargs='*', required=True) Rflags.add_argument("-o", "--output", action="store", default='instrainComparer', \ help='Output prefix') + Gflags = compare_parent.add_argument_group('GROUPS') + Gflags.add_argument('--list-groups', help="List all groups", action='store_true', + default=False, required=False) + Gflags.add_argument('--group-pkl', help="Pickled group object", + default=None, type=str, required=False) + Gflags.add_argument('--group', help="Specific group to compare", + default=None, type=int, required=False) + Gflags.add_argument('--comparisons', + help="All comparison pkl files (if previously used --group)", + default=None, required=False, nargs='+') + Gflags.add_argument('--comparisons-list', + help="A file list all pcomparison pkl files (if previously used --group). One file per line", + default=None, required=False) compare_parser = subparsers.add_parser("compare",formatter_class=SmartFormatter,\ - parents = [compare_parent, parent_parser, geneomewide_parent, variant_parent], add_help=False) + parents = [compare_parent, parent_parser, + geneomewide_parent, variant_parent], + add_help=False) # Database mode parameters Dflags = compare_parser.add_argument_group('DATABASE MODE PARAMETERS') @@ -255,7 +271,9 @@ def parse_args(args): ''' # Make a parent for profile to go above the system arguments genome_parser = subparsers.add_parser("genome_wide",formatter_class=SmartFormatter,\ - parents = [geneomewide_parent, genes_io, mm_parent, parent_parser], add_help=False) + parents = [geneomewide_parent, genes_io, + mm_parent, parent_parser], + add_help=False) ''' ####### Arguments for plot operation ###### diff --git a/inStrain/compare_controller.py b/inStrain/compare_controller.py index abbc4b4..6e6ed32 100644 --- a/inStrain/compare_controller.py +++ b/inStrain/compare_controller.py @@ -4,8 +4,10 @@ # Import packages import os +import sys import copy import logging +import pickle import pandas as pd from tqdm import tqdm import multiprocessing @@ -59,7 +61,7 @@ def main(self): *************************************************** """ logging.info(message) - self.run_comparisons() + self.gather_comparisons() # Do auxillary processing if needed message = """\ @@ -180,19 +182,58 @@ def run_comparisons(self): cdbs = [] # Store scaffold comparison information mdbs = [] # Store mismatch locations pair2mm2covOverlaps = [] # Store coverage overlap locations - order = [] # Store order of scaffolds - - groups = len(self.scaffold_comparison_groups) - for i, SCgroup in enumerate(self.scaffold_comparison_groups): - logging.info(f'Running group {i+1} of {groups}') + order = [] # Store order of scaffolds + # groups + groups = self.scaffold_comparison_groups + # listing roups + if self.args.list_groups: + print('--- groups ---') + for i,g in enumerate(self.scaffold_comparison_groups): + print('Group {}: '.format(i+1), end='') + print(g) + outfile = os.path.join(self.args.output, 'groups.pkl') + with open(outfile, 'wb') as outF: + pickle.dump(self.scaffold_comparison_groups, file=outF) + logging.info('Groups written to {}; exiting'.format(outfile)) + exit(0) + # selecting one group for analysis (if required) + if self.args.group is not None: + if self.args.group_pkl is not None: + self.scaffold_comparison_groups = pickle.load(open(self.args.group_pkl, 'rb')) + else: + msg = 'If using --group, you must provide the groups.pkl file via --group-pkl' + raise ValueError(msg) + logging.info('User provided group number {}'.format(self.args.group)) + self.scaffold_comparison_groups = [self.scaffold_comparison_groups[self.args.group-1]] + # processing groups + group_len = len(self.scaffold_comparison_groups) + for i,SCgroup in enumerate(self.scaffold_comparison_groups): + if self.args.group is not None: + logging.info('Running group {}'.format(self.args.group)) + else: + logging.info(f'Running group {i+1} of {group_len}') SCgroup.load_cache() - results = inStrain.compare_utils.run_compare_multiprocessing(SCgroup.cmd_queue, SCgroup.result_queue, - self.null_model, num_to_run=len(SCgroup.scaffolds), + results = inStrain.compare_utils.run_compare_multiprocessing(SCgroup.cmd_queue, + SCgroup.result_queue, + self.null_model, + num_to_run=len(SCgroup.scaffolds), **self.kwargs) + if self.args.group is not None: + outdir = os.path.join(self.args.output, 'groups') + if not os.path.isdir(outdir): + os.makedirs(outdir) + outfile = os.path.join(outdir, '{}.pkl'.format(self.args.group)) + with open(outfile, 'wb') as outF: + pickle.dump(results, file=outF) + logging.info('Group result written to {}'.format(outfile)) + logging.info('exiting') + exit(0) + for result in results: if result is not None: Cdb, Mdb, pair2mm2covOverlap, scaffold = result - for item, lis in zip([Cdb, Mdb, pair2mm2covOverlap, scaffold], [cdbs, mdbs, pair2mm2covOverlaps, order]): + for item, lis in zip([Cdb, Mdb, pair2mm2covOverlap, scaffold], + [cdbs, mdbs, pair2mm2covOverlaps, order]): lis.append(item) SCgroup.purge_cache() @@ -200,9 +241,45 @@ def run_comparisons(self): # Process results self.process_results(cdbs, mdbs, pair2mm2covOverlaps, order) + def gather_comparisons(self): + """ + Gather results of previous comparisons + """ + # comparisons provided? + ## as a file listing all comparison result pkl files? + if self.args.comparisons_list is not None: + self.args.comparisons = [] + with open(self.args.comparisions_list) as inF: + for line in inF: + line = line.rstrip() + if line != '': + self.args.comparisons.append(line) + ## as a list of files + if not isinstance(self.args.comparisons, list): + self.args.comparisons = [self.args.comparisons] + if self.args.comparisons[0] is None: + # run comparisons instead of just gathering/merging them + self.run_comparisons() + return None + # merging individual group-level comparisions + cdbs = [] # Store scaffold comparison information + mdbs = [] # Store mismatch locations + pair2mm2covOverlaps = [] # Store coverage overlap locations + order = [] # Store order of scaffolds + for F in self.args.comparisons: + results = pickle.load(open(F, 'rb')) + for result in results: + if result is not None: + Cdb, Mdb, pair2mm2covOverlap, scaffold = result + for item, lis in zip([Cdb, Mdb, pair2mm2covOverlap, scaffold], + [cdbs, mdbs, pair2mm2covOverlaps, order]): + lis.append(item) + # processing results + self.process_results(cdbs, mdbs, pair2mm2covOverlaps, order) + def process_results(self, cdbs, mdbs, pair2mm2covOverlaps, order): """ - Merge and store results + Merge and store results for all group comparisons """ if len(pair2mm2covOverlaps) > 0: scaff2pair2mm2overlap = {} @@ -563,4 +640,4 @@ def establish_SC_groups(valid_SCs, Cdb, group_length): # Just do it simple; would be better to use the distance matrix, but oh well SC_groups = simple_grouping(valid_SCs, group_length) - return SC_groups \ No newline at end of file + return SC_groups