api change

phipkit/antigen_analysis.py

@@ -66,22 +66,17 @@ def compute_coverage(antigen, sequence, blast_df):
 
 
 class AntigenAnalysis(object):
-    def __init__(self, blast_df, antigens_df, hits_df=None, sample_to_kind=None):
+    def __init__(self, blast_df, antigens_df, sample_to_hit_clones, sample_to_kind=None):
         self.blast_df = blast_df
         self.antigens_df = antigens_df
-        self.hits_df = hits_df
         self.sample_to_kind = sample_to_kind
 
-        say("Collecting hits.")
-        self.sample_to_hit_clones = collections.defaultdict(set)
-        all_clones = set()
-        for _, row in tqdm(self.hits_df.iterrows(), total=len(self.hits_df)):
-            self.sample_to_hit_clones[row.sample_id].add(row.clone1)
-            self.sample_to_hit_clones[row.sample_id].add(row.clone2)
-            all_clones.add(row.clone1)
-            all_clones.add(row.clone2)
         self.sample_to_hit_clones = dict(
-            (k, list(v)) for (k, v) in self.sample_to_hit_clones.items())
+            (k, list(v)) for (k, v) in sample_to_hit_clones.items())
+
+        all_clones = set()
+        for clones in self.sample_to_hit_clones.values():
+            all_clones.update(clones)
 
         self.clone_by_sample_hits_matrix = pandas.DataFrame(
             index=sorted(all_clones),

phipkit/call_antigens.py

@@ -65,7 +65,7 @@
 
 import pandas
 
-from . common import say, reconstruct_antigen_sequences
+from .common import say, reconstruct_antigen_sequences, hits_to_dict
 
 parser = argparse.ArgumentParser(
     description=__doc__,
@@ -168,18 +168,6 @@ def run(argv=sys.argv[1:]):
     say("Wrote: ", args.out)
 
 
-def hits_to_dict(hits_df):
-    """
-    Given a hits_df, return a dict of sample id -> list of hits
-    """
-    sample_to_clones = {}
-    for sample, sub_hits_df in hits_df.groupby("sample_id"):
-        sample_to_clones[sample] = sub_hits_df[
-            ["clone1", "clone2"]
-        ].stack().unique()
-    return sample_to_clones
-
-
 def find_consensus(sequences, threshold=0.7):
     """
     Given aligned sequences, return a string where each position i is the

phipkit/call_hits.py

@@ -45,7 +45,7 @@
 
 import pandas
 
-from . common import say
+from .common import say
 
 DEFAULTS = {
     'fdr': 0.01,

phipkit/common.py

@@ -44,4 +44,16 @@ def reconstruct_antigen_sequences(blast_df):
 
         antigen_sequences[title][hit_from - 1: hit_to] = hseq.encode('ascii')
     antigen_sequences = antigen_sequences.map(lambda arr: arr.decode())
-    return antigen_sequences
+    return antigen_sequences
+
+
+def hits_to_dict(hits_df):
+    """
+    Given a hits_df, return a dict of sample id -> list of hits
+    """
+    sample_to_clones = {}
+    for sample, sub_hits_df in hits_df.groupby("sample_id"):
+        sample_to_clones[sample] = sub_hits_df[
+            ["clone1", "clone2"]
+        ].stack().unique()
+    return sample_to_clones

phipkit/plot_antigens.py

@@ -20,7 +20,7 @@
 from matplotlib import pyplot
 from matplotlib.backends.backend_pdf import PdfPages
 
-from . common import say
+from .common import say, hits_to_dict
 from .antigen_analysis import AntigenAnalysis
 
 parser = argparse.ArgumentParser(
@@ -121,8 +121,9 @@ def plot_antigens(blast_df, hits_df, antigens_df, out, include_redundant=False):
 
     analyzer = AntigenAnalysis(
         blast_df=blast_df,
-        hits_df=hits_df,
-        antigens_df=antigens_df)
+        antigens_df=antigens_df,
+        sample_to_hit_clones=hits_to_dict(hits_df))
+
 
     say("Generating plots")
     antigens = antigens_df.antigen.unique()

test/test_integration.py

@@ -305,7 +305,6 @@ def test_integrated(save_dir=None):
     assert os.path.exists(out)
 
 
-
 if __name__ == "__main__":
     import sys
     import argparse