Store masked alignment in masked_base_composition

szhan · Sep 30, 2024 · ec7cc23 · ec7cc23
1 parent 8cb123e
commit ec7cc23
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/sc2ts/alignments.py b/sc2ts/alignments.py
@@ -155,12 +155,14 @@ class MaskedAlignment:
     masked_sites: np.ndarray
     original_base_composition: dict
     original_md5: str
+    masked_base_composition: str
 
     def qc_summary(self):
         return {
             "num_masked_sites": self.masked_sites.shape[0],
             "original_base_composition": self.original_base_composition,
             "original_md5": self.original_md5,
+            "masked_base_composition": self.masked_base_composition,
         }
 
 
@@ -171,9 +173,10 @@ def encode_and_mask(alignment, window_size=7):
     return MaskedAlignment(
         alignment=a,
         masked_sites=np.array(masked_sites, dtype=int),
-        original_base_composition=base_composition(
+        original_base_composition=base_composition(haplotype=alignment),
+        original_md5=hashlib.md5(alignment[1:]).hexdigest(),
+        masked_base_composition=base_composition(
             haplotype=alignment,
             excluded_sites=masked_sites,
         ),
-        original_md5=hashlib.md5(alignment[1:]).hexdigest(),
     )