Merge pull request jeromekelleher#336 from jeromekelleher/generalise-…

…mutation-plot Add abililty to plot classes of mutation along genome
szhan · Oct 7, 2024 · 88baf9a · 88baf9a
2 parents 0195a10 + d9a5dfd
commit 88baf9a
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/sc2ts/info.py b/sc2ts/info.py
@@ -1326,8 +1326,12 @@ def plot_ts_tv_per_site(self, annotate_threshold=0.9, xlim=None):
         if xlim is not None:
             plt.xlim(xlim)
 
-    def plot_mutations_per_site(self, annotate_threshold=0.9):
-        count = self.sites_num_mutations
+    def plot_mutations_per_site(self, annotate_threshold=0.9, select=None):
+        if select is None:
+            count = self.sites_num_mutations
+        else:
+            count = np.bincount(self.ts.mutations_site[select], minlength=self.ts.num_sites)
+
         pos = self.ts.sites_position
         zero_fraction = np.sum(count == 0) / self.ts.num_sites