refactoring DNSMDataset methods to be subclass-friendly

matsengrp · Sep 30, 2024 · 39e52c2 · 39e52c2
1 parent b342ced
commit 39e52c2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/netam/dnsm.py b/netam/dnsm.py
@@ -136,7 +136,7 @@ def of_pcp_df(cls, pcp_df, branch_length_multiplier=5.0):
 
     def clone(self):
         """Make a deep copy of the dataset."""
-        new_dataset = DNSMDataset(
+        new_dataset = self.__class__(
             self.nt_parents,
             self.nt_children,
             self.all_rates.copy(),
@@ -152,7 +152,7 @@ def subset_via_indices(self, indices):
         depends on `indices`: if `indices` is an iterable of integers, then we
         make a deep copy, otherwise we use slices to make a shallow copy.
         """
-        new_dataset = DNSMDataset(
+        new_dataset = self.__class__(
             self.nt_parents[indices].reset_index(drop=True),
             self.nt_children[indices].reset_index(drop=True),
             self.all_rates[indices],