Merge pull request #533 from jstoja/other/use_defaultdict

Use defaultdict where possible

augur/clades.py

@@ -6,6 +6,7 @@
 from Bio import Phylo
 import pandas as pd
 import numpy as np
+from collections import defaultdict
 from .utils import get_parent_name_by_child_name_for_tree, read_node_data, write_json, get_json_name
 
 def read_in_clade_definitions(clade_file):
@@ -31,15 +32,12 @@ def read_in_clade_definitions(clade_file):
         clade definitions as :code:`{clade_name:[(gene, site, allele),...]}`
     '''
 
-    clades = {}
-
+    clades = defaultdict(list)
     df = pd.read_csv(clade_file, sep='\t' if clade_file.endswith('.tsv') else ',')
     for index, row in df.iterrows():
         allele = (row.gene, row.site-1, row.alt)
-        if row.clade in clades:
-            clades[row.clade].append(allele)
-        else:
-            clades[row.clade] = [allele]
+        clades[row.clade].append(allele)
+    clades.default_factory = None
 
     return clades
 

augur/import_beast.py

@@ -7,6 +7,7 @@
 import json
 import datetime as dt
 from argparse import SUPPRESS
+from collections import defaultdict
 import numpy as np
 from Bio import Phylo
 from treetime import TreeAnc
@@ -506,17 +507,13 @@ def compute_entropies_for_discrete_traits(tree):
 
     Author: James Hadfield
     """
-    alphabets={} ## store alphabets
+    alphabets = defaultdict(list) ## store alphabets
     for clade in tree.find_clades(): ## iterate over branches
         for attr in [key for key in clade.attrs if isinstance(clade.attrs[key], dict)]: ## iterate over branch attributes
-            if attr in alphabets: ## if attr seen before
-                for val in clade.attrs[attr]: ## iterate over attribute values of the node
-                    if val not in alphabets[attr]: ## not seen this attribute value before
-                        alphabets[attr].append(val)
-            else:
-                alphabets[attr]=[] ## not seen trait before - start a list of its values
-                for val in clade.attrs[attr]: ## iterate over trait values for this branch
+            for val in clade.attrs[attr]: ## iterate over attribute values of the node
+                if val not in alphabets[attr]: ## not seen this attribute value before
                     alphabets[attr].append(val)
+    alphabets.default_factory = None
 
     for clade in tree.find_clades(): ## iterate over branches
         for trait in alphabets: ## iterate over traits