attempt at switch for calculating population frequencies

Author: savitakartik

tsqc/__main__.py

@@ -21,14 +21,15 @@
 logger = daiquiri.getLogger("tsqc")
 
 
-def load_data(path):
+def load_data(path, calc_population_frequencies):
     logger.info(f"Loading {path}")
     try:
         ts = tskit.load(path)
+
     except tskit.FileFormatError:
         ts = tszip.decompress(path)
 
-    tsm = model.TSModel(ts, path.name)
+    tsm = model.TSModel(ts, calc_population_frequencies, path.name)
     return tsm
 
 
@@ -152,13 +153,26 @@ def setup_logging(log_level, no_log_filter):
     is_flag=True,
     help="Do not filter the output log (advanced debugging only)",
 )
-def main(path, port, show, log_level, no_log_filter, annotations_file):
+@click.option(
+    "--calc-population-frequencies",
+    default=False,
+    help="Calculate population frequencies for sample nodes",
+)
+def main(
+    path,
+    port,
+    show,
+    log_level,
+    no_log_filter,
+    annotations_file,
+    calc_population_frequencies,
+):
     """
     Run the tsqc server.
     """
     setup_logging(log_level, no_log_filter)
 
-    tsm = load_data(pathlib.Path(path))
+    tsm = load_data(pathlib.Path(path), calc_population_frequencies)
     if annotations_file:
         config.ANNOTATIONS_FILE = annotations_file
 

tsqc/model.py

@@ -356,7 +356,7 @@ class TSModel:
     convenience methods for analysing the tree sequence.
     """
 
-    def __init__(self, ts, name=None):
+    def __init__(self, ts, calc_population_frequencies, name=None):
         self.ts = ts
         self.name = name
 
@@ -366,6 +366,7 @@ def __init__(self, ts, name=None):
         self.nodes_num_mutations = np.bincount(
             self.ts.mutations_node, minlength=self.ts.num_nodes
         )
+        self.calc_population_frequencies = calc_population_frequencies
 
     @property
     def file_uuid(self):
@@ -456,7 +457,7 @@ def mutations_df(self):
         self.mutations_inherited_state = inherited_state
 
         population_data = {}
-        if ts.num_populations > 0:
+        if ts.num_populations > 0 and self.calc_population_frequencies:
             pop_mutation_count = compute_population_mutation_counts(ts)
             for pop in ts.populations():
                 name = f"pop{pop.id}"