ecDNA simulator

cassiopeia/mixins/errors.py

@@ -27,6 +27,10 @@ class DistanceSolverError(Exception):
 
     pass
 
+class ecDNABirthDeathSimulatorError(Exception):
+    """An ExceptionClass for ecDNABirthDeathSimulator class."""
+
+    pass
 
 class FitchCountError(Exception):
     """An ExceptionClass for FitchCount."""

cassiopeia/plotting/local.py

@@ -185,6 +185,7 @@ def create_indel_heatmap(
     indel_colors: Optional[pd.DataFrame] = None,
     indel_priors: Optional[pd.DataFrame] = None,
     random_state: Optional[np.random.RandomState] = None,
+    clustered_linprof: Optional[pd.DataFrame] = None,
 ) -> Tuple[
     List[
         Dict[

cassiopeia/simulator/BirthDeathFitnessSimulator.py

@@ -94,6 +94,7 @@ class BirthDeathFitnessSimulator(TreeSimulator):
         collapse_unifurcations: Specifies whether to collapse unifurcations in
             the tree resulting from pruning dead lineages
         random_seed: A seed for reproducibility
+        initial_tree: A tree used for initializing the simulation.
 
     Raises:
         TreeSimulatorError if invalid stopping conditions are provided or if a
@@ -114,6 +115,7 @@ def __init__(
         experiment_time: Optional[float] = None,
         collapse_unifurcations: bool = True,
         random_seed: int = None,
+        initial_tree: Optional[CassiopeiaTree] = None,
     ):
         if num_extant is None and experiment_time is None:
             raise TreeSimulatorError(
@@ -147,6 +149,85 @@ def __init__(
         self.collapse_unifurcations = collapse_unifurcations
         self.random_seed = random_seed
 
+        # useful for resuming a simulation, perhaps under different pressures.
+        self.initial_tree = initial_tree
+
+    def initialize_tree(self, names) -> nx.DiGraph:
+        """initializes a tree (nx.DiGraph() object with one node). Auxiliary data for each node is grabbed from self (initial conditions / params) or hardcoded
+
+        Args: names is a generator (function object that stores internal state) that will be used to generate names for the tree nodes
+        Output: tree (DiGraph object with one node, the root)
+                root (name of root node in tree)
+        """
+        if self.initial_tree:
+            tree = self.initial_tree.get_tree_topology()
+            for node in self.initial_tree.nodes:
+                tree.nodes[node]['birth_scale'] = self.initial_tree.get_attribute(node, 'birth_scale')
+                tree.nodes[node]['time'] = self.initial_tree.get_attribute(node, 'time')
+            return tree
+
+        tree = nx.DiGraph()
+        root = next(names)
+        tree.add_node(root)
+        tree.nodes[root]["birth_scale"] = self.initial_birth_scale
+        tree.nodes[root]["time"] = 0
+
+        return tree
+
+    def make_initial_lineage_dict(self, tree: nx.DiGraph):
+        """
+        uses self initial-conditions and hardcoded default parameters to create an initial lineage dict
+        Args: id_value: name of new lineage
+
+        Output: a lineage dict
+        """
+
+        leaves = [node for node in tree if tree.out_degree(node) == 0]
+        current_lineages = PriorityQueue()
+        for leaf in leaves:
+
+            lineage_dict = self.make_lineage_dict(
+                    leaf, tree.nodes[leaf]['birth_scale'], tree.nodes[leaf]['time'], True
+                )
+
+            if len(tree.nodes) == 1:
+                return lineage_dict
+
+            current_lineages.put(
+                (
+                    tree.nodes[leaf]['time'],
+                    leaf,
+                    lineage_dict
+                )
+            )
+
+        return current_lineages
+
+    def make_lineage_dict(
+        self,
+        id_value,
+        birth_scale,
+        total_time,
+        active_flag,
+    ):
+        """makes a dict (lineage) from the given parameters. keys are hardcoded.
+
+        Args: id_value: id of new lineage
+              birth_scale: birth_scale parameter of new lineage
+              total_time: age of lineage
+              active_flag: bool to indicate whether lineage is active
+
+        Returns: a dict (lineage) with the parameter values under the hard-coded keys
+
+        """
+        lineage_dict = {
+            "id": id_value,
+            "birth_scale": birth_scale,
+            "total_time": total_time,
+            "active": active_flag,
+        }
+        return lineage_dict
+
     def simulate_tree(
         self,
     ) -> CassiopeiaTree:
@@ -169,39 +250,39 @@ def simulate_tree(
             TreeSimulatorError if all lineages die before a stopping condition
         """
 
-        def node_name_generator() -> Generator[str, None, None]:
+        def node_name_generator(start = 0) -> Generator[str, None, None]:
             """Generates unique node names for the tree."""
-            i = 0
+            i = start
             while True:
                 yield str(i)
                 i += 1
 
-        names = node_name_generator()
+        starting_index = 0
+        if self.initial_tree:
+           starting_index = (np.max([int(l) for l in self.initial_tree.leaves]) + 1)
+        names = node_name_generator(starting_index)
 
         # Set the seed
         if self.random_seed:
             np.random.seed(self.random_seed)
+
+        tree = self.initialize_tree(names)
 
-        # Instantiate the implicit root
-        tree = nx.DiGraph()
-        root = next(names)
-        tree.add_node(root)
-        tree.nodes[root]["birth_scale"] = self.initial_birth_scale
-        tree.nodes[root]["time"] = 0
-        current_lineages = PriorityQueue()
+        current_lineages = PriorityQueue() # instantiate queue
         # Records the nodes that are observed at the end of the experiment
+
+        # TO DO: update to accept arbitrary fields in the dict.
         observed_nodes = []
-        starting_lineage = {
-            "id": root,
-            "birth_scale": self.initial_birth_scale,
-            "total_time": 0,
-            "active": True,
-        }
 
-        # Sample the waiting time until the first division
-        self.sample_lineage_event(
-            starting_lineage, current_lineages, tree, names, observed_nodes
-        )
+        starting_lineage = self.make_initial_lineage_dict(tree)
+
+        if len(tree.nodes) == 1:
+            # Sample the waiting time until the first division
+            self.sample_lineage_event(
+                starting_lineage, current_lineages, tree, names, observed_nodes
+            )
+        else:
+            current_lineages = starting_lineage
 
         # Perform the process until there are no active extant lineages left
         while not current_lineages.empty():
@@ -219,6 +300,7 @@ def node_name_generator() -> Generator[str, None, None]:
                     min_total_time = remaining_lineages[0]["total_time"]
                     for lineage in remaining_lineages:
                         parent = list(tree.predecessors(lineage["id"]))[0]
+
                         tree.nodes[lineage["id"]]["time"] += (
                             min_total_time - lineage["total_time"]
                         )
@@ -229,31 +311,18 @@ def node_name_generator() -> Generator[str, None, None]:
                     break
             # Pop the minimum age lineage to simulate forward time
             _, _, lineage = current_lineages.get()
+
             # If the lineage is no longer active, just remove it from the queue.
             # This represents the time at which the lineage dies.
             if lineage["active"]:
-                for _ in range(2):
+                for i in range(2):
                     self.sample_lineage_event(
                         lineage, current_lineages, tree, names, observed_nodes
                     )
 
-        cassiopeia_tree = CassiopeiaTree(tree=tree)
-        time_dictionary = {}
-        for i in tree.nodes:
-            time_dictionary[i] = tree.nodes[i]["time"]
-        cassiopeia_tree.set_times(time_dictionary)
-
-        # Prune dead lineages and collapse resulting unifurcations
-        to_remove = list(set(cassiopeia_tree.leaves) - set(observed_nodes))
-        cassiopeia_tree.remove_leaves_and_prune_lineages(to_remove)
-        if self.collapse_unifurcations and len(cassiopeia_tree.nodes) > 1:
-            cassiopeia_tree.collapse_unifurcations(source="1")
-
-        # If only implicit root remains after pruning dead lineages, error
-        if len(cassiopeia_tree.nodes) == 1:
-            raise TreeSimulatorError(
-                "All lineages died before stopping condition"
-            )
+        cassiopeia_tree = self.populate_tree_from_simulation(
+            tree, observed_nodes
+        )
 
         return cassiopeia_tree
 
@@ -266,7 +335,6 @@ def sample_lineage_event(
         observed_nodes: List[str],
     ) -> None:
         """A helper function that samples an event for a lineage.
-
         Takes a lineage and determines the next event in that lineage's
         future. Simulates the lifespan of a new descendant. Birth and
         death waiting times are sampled, representing how long the
@@ -285,7 +353,6 @@ def sample_lineage_event(
         the lifespan is cut off at the experiment time and a final observed
         sample is added to the tree. In this case the lineage is marked as
         inactive as well.
-
         Args:
             unique_id: The unique ID number to be used to name a new node
                 added to the tree
@@ -300,7 +367,6 @@ def sample_lineage_event(
             names: A generator providing unique names for tree nodes
             observed_nodes: A list of nodes that are observed at the end of
                 the experiment
-
         Raises:
             TreeSimulatorError if a negative waiting time is sampled or a
             non-active lineage is passed in
@@ -378,6 +444,7 @@ def sample_lineage_event(
                     )
                 )
 
+
             else:
                 tree.add_node(unique_id)
                 tree.nodes[unique_id]["birth_scale"] = lineage["birth_scale"]
@@ -435,3 +502,38 @@ def update_fitness(self, birth_scale: float) -> float:
                     self.fitness_base ** self.fitness_distribution()
                 )
         return birth_scale * base_selection_coefficient
+
+    def populate_tree_from_simulation(
+        self, tree: nx.DiGraph, observed_nodes: List[str]
+    ) -> CassiopeiaTree:
+        """Populates tree with appropriate meta data.
+
+        Args:
+            tree: The tree simulated with ecDNA and fitness values populated as attributes.
+            observed_nodes: The observed leaves of the tree.
+
+        Returns:
+            A CassiopeiaTree with relevant node attributes filled in.
+        """
+
+        cassiopeia_tree = CassiopeiaTree(tree=tree)
+
+        time_dictionary = {}
+        for i in tree.nodes:
+            time_dictionary[i] = tree.nodes[i]["time"]
+            cassiopeia_tree.set_attribute(i, 'birth_scale', tree.nodes[i]['birth_scale'])
+        cassiopeia_tree.set_times(time_dictionary)
+
+        # Prune dead lineages and collapse resulting unifurcations
+        to_remove = list(set(cassiopeia_tree.leaves) - set(observed_nodes))
+        cassiopeia_tree.remove_leaves_and_prune_lineages(to_remove)
+        if self.collapse_unifurcations and len(cassiopeia_tree.nodes) > 1:
+            cassiopeia_tree.collapse_unifurcations(source="1")
+
+        # If only implicit root remains after pruning dead lineages, error
+        if len(cassiopeia_tree.nodes) == 1:
+            raise TreeSimulatorError(
+                "All lineages died before stopping condition"
+            )
+
+        return cassiopeia_tree

cassiopeia/simulator/UniformLeafSubsampler.py

@@ -120,4 +120,5 @@ def subsample_leaves(
                 [(node, tree.get_time(node)) for node in subsampled_tree.nodes]
             )
         )
+
         return subsampled_tree

cassiopeia/simulator/__init__.py

@@ -9,6 +9,7 @@
 from .ClonalSpatialDataSimulator import ClonalSpatialDataSimulator
 from .CompleteBinarySimulator import CompleteBinarySimulator
 from .DataSimulator import DataSimulator
+from .ecDNABirthDeathSimulator import ecDNABirthDeathSimulator
 from .LeafSubsampler import LeafSubsampler
 from .LineageTracingDataSimulator import LineageTracingDataSimulator
 from .SimpleFitSubcloneSimulator import SimpleFitSubcloneSimulator
@@ -25,6 +26,7 @@
     "SeqeuntialLineageTracingDataSimulator",
     "CompleteBinarySimulator",
     "DataSimulator",
+    "ecDNABirthDeathSimulator",
     "LeafSubsampler",
     "LineageTracingDataSimulator",
     "SimpleFitSubcloneSimulator",