update stairway.py

popsim-consortium · Oct 3, 2023 · 433e584 · 433e584
1 parent 6769863
commit 433e584
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/workflows/stairway.py b/workflows/stairway.py
@@ -73,9 +73,10 @@ def ts_to_stairway(self, ts_path, pop_name, mask_intervals, num_bootstraps=1):
             class_muts = {}
             for dfe in ts.metadata["stdpopsim"]["DFEs"]:
                 for mt in dfe["mutation_types"]:
-                    mid = mt["slim_mutation_type_id"]
-                    if not mid in class_muts:
-                        class_muts[mid] = "neutral" if mt["is_neutral"] else "non_neutral"
+                    mids = mt["slim_mutation_type_id"]
+                    for mid in mids:
+                        if not mid in class_muts:
+                            class_muts[mid] = "neutral" if mt["is_neutral"] else "non_neutral"
 
             site_class = np.empty(ts.num_sites, dtype=object)