Repository: YosefLab/Cassiopeia (language: python)

Issue:
Supercellular sampler
`SupercellularSampler` implementation, along with several utility functions in `CassiopeiaTree` that are useful in dealing with ambiguous states and branch distances.

Some important changes are the following.
* Ambiguous characters are denoted as a tuple in the character matrix and/or the network. Updates to some functions and addition of new functions in `CassiopeaTree` to support ambiguous characters.
* When computing dissimilarities with `CassiopeiaTree.compute_dissimilarity_map`, a new `cluster_dissimilarity` function should be used when there are ambiguous characters. Here is an example snippet

```python
from functools import partial
dissimilarity_function = partial(cluster_dissimilarity, weighted_hamming_distance)
tree.compute_dissimilarity_map(dissimilarity_function)
```
I had considered whether this should be done automatically in ``compute_dissimilarity_map``, but felt that the arguments should always be explicit for the user.

* Updated `n_cell`, `n_character` and `populate_tree` in `CassiopeiaTree` to use the current character matrix instead of the original one because these properties/function should always return/operate on the current object state, instead of a possibly-outdated state (i.e. the original character matrix). Previously, the original (old) state was being used and this was causing problems for `GreedySolver` when trying to call `collapse_mutationless_edges`.

This PR is big enough for reviews from both @mattjones315 and @richardyz98.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 5b4df2e7..aa137100 100755
--- a/.gitignore
+++ b/.gitignore
@@ -5,3 +5,8 @@
 _build
 _static
 _templates
+build
+*.egg-info
+*.c
+stdout.log
+notebooks/.ipynb_checkpoints
diff --git a/cassiopeia/__main__.py b/cassiopeia/__main__.py
deleted file mode 100755
index b61d36aa..00000000
--- a/cassiopeia/__main__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-#!/usr/bin/env python
-
-import cassiopeia
-
-
-def main():
-    print("Congrats, you've entered the program!")
diff --git a/cassiopeia/critique/compare.py b/cassiopeia/critique/compare.py
index fa10f281..6cc63da5 100644
--- a/cassiopeia/critique/compare.py
+++ b/cassiopeia/critique/compare.py
@@ -16,13 +16,16 @@
 
 
 def triplets_correct(
-    tree1: CassiopeiaTree, tree2: CassiopeiaTree, number_of_trials: int = 1000, min_triplets_at_depth: int = 1
+    tree1: CassiopeiaTree,
+    tree2: CassiopeiaTree,
+    number_of_trials: int = 1000,
+    min_triplets_at_depth: int = 1,
 ) -> Tuple[
     Dict[int, float], Dict[int, float], Dict[int, float], Dict[int, float]
 ]:
     """Calculate the triplets correct accuracy between two trees.
 
-    Takes in two newick strings and computes the proportion of triplets in the 
+    Takes in two newick strings and computes the proportion of triplets in the
     tree (defined as a set of three leaves) that are the same across the two
     trees. This procedure samples the same number of triplets at every depth
     such as to reduce the amount of bias of sampling triplets randomly.
@@ -50,7 +53,7 @@ def triplets_correct(
     unresolved_triplets_correct = defaultdict(int)
     resolvable_triplets_correct = defaultdict(int)
     proportion_unresolvable = defaultdict(int)
-    
+
     # create copies of the trees and collapse process
     T1 = copy.copy(tree1)
     T2 = copy.copy(tree2)
@@ -70,7 +73,9 @@ def triplets_correct(
 
         # check that there are enough triplets at this depth
         candidate_nodes = depth_to_nodes[depth]
-        total_triplets = sum([T1.get_attribute(v, "number_of_triplets") for v in candidate_nodes])
+        total_triplets = sum(
+            [T1.get_attribute(v, "number_of_triplets") for v in candidate_nodes]
+        )
         if total_triplets < min_triplets_at_depth:
             continue
 
@@ -80,13 +85,15 @@ def triplets_correct(
                 T1, depth, depth_to_nodes
             )
 
-            reconstructed_outgroup = critique_utilities.get_outgroup(T2, (i, j, k))
+            reconstructed_outgroup = critique_utilities.get_outgroup(
+                T2, (i, j, k)
+            )
 
             is_resolvable = True
             if out_group == "None":
                 number_unresolvable_triplets += 1
                 is_resolvable = False
-            
+
             # increment score if the reconstructed outgroup is the same as the
             # ground truth
             score = int(reconstructed_outgroup == out_group)
diff --git a/cassiopeia/critique/critique_utilities.py b/cassiopeia/critique/critique_utilities.py
index e5e5f38d..8b48fac5 100644
--- a/cassiopeia/critique/critique_utilities.py
+++ b/cassiopeia/critique/critique_utilities.py
@@ -12,7 +12,7 @@
 from cassiopeia.data import CassiopeiaTree
 
 
-def nCr(n:int , r:int) -> float:
+def nCr(n: int, r: int) -> float:
     """Computes nCr
 
     Args:
@@ -25,7 +25,8 @@ def nCr(n:int , r:int) -> float:
     if r > n or n < 0 or r < 0:
         return 0
     f = math.factorial
-    return f(n) // f(r) // f(n-r)
+    return f(n) // f(r) // f(n - r)
+
 
 def annotate_tree_depths(tree: CassiopeiaTree) -> None:
     """Annotates tree depth at every node.
@@ -36,17 +37,19 @@ def annotate_tree_depths(tree: CassiopeiaTree) -> None:
 
     Args:
         tree: An ete3 Tree
-        
+
     Returns:
         A dictionary mapping depth to the list of nodes at that depth.
     """
 
     depth_to_nodes = defaultdict(list)
-    for n in tree.depth_first_traverse_nodes(source = tree.root, postorder=False):
+    for n in tree.depth_first_traverse_nodes(source=tree.root, postorder=False):
         if tree.is_root(n):
             tree.set_attribute(n, "depth", 0)
         else:
-            tree.set_attribute(n, "depth", tree.get_attribute(tree.parent(n), "depth") + 1)
+            tree.set_attribute(
+                n, "depth", tree.get_attribute(tree.parent(n), "depth") + 1
+            )
 
         depth_to_nodes[tree.get_attribute(n, "depth")].append(n)
 
@@ -56,10 +59,13 @@ def annotate_tree_depths(tree: CassiopeiaTree) -> None:
             number_of_leaves += len(tree.leaves_in_subtree(child))
             correction += nCr(len(tree.leaves_in_subtree(child)), 3)
 
-        tree.set_attribute(n, "number_of_triplets", nCr(number_of_leaves, 3) - correction)
+        tree.set_attribute(
+            n, "number_of_triplets", nCr(number_of_leaves, 3) - correction
+        )
 
     return depth_to_nodes
 
+
 def get_outgroup(tree: CassiopeiaTree, triplet: Tuple[str, str, str]) -> str:
     """Infers the outgroup of a triplet from a CassioepiaTree.
 
@@ -95,13 +101,16 @@ def get_outgroup(tree: CassiopeiaTree, triplet: Tuple[str, str, str]) -> str:
         out_group = i
     return out_group
 
+
 def sample_triplet_at_depth(
-    tree: CassiopeiaTree, depth: int, depth_to_nodes: Optional[Dict[int, List[str]]] = None,
+    tree: CassiopeiaTree,
+    depth: int,
+    depth_to_nodes: Optional[Dict[int, List[str]]] = None,
 ) -> Tuple[List[int], str]:
     """Samples a triplet at a given depth.
 
     Samples a triplet of leaves such that the depth of the LCA of the triplet
-    is at the specified depth. 
+    is at the specified depth.
 
     Args:
         tree: CassiopeiaTree
@@ -115,15 +124,22 @@ def sample_triplet_at_depth(
     """
 
     if depth_to_nodes is None:
-        candidate_nodes = tree.filter_nodes(lambda x: tree.get_attribute(x, "depth") == depth)
+        candidate_nodes = tree.filter_nodes(
+            lambda x: tree.get_attribute(x, "depth") == depth
+        )
     else:
         candidate_nodes = depth_to_nodes[depth]
 
-    total_triplets = sum([tree.get_attribute(v, "number_of_triplets") for v in candidate_nodes])
+    total_triplets = sum(
+        [tree.get_attribute(v, "number_of_triplets") for v in candidate_nodes]
+    )
 
     # sample a  node from this depth with probability proportional to the number
     # of triplets underneath it
-    probs = [tree.get_attribute(v, "number_of_triplets") / total_triplets for v in candidate_nodes]
+    probs = [
+        tree.get_attribute(v, "number_of_triplets") / total_triplets
+        for v in candidate_nodes
+    ]
     node = np.random.choice(candidate_nodes, size=1, replace=False, p=probs)[0]
 
     # Generate the probilities to sample each combination of 3 daughter clades
diff --git a/cassiopeia/data/CassiopeiaTree.py b/cassiopeia/data/CassiopeiaTree.py
old mode 100644
new mode 100755
index bef7e8d6..245a3e80
--- a/cassiopeia/data/CassiopeiaTree.py
+++ b/cassiopeia/data/CassiopeiaTree.py
@@ -6,14 +6,15 @@
 stored here, like the priors for given character states as well any meta data
 associated with this clonal  population.
 
-When a solver has been called on this object, a tree will be added to the data 
-structure at which point basic properties can be queried like the average tree 
+When a solver has been called on this object, a tree will be added to the data
+structure at which point basic properties can be queried like the average tree
 depth or agreement between character states and phylogeny.
 
 This object can be passed to any CassiopeiaSolver subclass as well as any
 analysis module, like a branch length estimator or rate matrix estimator
 """
 import copy
+import itertools
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 import warnings
 
@@ -58,7 +59,7 @@ class CassiopeiaTree:
     track both an unmodified version of the character matrix (obtainable via
     the `get_original_character_matrix` method) that does not maintain
     consistency with the character states of the leaves, and a working character
-    matrix (obtainable via the `get_modified_character_matrix` method) that
+    matrix (obtainable via the `get_current_character_matrix` method) that
     is updated when the character states of leaves are changed.
 
     Some reconstruction algorithms will make use of dissimilarities between
@@ -83,6 +84,8 @@ class CassiopeiaTree:
     TODO(mattjones315): Add bulk set_states method.
     TODO(mattjones): Add boolean to `get_tree_topology` which will include
         all attributes (e.g., node times)
+    TODO(Lioscro): Add support for a tree to hold multiple character matrices,
+        similar to "layers" of an AnnData object.
 
     Args:
         character_matrix: The character matrix for the lineage.
@@ -162,12 +165,12 @@ def populate_tree(self, tree: Union[str, ete3.Tree, nx.DiGraph]) -> None:
         # add character states
         for n in self.nodes:
             if (
-                self.__original_character_matrix is not None
-                and n in self.__original_character_matrix.index.tolist()
+                self.__current_character_matrix is not None
+                and n in self.__current_character_matrix.index.tolist()
             ):
                 self.__network.nodes[n][
                     "character_states"
-                ] = self.__original_character_matrix.loc[n].to_list()
+                ] = self.__current_character_matrix.loc[n].to_list()
             else:
                 self.__network.nodes[n]["character_states"] = []
 
@@ -190,8 +193,9 @@ def set_character_matrix(self, character_matrix: pd.DataFrame):
         """Initializes a character matrix in the object."""
 
         if not all(type(i) == str for i in character_matrix.index):
-            raise CassiopeiaTreeError("Index of character matrix must consist" 
-            " of strings.")
+            raise CassiopeiaTreeError(
+                "Index of character matrix must consist" " of strings."
+            )
 
         self.__original_character_matrix = character_matrix.copy()
         self.__current_character_matrix = character_matrix.copy()
@@ -316,13 +320,13 @@ def n_cell(self) -> int:
             CassiopeiaTreeError if the object is empty (i.e. no tree or
             character matrix).
         """
-        if self.__original_character_matrix is None:
+        if self.__current_character_matrix is None:
             if self.__network is None:
                 raise CassiopeiaTreeError(
                     "This is an empty object with no tree or character matrix."
                 )
             return len(self.leaves)
-        return self.__original_character_matrix.shape[0]
+        return self.__current_character_matrix.shape[0]
 
     @property
     def n_character(self) -> int:
@@ -333,7 +337,7 @@ def n_character(self) -> int:
             character matrix) or if the character states have not been
             initialized.
         """
-        if self.__original_character_matrix is None:
+        if self.__current_character_matrix is None:
             if self.__network is None:
                 raise CassiopeiaTreeError(
                     "This is an empty object with no tree or character matrix."
@@ -343,7 +347,7 @@ def n_character(self) -> int:
             raise CassiopeiaTreeError(
                 "Character states have not been initialized."
             )
-        return self.__original_character_matrix.shape[1]
+        return self.__current_character_matrix.shape[1]
 
     @property
     def root(self) -> str:
@@ -469,13 +473,104 @@ def is_internal_node(self, node: str) -> bool:
         self.__check_network_initialized()
         return self.__network.out_degree(node) > 0
 
+    def is_ambiguous(self, node: str) -> bool:
+        """Returns whether the node is ambiguous.
+
+        This is detected by checking if any of the characters are tuples.
+
+        Args:
+            node: Name of the node to check if it has ambiguous character(s)
+
+        Returns:
+            True if the node is ambiguous, False otherwise.
+
+        Raises:
+            CassiopeiaTreeError if the tree has not been initialized.
+        """
+        self.__check_network_initialized()
+        states = self.get_character_states(node)
+        return any(utilities.is_ambiguous_state(state) for state in states)
+
+    def collapse_ambiguous_characters(self) -> None:
+        """Only retain unique characters for ambiguous nodes.
+
+        Ambiguous nodes have character strings represented as a list of tuples
+        of integers. The inner list may contain multiple of the same states, encoding
+        the relative abundance of a certain state in the ambiguous state distribution.
+        Calling this function removes such duplicates and only retains unique
+        characters. This function is idempotent and does nothing for trees that
+        have no ambiguous characters.
+
+        Raises:
+            CassiopeiaTreeError if the tree has not been initialized.
+        """
+        self.__check_network_initialized()
+
+        for node in self.nodes:
+            if self.is_ambiguous(node):
+                states = self.get_character_states(node)
+                # Modify states in place. This is okay because get_character_states
+                # returns a copy.
+                modified = False
+                for i in range(len(states)):
+                    if utilities.is_ambiguous_state(states[i]):
+                        new_states = tuple(set(states[i]))
+                        if states != new_states:
+                            states[i] = new_states
+                            modified = True
+                if modified:
+                    self.set_character_states(node, states)
+
+    def resolve_ambiguous_characters(
+        self,
+        resolve_function: Callable[
+            [Tuple[int, ...]], int
+        ] = utilities.resolve_most_abundant,
+    ) -> None:
+        """Resolve all nodes with ambiguous characters.
+
+        A custom ``resolve_function`` may be provided to perform the resolution.
+        By default, the most abundant state is selected.
+        One is randomly selected on ties. Modifies the tree in-place.
+
+        Args:
+            resolve_function: Function that performs character resolution. This
+                function is called once per ambiguous character state, and thus
+                takes a single integer tuple as its argument and returns the
+                resolved character state.
+
+        Raises:
+            CassiopeiaTreeError if the tree has not been initialized.
+        """
+        self.__check_network_initialized()
+
+        for node in self.nodes:
+            if self.is_ambiguous(node):
+                states = self.get_character_states(node)
+
+                # Modify states in place. This is okay because get_character_states
+                # returns a copy.
+                modified = False
+                for i in range(len(states)):
+                    if utilities.is_ambiguous_state(states[i]):
+                        states[i] = resolve_function(states[i])
+                        modified = True
+                if modified:
+                    self.set_character_states(node, states)
+
+        # All nodes have a single character string, so we cast the character matrix
+        # back to an integer.
+        self.__current_character_matrix = (
+            self.__current_character_matrix.astype(int, copy=False)
+        )
+
     def reconstruct_ancestral_characters(self) -> None:
         """Reconstruct ancestral character states.
 
         Reconstructs ancestral states (i.e., those character states in the
         internal nodes) using the Camin-Sokal parsimony criterion (i.e.,
         irreversibility). Operates on the tree in place.
-        
+
         Raises:
             CassiopeiaTreeError if the tree has not been initialized.
         """
@@ -484,9 +579,10 @@ def reconstruct_ancestral_characters(self) -> None:
         for n in self.depth_first_traverse_nodes(postorder=True):
             if self.is_leaf(n):
                 if len(self.get_character_states(n)) == 0:
-                    raise CassiopeiaTreeError("Character states not annotated "
-                    "at a leaf node, initialize character states at leaves "
-                    "before reconstructing ancestral characters."
+                    raise CassiopeiaTreeError(
+                        "Character states not annotated "
+                        "at a leaf node, initialize character states at leaves "
+                        "before reconstructing ancestral characters."
                     )
                 continue
             children = self.children(n)
@@ -509,7 +605,7 @@ def parent(self, node: str) -> str:
             CassiopeiaTreeError if the tree is not initialized.
         """
         self.__check_network_initialized()
-        
+
         return [u for u in self.__network.predecessors(node)][0]
 
     def children(self, node: str) -> List[str]:
@@ -529,7 +625,7 @@ def children(self, node: str) -> List[str]:
 
     def __remove_node(self, node) -> None:
         """Private method to remove node from tree.
-        
+
         Args:
             node: A node in the tree to be removed
 
@@ -539,27 +635,27 @@ def __remove_node(self, node) -> None:
         self.__check_network_initialized()
 
         self.__network.remove_node(node)
-    
+
     def __add_node(self, node) -> None:
         """Private method to add node to tree.
-        
+
         Args:
             node: A node to be added to the tree.
-            
+
         Raises:
             CassiopeiaTreeError if the tree is not initialized.
         """
         self.__check_network_initialized()
-        
+
         self.__network.add_node(node)
 
     def __remove_edge(self, u, v) -> None:
         """Private method to remove edge from tree.
-        
+
         Args:
             u: The source node of the directed edge to be removed
             v: The sink node of the directed edge to be removed
-            
+
         Raises:
             CassiopeiaTreeError if the tree is not initialized.
         """
@@ -573,7 +669,7 @@ def __add_edge(self, u, v) -> None:
         Args:
             u: The source node of the directed edge to be added
             v: The sink node of the directed edge to be added
-            
+
         Raises:
             CassiopeiaTreeError if the tree is not initialized.
         """
@@ -624,6 +720,9 @@ def set_time(self, node: str, new_time: float) -> None:
                 self.get_time(child) - new_time
             )
 
+        if "distances" in self.__cache:
+            del self.__cache["distances"]
+
     def set_times(self, time_dict: Dict[str, float]) -> None:
         """Sets the time of all nodes in the tree.
 
@@ -658,6 +757,9 @@ def set_times(self, time_dict: Dict[str, float]) -> None:
         for node, time in time_dict.items():
             self.__network.nodes[node]["time"] = time
 
+        if "distances" in self.__cache:
+            del self.__cache["distances"]
+
     def get_time(self, node: str) -> float:
         """Gets the time of a node.
 
@@ -684,7 +786,9 @@ def get_times(self) -> Dict[str, float]:
 
         return dict([(node, self.get_time(node)) for node in self.nodes])
 
-    def __set_branch_length(self, parent: str, child: str, length: float) -> None:
+    def __set_branch_length(
+        self, parent: str, child: str, length: float
+    ) -> None:
         """A private method for setting branch lengths.
 
         A private method for setting branch lengths with no checks. Useful
@@ -698,7 +802,6 @@ def __set_branch_length(self, parent: str, child: str, length: float) -> None:
 
         self.__network[parent][child]["length"] = length
 
-
     def set_branch_length(self, parent: str, child: str, length: float) -> None:
         """Sets the length of a branch.
 
@@ -732,18 +835,23 @@ def set_branch_length(self, parent: str, child: str, length: float) -> None:
                 self.__network.nodes[u]["time"] + self.__network[u][v]["length"]
             )
 
-    def set_branch_lengths(self, branch_length_dict: Dict[Tuple[str, str], float]) -> None:
+        if "distances" in self.__cache:
+            del self.__cache["distances"]
+
+    def set_branch_lengths(
+        self, branch_length_dict: Dict[Tuple[str, str], float]
+    ) -> None:
         """Sets the length of multiple branches on a tree.
 
         Adjusts the branch length of specified parent-child relationships.
         This procedure maintains the consistency with the rest of the times in
         the tree. Namely, by changing branch lengths here, it will change
-        the times of all the nodes in the tree such that the times are 
+        the times of all the nodes in the tree such that the times are
         representative of the new branch lengths.
 
         Args:
             branch_dict: A dictionary of edges to updated branch lengths
-        
+
         Raises:
             CassiopeiaTreeError if the tree has not been initialized.
         """
@@ -762,6 +870,9 @@ def set_branch_lengths(self, branch_length_dict: Dict[Tuple[str, str], float]) -
                 self.__network.nodes[u]["time"] + self.__network[u][v]["length"]
             )
 
+        if "distances" in self.__cache:
+            del self.__cache["distances"]
+
     def get_branch_length(self, parent: str, child: str) -> float:
         """Gets the length of a branch.
 
@@ -776,7 +887,9 @@ def get_branch_length(self, parent: str, child: str) -> float:
 
         return self.__network[parent][child]["length"]
 
-    def set_character_states(self, node: str, states: List[int]) -> None:
+    def set_character_states(
+        self, node: str, states: Union[List[int], List[Tuple[int, ...]]]
+    ) -> None:
         """Sets the character states for a particular node.
 
         Args:
@@ -803,9 +916,21 @@ def set_character_states(self, node: str, states: List[int]) -> None:
         self.__set_character_states(node, states)
 
         if self.is_leaf(node):
-            self.__current_character_matrix.loc[node] = states
-
-    def __set_character_states(self, node: str, states: List[int]) -> None:
+            # Cast entire character matrix to object dtype if we get an ambiguous
+            # character string.
+            if self.is_ambiguous(node):
+                self.__current_character_matrix = (
+                    self.__current_character_matrix.astype(object, copy=False)
+                )
+            # Pass in as numpy array of tuples to bypass the VisibleDeprecationWarning
+            # from creating an ndarray from ragged nested sequences
+            states_arr = np.empty(len(states), dtype=object)
+            states_arr[:] = states
+            self.__current_character_matrix.at[node] = states_arr
+
+    def __set_character_states(
+        self, node: str, states: Union[List[int], List[Tuple[int, ...]]]
+    ) -> None:
         """A private method for setting states.
 
         A private method for setting states of nodes with no checks. Useful
@@ -818,7 +943,9 @@ def __set_character_states(self, node: str, states: List[int]) -> None:
 
         self.__network.nodes[node]["character_states"] = states
 
-    def get_character_states(self, node: str) -> List[int]:
+    def get_character_states(
+        self, node: str
+    ) -> Union[List[int], List[Tuple[int, ...]]]:
         """Gets all the character states for a particular node.
 
         Args:
@@ -834,12 +961,20 @@ def get_character_states(self, node: str) -> List[int]:
 
         return self.__network.nodes[node]["character_states"][:]
 
-    def get_all_ancestors(self, node: str) -> List[str]:
+    def get_all_ancestors(
+        self, node: str, include_node: bool = False
+    ) -> List[str]:
         """Gets all the ancestors of a particular node.
 
+        Nodes that are closest to the given node appear first in the list.
+        The ancestors, including those of all intermediate nodes, are cached.
+
         Args:
             node: Node in the tree
-            
+            include_node: Whether or not to include the node itself in the list
+                of ancestors. If True, the first element of the list is the node
+                itself.
+
         Returns:
             The list of nodes along the path from the root to the node.
         """
@@ -848,13 +983,21 @@ def get_all_ancestors(self, node: str) -> List[str]:
 
         if "ancestors" not in self.__cache:
             self.__cache["ancestors"] = {}
-        
+
+        if self.is_root(node):
+            return []
+
         if node not in self.__cache["ancestors"]:
-            self.__cache["ancestors"][node] = [n
-                for n in nx.ancestors(self.__network, node)
-            ]
+            parent = self.parent(node)
+            self.__cache["ancestors"][node] = [parent] + self.get_all_ancestors(
+                parent
+            )
 
-        return self.__cache["ancestors"][node]
+        # Note that the cache never includes the node itself.
+        ancestors = self.__cache["ancestors"][node]
+        if include_node:
+            ancestors = [node] + ancestors
+        return ancestors
 
     def depth_first_traverse_nodes(
         self, source: Optional[int] = None, postorder: bool = True
@@ -900,7 +1043,7 @@ def depth_first_traverse_edges(
         Raises:
             CassiopeiaTreeError if the tree has not been initialized.
         """
-    
+
         self.__check_network_initialized()
 
         if source is None:
@@ -933,13 +1076,12 @@ def leaves_in_subtree(self, node) -> List[str]:
                     for child in self.children(n):
                         leaves += self.leaves_in_subtree(child)
                     self.__cache["subtree"][n] = leaves
-    
+
         return self.__cache["subtree"][node]
-            
 
-    def get_newick(self, record_branch_lengths = False) -> str:
+    def get_newick(self, record_branch_lengths=False) -> str:
         """Returns newick format of tree.
-        
+
         Args:
             record_branch_lengths: Whether to record branch lengths on the tree
             in the newick string
@@ -952,11 +1094,21 @@ def get_newick(self, record_branch_lengths = False) -> str:
         """
         self.__check_network_initialized()
 
+        # No nodes should have a comma in the name. Otherwise, they will be
+        # displayed as disconnected nodes. For instance, if there is a node with
+        # name "1,2,3", then three separate nodes will be output. We decide to
+        # raise an exception instead of a warning because the returned newick string
+        # will be a WRONG tree.
+        if any("," in node for node in self.nodes):
+            raise CassiopeiaTreeError(
+                "No nodes may have the comma (,) character in its name."
+            )
+
         return utilities.to_newick(self.__network, record_branch_lengths)
 
     def get_tree_topology(self) -> nx.DiGraph:
         """Returns the tree in Networkx format.
-        
+
         Raises:
             CassiopeiaTreeError if the tree has not been initialized.
         """
@@ -1053,28 +1205,141 @@ def relabel_nodes(self, relabel_map: Dict[str, str]) -> None:
     def __register_data_with_tree(self) -> None:
         """Makes the leaf data consistent with the leaves in the tree.
 
-        Removes any leaves from the character matrix, cell metadata, and 
+        Removes any leaves from the character matrix, cell metadata, and
         dissimilarity maps that do not appear in the tree.
+        Additionally, adds any leaves that appear in the tree but not in the
+        character matrix, cell metadata, or dissimilarity map with default values.
+        The default values for each table is as follows:
+        * character matrix: all states are missing values (``missing_state_indicator``)
+        * cell metadata: None
+        * dissimilarity map: ``np.inf`` distance from the leaf to all other leaves
         """
+        leaves_set = set(self.leaves)
         if self.__current_character_matrix is not None:
-            remove_from_charater_matrix = set(self.__current_character_matrix.index) - set(self.leaves)
-            self.__current_character_matrix = self.__current_character_matrix.drop(index = remove_from_charater_matrix)
-                
+            remove_from_charater_matrix = (
+                set(self.__current_character_matrix.index) - leaves_set
+            )
+            self.__current_character_matrix = (
+                self.__current_character_matrix.drop(
+                    index=remove_from_charater_matrix
+                )
+            )
+
+            add_to_character_matrix = leaves_set - set(
+                self.__current_character_matrix.index
+            )
+            for to_add in add_to_character_matrix:
+                # Note that we initialize state each iteration so that each state
+                # is a new object.
+                char_vector = [self.missing_state_indicator] * self.n_character
+                self.__set_character_states(to_add, char_vector)
+                self.__current_character_matrix.loc[to_add] = char_vector
+
         if self.cell_meta is not None:
-            remove_from_cell_meta = set(self.cell_meta.index) - set(self.leaves)
-            self.__cell_meta = self.cell_meta.drop(index = remove_from_cell_meta)
+            remove_from_cell_meta = set(self.cell_meta.index) - leaves_set
+            self.cell_meta = self.cell_meta.drop(index=remove_from_cell_meta)
+
+            add_to_cell_meta = leaves_set - set(self.cell_meta.index)
+            for to_add in add_to_cell_meta:
+                self.cell_meta.loc[to_add] = None
 
         if self.__dissimilarity_map is not None:
-            remove_from_dissimilarity_map = set(self.__dissimilarity_map.index) - set(self.leaves)
-            self.__dissimilarity_map = self.__dissimilarity_map.drop(index = remove_from_dissimilarity_map, columns = remove_from_dissimilarity_map)
+            remove_from_dissimilarity_map = (
+                set(self.__dissimilarity_map.index) - leaves_set
+            )
+            self.__dissimilarity_map = self.__dissimilarity_map.drop(
+                index=remove_from_dissimilarity_map,
+                columns=remove_from_dissimilarity_map,
+            )
+
+            add_to_dissimilarity_map = leaves_set - set(
+                self.__dissimilarity_map.index
+            )
+            for to_add in add_to_dissimilarity_map:
+                self.__dissimilarity_map.loc[to_add] = np.inf
+                self.__dissimilarity_map[to_add] = np.inf
+
+    def add_leaf(
+        self,
+        parent: str,
+        node: str,
+        states: Optional[Union[List[int], List[Tuple[int, ...]]]] = None,
+        time: Optional[float] = None,
+        branch_length: Optional[float] = None,
+        dissimilarity: Optional[Dict[str, float]] = None,
+    ) -> None:
+        """Add a leaf to the given parent node.
+
+        The parent node may NOT also be a leaf, as this makes bookkeeping the
+        character, dissimilarity and cell meta quite complicated. Also, adding a
+        leaf to an existing leaf is probably not a common use case.
+
+        Optional arguments ``states``, ``time``, ``branch_length``,
+        ``dissimilarity`` may be provided to initialize the character
+        string, time, branch length, dissimilarity, and cell meta for this leaf
+        respectively. See :func:`__register_data_with_tree`
+        for default values that are used when any of these are not provided.
+
+        Note that only one of ``time`` and ``branch_length`` may be provided, and
+        when neither are provided, the new leaf is added at the same time as the
+        ``parent``.
+
+        Args:
+            parent: Parent node, to which to connect the leaf
+            node: Name of the leaf to add
+            states: Character states to initialize the new leaf with
+            time: Time to place the new leaf
+            dissimilarity: Indel dissimilarity between the new leaf and all
+                other leaves
+            meta: Cell metadata
+
+        Raises:
+            CassiopeiaTreeError if ``parent`` does not exist or ``node``
+                already exists, if ``parent`` is a leaf, if the tree has not
+                been initialized, or if both ``time`` and ``branch_length``
+                are provided.
+        """
+        self.__check_network_initialized()
+
+        # Make sure the node is not already in the tree and the parent exists
+        if node in self.nodes:
+            raise CassiopeiaTreeError(f"Node {node} already exists.")
+        if parent not in self.nodes:
+            raise CassiopeiaTreeError(f"Node {parent} does not exist.")
+        if parent in self.leaves:
+            raise CassiopeiaTreeError("Can not add a leaf to a leaf.")
+        if time and branch_length:
+            raise CassiopeiaTreeError(
+                "Only one of `time` and `branch_length` may be provided."
+            )
+
+        self.__add_edge(parent, node)
+        if states:
+            self.__set_character_states(node, states)
+            self.__current_character_matrix.loc[node] = states
+        if time:
+            self.set_time(node, time)
+        elif branch_length:
+            self.set_branch_length(parent, node, branch_length)
+        else:
+            # By default, add leaf at the same time as parent.
+            self.set_branch_length(parent, node, 0)
+        if dissimilarity:
+            self.set_dissimilarity(node, dissimilarity)
+
+        # reset cache because we've changed the tree topology
+        self.__cache = {}
+
+        # Update new leaf data with defaults
+        self.__register_data_with_tree()
 
     def remove_leaf_and_prune_lineage(self, node: str) -> None:
         """Removes a leaf from the tree and prunes the lineage.
 
         Removes a leaf and all ancestors of that leaf that are no longer the
         ancestor of any leaves. In the context of a phylogeny, this prunes the
-        lineage of all nodes no longer relevant to observed samples. 
-        Additionally, maintains consistency with the updated tree by removing 
+        lineage of all nodes no longer relevant to observed samples.
+        Additionally, maintains consistency with the updated tree by removing
         the node from all leaf data.
 
         Args:
@@ -1094,9 +1359,8 @@ def remove_leaf_and_prune_lineage(self, node: str) -> None:
         else:
             curr_parent = self.parent(node)
             self.__remove_node(node)
-            while (
-                len(self.children(curr_parent)) < 1
-                and not self.is_root(curr_parent)
+            while len(self.children(curr_parent)) < 1 and not self.is_root(
+                curr_parent
             ):
                 next_parent = self.parent(curr_parent)
                 self.__remove_node(curr_parent)
@@ -1127,7 +1391,9 @@ def collapse_unifurcations(self, source: Optional[int] = None) -> None:
         if source is None:
             source = self.root
 
-        for node in self.depth_first_traverse_nodes(postorder = True, source = source):
+        for node in list(
+            self.depth_first_traverse_nodes(postorder=True, source=source)
+        ):
             if self.is_leaf(node):
                 continue
             elif node == source:
@@ -1161,16 +1427,16 @@ def collapse_mutationless_edges(
         """Collapses mutationless edges in the tree in-place.
 
         Uses the internal node annotations of a tree to collapse edges with no
-        mutations. The introduction of a missing data event is considered a 
+        mutations. The introduction of a missing data event is considered a
         mutation in this context. Either takes the existing character states on
-        the tree or infers the annotations bottom-up from the samples obeying 
+        the tree or infers the annotations bottom-up from the samples obeying
         Camin-Sokal Parsimony. Preserves the times of nodes that are not removed
-        by connecting the parent and children of removed nodes by branchs with 
+        by connecting the parent and children of removed nodes by branchs with
         lengths equal to the total time elapsed from parent to each child.
 
         Args:
             tree: A networkx DiGraph object representing the tree
-            infer_ancestral_characters: Infer the ancestral characters states 
+            infer_ancestral_characters: Infer the ancestral characters states
                 of the tree
 
         Raises:
@@ -1179,13 +1445,15 @@ def collapse_mutationless_edges(
         if infer_ancestral_characters:
             self.reconstruct_ancestral_characters()
 
-        for n in self.depth_first_traverse_nodes(postorder = True):
+        for n in list(self.depth_first_traverse_nodes(postorder=True)):
             if self.is_leaf(n):
                 continue
             for child in self.children(n):
                 if not self.is_leaf(child):
                     t = self.get_branch_length(n, child)
-                    if self.get_character_states(n) == self.get_character_states(child):
+                    if self.get_character_states(
+                        n
+                    ) == self.get_character_states(child):
                         for grandchild in self.children(child):
                             t_ = self.get_branch_length(child, grandchild)
                             self.__add_edge(n, grandchild)
@@ -1210,15 +1478,15 @@ def set_dissimilarity_map(self, dissimilarity_map: pd.DataFrame) -> None:
             dissimilarity_map: Dissimilarity map relating all N x N distances
                 between leaves.
         """
-        character_matrix = self.__original_character_matrix
+        character_matrix = self.__current_character_matrix
         if character_matrix is not None:
 
-            if character_matrix.shape[0] != dissimilarity_map.shape[
-                0
-            ] or collections.Counter(
-                character_matrix.index
-            ) != collections.Counter(
-                dissimilarity_map.index
+            if (
+                character_matrix.shape[0] != dissimilarity_map.shape[0]
+                or collections.Counter(character_matrix.index)
+                != collections.Counter(dissimilarity_map.index)
+                or collections.Counter(character_matrix.index)
+                != collections.Counter(dissimilarity_map.columns)
             ):
                 warnings.warn(
                     "The samples in the existing character matrix and "
@@ -1228,6 +1496,60 @@ def set_dissimilarity_map(self, dissimilarity_map: pd.DataFrame) -> None:
 
         self.__dissimilarity_map = dissimilarity_map.copy()
 
+    def get_dissimilarity(self, node: str) -> Dict[str, float]:
+        """Get the dissimilarity of a single leaf to all other leaves.
+
+        Args:
+            node: Name of the leaf node to get the dissimilarity to/from
+
+        Raises:
+            CassiopeiaTreeError if the dissimilarity map is not defined or the
+                provided node is not in the dissimilarity map
+        """
+        if self.__dissimilarity_map is None:
+            raise CassiopeiaTreeError("No dissimilarity map is defined.")
+        if node not in self.__dissimilarity_map.index:
+            raise CassiopeiaTreeError(
+                f"`{node}` is not in the dissimilarity map."
+            )
+        return self.__dissimilarity_map.loc[node].to_dict()
+
+    def set_dissimilarity(
+        self, node: str, dissimilarity: Dict[str, float]
+    ) -> None:
+        """Set the dissimilarity of a single leaf.
+
+        This function may be used only when a dissimilarity map already exists.
+
+        Args:
+            node: Name of the leaf node to set the dissimilarity to/from
+            dissimilarity: Dictionary containing other leaf names as keys and
+                dissimilarities as values
+
+        Raises:
+            CassiopeiaTreeError if the dissimilarity map does not exist, if
+                the ``dissimilarity`` dictionary does not contain all other leaves,
+        """
+        if self.__dissimilarity_map is None:
+            raise CassiopeiaTreeError(
+                "This function may only be used with an existing dissimilarity map. "
+                "Use `compute_dissimilarity_map` or `set_dissimilarity_map` to set "
+                "the dissimilarity map first."
+            )
+        if set(dissimilarity.keys()) - {node} != set(
+            self.__dissimilarity_map.index
+        ) - {node}:
+            raise CassiopeiaTreeError(
+                "Provided dissimilarity dictionary does not provide dissimilarities "
+                "to all other leaves."
+            )
+
+        self.__dissimilarity_map.loc[node] = 0
+        self.__dissimilarity_map[node] = 0
+        for other, dissim in dissimilarity.items():
+            self.__dissimilarity_map.loc[node, other] = dissim
+            self.__dissimilarity_map.loc[other, node] = dissim
+
     def compute_dissimilarity_map(
         self,
         dissimilarity_function: Optional[
@@ -1243,6 +1565,12 @@ def compute_dissimilarity_map(
         will be computed over the samples in the character matrix. Populates
         the dissimilarity_map attribute in the object.
 
+        If any of the leaves have ambiguous character states (detected by checking
+        if any of the states are lists), then
+        :func:`cassiopeia.solver.dissimilarity_functions.cluster_dissimilarity`
+        is called instead, with the provided ``dissimilarity_function`` as
+        the first argument.
+
         Args:
             dissimilarity_function: A function that will take in two character
                 vectors and priors and produce a dissimilarity.
@@ -1290,7 +1618,7 @@ def compute_dissimilarity_map(
 
     def set_attribute(self, node: str, attribute_name: str, value: Any) -> None:
         """Sets an attribute in the tree.
-        
+
         Args:
             node: Node name
             attribute_name: Name for the new attribute
@@ -1305,11 +1633,11 @@ def set_attribute(self, node: str, attribute_name: str, value: Any) -> None:
 
     def get_attribute(self, node: str, attribute_name: str) -> Any:
         """Retrieves the value of an attribute for a node.
-        
+
         Args:
             node: Node name
             attribute_name: Name of the attribute.
-        
+
         Returns:
             The value of the attribute for that node.
         Raises:
@@ -1320,8 +1648,9 @@ def get_attribute(self, node: str, attribute_name: str) -> Any:
         try:
             return self.__network.nodes[node][attribute_name]
         except KeyError:
-            raise CassiopeiaTreeError(f"Attribute {attribute_name} not " 
-                                    "detected for this node.")
+            raise CassiopeiaTreeError(
+                f"Attribute {attribute_name} not " "detected for this node."
+            )
 
     def filter_nodes(self, condition: Callable[[str], bool]) -> List[str]:
 
@@ -1334,14 +1663,130 @@ def filter_nodes(self, condition: Callable[[str], bool]) -> List[str]:
 
         return _filter
 
-    def find_lcas_of_pairs(self, pairs: Union[Iterator[str], List[str]]) -> Iterator[Tuple[Tuple[str, str], str]]:
-        """Finds LCAs of all pairs.
+    def find_lcas_of_pairs(
+        self,
+        pairs: Optional[
+            Union[Iterator[Tuple[str, str]], List[Tuple[str, str]]]
+        ] = None,
+    ) -> Iterator[Tuple[Tuple[str, str], str]]:
+        """Finds LCAs of all (provided) pairs.
 
         Args:
-            pairs: Pairs of nodes for which to find LCAs
+            pairs: Pairs of nodes for which to find LCAs. If not provided, LCAs of
+                all pairs are computed. Defaults to None.
 
         Returns:
             A generator of ((u, v), LCA) tuples.
+
+        Raises:
+            CassiopeiaTreeError if the tree has not been initialized.
+        """
+        self.__check_network_initialized()
+        return nx.tree_all_pairs_lowest_common_ancestor(
+            self.__network, root=self.root, pairs=pairs
+        )
+
+    def find_lca(self, *nodes: str) -> str:
+        """Finds the LCA of all provided nodes.
+
+        Internally, this function calls :func:`get_all_ancestors` for each provided
+        node and finds the deepest node that is shared.
+
+        Args:
+            *nodes: Nodes for which to find the LCA. At least two must be provided.
+
+        Returns:
+            The LCA node
+
+        Raises:
+            CassiopeiaTreeError if less than two nodes were provided or if the
+            tree has not been initialized.
+        """
+        self.__check_network_initialized()
+        nodes = set(nodes)
+        if len(nodes) < 2:
+            raise CassiopeiaTreeError(
+                "At least two distinct nodes must be provided"
+            )
+
+        # Reversing get_all_ancestors gives a list of nodes from the root to each
+        # node. Since get_all_ancestors doesn't include the node itself, we add
+        # those manually.
+        all_ancestors = [
+            reversed(self.get_all_ancestors(node, include_node=True))
+            for node in nodes
+        ]
+        last_ancestor = self.root
+        for ancestors in zip(*all_ancestors):
+            if len(set(ancestors)) > 1:
+                break
+            last_ancestor = ancestors[0]
+        return last_ancestor
+
+    def get_distance(self, node1: str, node2: str) -> float:
+        """Compute the branch distance between two nodes.
+
+        Internally, this function converts the tree to an undirected graph and
+        finds the shortest path by calling :func:`nx.shortest_path_length`.
+
+        Args:
+            node1: First node
+            node2: Second node
+
+        Returns:
+            The branch distance between the two nodes
+
+        Raises:
+            CassiopeiaTreeError if the tree has not been initialized.
+        """
+        self.__check_network_initialized()
+
+        key = frozenset({node1, node2})
+        if (
+            "distances" not in self.__cache
+            or key not in self.__cache["distances"]
+        ):
+            distances = self.__cache.setdefault("distances", {})
+            undirected_network = self.__network.to_undirected()
+            distances[key] = nx.shortest_path_length(
+                undirected_network, source=node1, target=node2, weight="length"
+            )
+
+        return self.__cache["distances"][key]
+
+    def get_distances(
+        self, node: str, leaves_only: bool = False
+    ) -> Dict[str, float]:
+        """Compute the distances between the given node and every other node.
+
+        Internally, this function converts the tree to an undirected graph and
+        finds the shortest path by calling :func:`nx.shortest_path_length`.
+
+        Args:
+            node: Node from which to compute distance to all other nodes
+            leaves_only: Calculate distances to leaves only and not internal nodes.
+                Defaults to False.
+
+        Returns:
+            Dictionary of distances
+
+        Raises:
+            CassiopeiaTreeError if the tree has not been initialized.
         """
         self.__check_network_initialized()
-        return nx.tree_all_pairs_lowest_common_ancestor(self.__network, root=self.root, pairs=pairs)
+
+        if "distances" not in self.__cache or any(
+            frozenset({node, _node}) not in self.__cache["distances"]
+            for _node in self.nodes
+        ):
+            distances = self.__cache.setdefault("distances", {})
+            undirected_network = self.__network.to_undirected()
+            for other_node, distance in nx.shortest_path_length(
+                undirected_network, source=node, weight="length"
+            ).items():
+                distances[frozenset({node, other_node})] = distance
+
+        return {
+            _node: self.__cache["distances"][frozenset({node, _node})]
+            for _node in (self.leaves if leaves_only else self.nodes)
+        }
diff --git a/cassiopeia/data/utilities.py b/cassiopeia/data/utilities.py
old mode 100644
new mode 100755
index 14ec0fba..5b3d8d80
--- a/cassiopeia/data/utilities.py
+++ b/cassiopeia/data/utilities.py
@@ -1,7 +1,9 @@
 """
 General utilities for the datasets encountered in Cassiopeia.
 """
-from typing import Callable, Dict, List, Optional, Tuple
+import collections
+from typing import Callable, Dict, List, Optional, Tuple, Union
+import warnings
 
 import ete3
 import networkx as nx
@@ -12,9 +14,16 @@
 
 from cassiopeia.preprocess import utilities as preprocessing_utilities
 
+# Can't import from CassiopeiaTree.py due to circular imports
+class CassiopeiaTreeWarning(UserWarning):
+    """A Warning for the CassiopeiaTree class."""
+
+    pass
+
 
 def get_lca_characters(
-    vecs: List[List[int]], missing_state_indicator: int
+    vecs: List[Union[List[int], List[Tuple[int, ...]]]],
+    missing_state_indicator: int,
 ) -> List[int]:
     """Builds the character vector of the LCA of a list of character vectors,
     obeying Camin-Sokal Parsimony.
@@ -36,7 +45,12 @@ def get_lca_characters(
         assert len(i) == k
     lca_vec = [0] * len(vecs[0])
     for i in range(k):
-        chars = set([vec[i] for vec in vecs])
+        chars = set()
+        for vec in vecs:
+            if is_ambiguous_state(vec[i]):
+                chars = chars.union(vec[i])
+            else:
+                chars.add(vec[i])
         if len(chars) == 1:
             lca_vec[i] = list(chars)[0]
         else:
@@ -105,11 +119,11 @@ def to_newick(tree: nx.DiGraph, record_branch_lengths: bool = False) -> str:
     def _to_newick_str(g, node):
         is_leaf = g.out_degree(node) == 0
         weight_string = ""
-        
+
         if record_branch_lengths and g.in_degree(node) > 0:
             parent = list(g.predecessors(node))[0]
             weight_string = ":" + str(g[parent][node]["length"])
-            
+
         _name = str(node)
         return (
             "%s" % (_name,) + weight_string
@@ -119,7 +133,8 @@ def _to_newick_str(g, node):
                 + ",".join(
                     _to_newick_str(g, child) for child in g.successors(node)
                 )
-                + ")" + weight_string
+                + ")"
+                + weight_string
             )
         )
 
@@ -148,8 +163,20 @@ def compute_dissimilarity_map(
     Returns:
         A dissimilarity mapping as a flattened array.
     """
-
-    nb_dissimilarity = numba.jit(dissimilarity_function, nopython=True)
+    # Try to numbaize the dissimilarity function, but fallback to python
+    numbaize = True
+    try:
+        dissimilarity_func = numba.jit(dissimilarity_function, nopython=True)
+    # When cluster_dissimilarity is used, the dissimilarity_function is wrapped
+    # in a partial, which raises a TypeError when trying to numbaize.
+    except TypeError:
+        warnings.warn(
+            "Failed to numbaize dissimilarity function. "
+            "Falling back to Python.",
+            CassiopeiaTreeWarning,
+        )
+        numbaize = False
+        dissimilarity_func = dissimilarity_function
 
     nb_weights = numba.typed.Dict.empty(
         numba.types.int64,
@@ -165,23 +192,30 @@ def compute_dissimilarity_map(
                 nb_char_weights[state] = prior
             nb_weights[k] = nb_char_weights
 
-    @numba.jit(nopython=True)
     def _compute_dissimilarity_map(cm, C, missing_state_indicator, nb_weights):
 
-        dm = np.zeros(C * (C - 1) // 2, dtype=numba.float64)
+        dm = np.zeros(C * (C - 1) // 2, dtype=np.float64)
         k = 0
         for i in range(C - 1):
             for j in range(i + 1, C):
 
                 s1 = cm[i, :]
                 s2 = cm[j, :]
-                dm[k] = nb_dissimilarity(
+                dm[k] = dissimilarity_func(
                     s1, s2, missing_state_indicator, nb_weights
                 )
                 k += 1
 
         return dm
 
+    # Don't numbaize _compute_dissimilarity_map when we failed to numbaize
+    # the dissimilarity function. Otherwise, if we try to call a Python
+    # function from a numbaized (in object mode) a LOT of warnings are raised.
+    if numbaize:
+        _compute_dissimilarity_map = numba.jit(
+            _compute_dissimilarity_map, nopython=True
+        )
+
     return _compute_dissimilarity_map(
         cm, C, missing_state_indicator, nb_weights
     )
@@ -281,8 +315,10 @@ def sample_bootstrap_allele_tables(
             allele_table
         )
 
-    lineage_profile = preprocessing_utilities.convert_alleletable_to_lineage_profile(
-        allele_table, cut_sites
+    lineage_profile = (
+        preprocessing_utilities.convert_alleletable_to_lineage_profile(
+            allele_table, cut_sites
+        )
     )
 
     intbcs = allele_table["intBC"].unique()
@@ -324,3 +360,37 @@ def sample_bootstrap_allele_tables(
         )
 
     return bootstrap_samples
+
+
+def is_ambiguous_state(state: Union[int, Tuple[int, ...]]) -> bool:
+    """Determine whether the provided state is ambiguous.
+
+    Note that this function operates on a single (indel) state.
+
+    Args:
+        state: Single, possibly ambiguous, character state
+
+    Returns:
+        True if the state is ambiguous, False otherwise.
+    """
+    return isinstance(state, tuple)
+
+
+def resolve_most_abundant(state: Tuple[int, ...]) -> int:
+    """Resolve an ambiguous character by selecting the most abundant.
+
+    This function is designed to be used with
+    :func:`CassiopeiaTree.resolve_ambiguous_characters`. It resolves an ambiguous
+    character, represented as a tuple of integers, by selecting the most abundant,
+    where ties are resolved randomly.
+
+    Args:
+        state: Ambiguous state as a tuple of integers
+
+    Returns:
+        Selected state as a single integer
+    """
+    most_common = collections.Counter(state).most_common()
+    return np.random.choice(
+        [state for state, count in most_common if count == most_common[0][1]]
+    )
diff --git a/cassiopeia/plotting/itol_utilities.py b/cassiopeia/plotting/itol_utilities.py
index f8469d07..d68ee2c4 100644
--- a/cassiopeia/plotting/itol_utilities.py
+++ b/cassiopeia/plotting/itol_utilities.py
@@ -23,8 +23,7 @@
 
 
 class iTOLError(Exception):
-    """Raises errors related to iTOL plotting
-    """
+    """Raises errors related to iTOL plotting"""
 
     pass
 
@@ -47,7 +46,7 @@ def upload_and_export_itol(
     random_state: Optional[np.random.RandomState] = None,
     user_dataset_files: Optional[List[str]] = None,
     verbose: bool = True,
-    **kwargs
+    **kwargs,
 ):
     """Uploads a tree to iTOL and exports it.
 
@@ -110,7 +109,7 @@ def upload_and_export_itol(
     # create temporary directory for storing files we'll upload to iTOL
     temporary_directory = tempfile.mkdtemp()
 
-    if (api_key is None or project_name is None):
+    if api_key is None or project_name is None:
         if os.path.exists(os.path.expanduser(itol_config)):
 
             config = configparser.ConfigParser()
@@ -131,7 +130,7 @@ def upload_and_export_itol(
             )
 
     with open(os.path.join(temporary_directory, "tree_to_plot.tree"), "w") as f:
-        f.write(cassiopeia_tree.get_newick(record_branch_lengths = True))
+        f.write(cassiopeia_tree.get_newick(record_branch_lengths=True))
 
     file_format = export_filepath.split("/")[-1].split(".")[-1]
 
diff --git a/cassiopeia/preprocess/UMI_utils.py b/cassiopeia/preprocess/UMI_utils.py
old mode 100644
new mode 100755
index b4c40386..d59a7063
--- a/cassiopeia/preprocess/UMI_utils.py
+++ b/cassiopeia/preprocess/UMI_utils.py
@@ -71,7 +71,34 @@
 ####################Utils for Collapsing UMIs#####################
 
 
-def sort_cellranger_bam(
+def detect_cell_bc_tag(bam_fp: str) -> str:
+    """Detect what BAM tag corresponds to the cell barcode.
+
+    Use this function to detect either ``BAM_CONSTANTS["RAW_CELL_BC_TAG"]`` or
+    ``BAM_CONSTANTS["CELL_BC_TAG"]`` to use as the cell barcodes for downstream
+    processing. This is because for some chemistries it may not be possible to
+    perform barcode correction (i.e. chemistries without a whitelist) and
+    therefore the pipeline needs to use the raw barcode sequences.
+
+    Args:
+        bam_fp: Path to the BAM file
+
+    Returns:
+        The BAM tag to use as the barcode
+    """
+    raw_tag = BAM_CONSTANTS["RAW_CELL_BC_TAG"]
+    corrected_tag = BAM_CONSTANTS["CELL_BC_TAG"]
+
+    # Iterating through reads in a BAM file is fast, so it's okay to check all
+    # reads
+    with pysam.AlignmentFile(bam_fp, "rb", check_sq=False) as f:
+        for read in f:
+            if read.has_tag(corrected_tag):
+                return corrected_tag
+    return raw_tag
+
+
+def sort_bam(
     bam_fp: str,
     sorted_fn: str,
     sort_key: Callable[[pysam.AlignedSegment], str] = sort_key,
@@ -144,6 +171,8 @@ def form_collapsed_clusters(
     sorted_fn: Callable[[pysam.AlignedSegment], str],
     max_hq_mismatches: int,
     max_indels: int,
+    cell_key: Callable[[pysam.AlignedSegment], str] = cell_key,
+    UMI_key: Callable[[pysam.AlignedSegment], str] = UMI_key,
 ):
     """Aggregates together aligned segments (reads) that share UMIs if their
     sequences are close.
@@ -168,6 +197,8 @@ def form_collapsed_clusters(
             collapsed.
         max_indels: A threshold specifying the maximum number of differing indels
             allowed between the sequences of 2 aligned segments to be collapsed.
+        cell_key: A function that takes an alignment and returns the cell barcode
+        UMI_key: A function that takes an alignment and returns the UMI sequence
 
     None:
         Saves the sorted bam to file
@@ -501,7 +532,6 @@ def merge_annotated_clusters(
 
 def correct_umis_in_group(
     cell_group: pd.DataFrame,
-    sampleID: str,
     max_umi_distance: int = 2,
     verbose=False,
 ) -> Tuple[pd.DataFrame, int, int]:
@@ -521,7 +551,6 @@ def correct_umis_in_group(
 
     Args:
         input_df: Input DataFrame of alignments.
-        _id: Identification of sample.
         max_umi_distance: Maximum Hamming distance between UMIs
             for error correction.
 
diff --git a/cassiopeia/preprocess/cassiopeia_preprocess.py b/cassiopeia/preprocess/cassiopeia_preprocess.py
old mode 100644
new mode 100755
index 7f7fde95..1b5b8923
--- a/cassiopeia/preprocess/cassiopeia_preprocess.py
+++ b/cassiopeia/preprocess/cassiopeia_preprocess.py
@@ -1,7 +1,7 @@
 """
 Main logic behind Cassiopeia-preprocess.
 
-This file stores the main entry point for Cassiopeia-preprocess, and makes 
+This file stores the main entry point for Cassiopeia-preprocess, and makes
 heavy use of the high level functionality in
 cassiopeia.preprocess.pipeline. Here, we assume that the user
 has already run CellRanger Count, or some equivalent, to obtain a BAM file that
@@ -20,11 +20,13 @@
 from cassiopeia.preprocess import pipeline, setup_utilities
 
 STAGES = {
+    "convert": pipeline.convert_fastqs_to_unmapped_bam,
+    "error_correct_barcodes": pipeline.error_correct_barcodes,
     "collapse": pipeline.collapse_umis,
     "resolve": pipeline.resolve_umi_sequence,
     "align": pipeline.align_sequences,
     "call_alleles": pipeline.call_alleles,
-    "error_correct": pipeline.error_correct_umis,
+    "error_correct_umis": pipeline.error_correct_umis,
     "filter_molecule_table": pipeline.filter_molecule_table,
     "call_lineages": pipeline.call_lineage_groups,
 }
@@ -47,7 +49,7 @@ def main():
 
     # pull out general parameters
     output_directory = pipeline_parameters["general"]["output_directory"]
-    data_filepath = pipeline_parameters["general"]["input_file"]
+    data_filepaths = pipeline_parameters["general"]["input_files"]
     entry_point = pipeline_parameters["general"]["entry"]
     exit_point = pipeline_parameters["general"]["exit"]
 
@@ -58,22 +60,46 @@ def main():
     pipeline_stages = setup_utilities.create_pipeline(
         entry_point, exit_point, STAGES
     )
-    if entry_point == "collapse":
-        data = data_filepath
+    if entry_point == "convert":
+        data = data_filepaths
     else:
-        data = pd.read_csv(data_filepath, sep="\t")
+        if len(data_filepaths) != 1:
+            raise pipeline.PreprocessError(
+                "`input_files` must contain exactly one input file for pipeline "
+                f"stage `{entry_point}`"
+            )
 
-    data_file_stem = ".".join(data_filepath.split("/")[-1].split(".")[:-1])
+        if entry_point in ("error_correct_barcodes", "collapse"):
+            data = data_filepaths[0]
+        else:
+            data = pd.read_csv(data_filepaths[0], sep="\t")
+
+    data_file_stem = os.path.splitext(os.path.basename(data_filepaths[0]))[0]
 
     # ---------------------- Run Pipeline ---------------------- #
     for stage in pipeline_stages:
+        # Skip barcode correction if whitelist_fp was not provided
+        if (
+            stage == "error_correct_barcodes"
+            and not pipeline_parameters[stage]["whitelist_fp"]
+        ):
+            logging.warning(
+                "Skipping barcode error correction because no whitelist was "
+                "provided in the configuration."
+            )
+            continue
 
         procedure = STAGES[stage]
         data = procedure(data, **pipeline_parameters[stage])
-        data.to_csv(
-            os.path.join(output_directory, data_file_stem + f".{stage}.txt"),
-            sep="\t",
-        )
+
+        # Write to CSV only if it is a pandas dataframe
+        if isinstance(data, pd.DataFrame):
+            data.to_csv(
+                os.path.join(
+                    output_directory, data_file_stem + f".{stage}.txt"
+                ),
+                sep="\t",
+            )
 
 
 if __name__ == "__main__":
diff --git a/cassiopeia/preprocess/constants.py b/cassiopeia/preprocess/constants.py
old mode 100644
new mode 100755
index 0a9071c9..405972ed
--- a/cassiopeia/preprocess/constants.py
+++ b/cassiopeia/preprocess/constants.py
@@ -3,8 +3,11 @@
 """
 
 BAM_CONSTANTS = {
+    "RAW_CELL_BC_TAG": "CR",
+    "RAW_CELL_BC_QUALITY_TAG": "CY",
     "CELL_BC_TAG": "CB",
     "UMI_TAG": "UR",
+    "UMI_QUALITY_TAG": "UY",
     "NUM_READS_TAG": "ZR",
     "CLUSTER_ID_TAG": "ZC",
     "N_Q": 2,
@@ -12,6 +15,29 @@
     "LOW_Q": 10,
 }
 
+SINGLE_CELL_BAM_TAGS = {
+    "umi": (BAM_CONSTANTS["UMI_TAG"], BAM_CONSTANTS["UMI_QUALITY_TAG"]),
+    "cell_barcode": (
+        BAM_CONSTANTS["RAW_CELL_BC_TAG"],
+        BAM_CONSTANTS["RAW_CELL_BC_QUALITY_TAG"],
+    ),
+}
+SPATIAL_BAM_TAGS = {
+    "umi": (BAM_CONSTANTS["UMI_TAG"], BAM_CONSTANTS["UMI_QUALITY_TAG"]),
+    "spot_barcode": (
+        BAM_CONSTANTS["RAW_CELL_BC_TAG"],
+        BAM_CONSTANTS["RAW_CELL_BC_QUALITY_TAG"],
+    ),
+}
+CHEMISTRY_BAM_TAGS = {
+    "dropseq": SINGLE_CELL_BAM_TAGS,
+    "10xv2": SINGLE_CELL_BAM_TAGS,
+    "10xv3": SINGLE_CELL_BAM_TAGS,
+    "indropsv3": SINGLE_CELL_BAM_TAGS,
+    "slideseq2": SPATIAL_BAM_TAGS,
+}
+
+
 DNA_SUBSTITUTION_MATRIX = {
     "A": {"A": 5, "T": -4, "C": -4, "G": -4, "Z": 0, "N": 0},
     "T": {"A": -4, "T": 5, "C": -4, "G": -4, "Z": 0, "N": 0},
@@ -22,8 +48,10 @@
 }
 
 DEFAULT_PIPELINE_PARAMETERS = {
-    "general": {"entry": "'collapse'", "exit": "'call_lineages'"},
-    "collapse": {"max_hq_mismatches": 3, "max_indels": 2, "force_sort": True},
+    "general": {"entry": "'convert'", "exit": "'call_lineages'"},
+    "convert": {},
+    "error_correct_barcodes": {},
+    "collapse": {"max_hq_mismatches": 3, "max_indels": 2},
     "resolve": {
         "min_avg_reads_per_umi": 2.0,
         "min_umi_per_cell": 10,
@@ -37,8 +65,7 @@
         "context": True,
         "context_size": 5,
     },
-    "error_correct": {
-        "_id": "'batch1'",
+    "error_correct_umis": {
         "max_umi_distance": 2,
         "verbose": False,
     },
diff --git a/cassiopeia/preprocess/pipeline.py b/cassiopeia/preprocess/pipeline.py
index e02b0f09..fb7388de 100755
--- a/cassiopeia/preprocess/pipeline.py
+++ b/cassiopeia/preprocess/pipeline.py
@@ -4,21 +4,22 @@
 is mainly invoked by cassiopeia_preprocess.py.
 """
 
+from functools import partial
+import logging
 import os
+from pathlib import Path
 import time
-
 from typing import List, Optional, Tuple
 
 from Bio import SeqIO
-from functools import partial
-import logging
 import matplotlib.pyplot as plt
+import ngs_tools as ngs
 import numpy as np
 import pandas as pd
+import pysam
 from skbio import alignment
-
-from pathlib import Path
 from tqdm.auto import tqdm
+from typing_extensions import Literal
 
 from cassiopeia.preprocess import alignment_utilities
 from cassiopeia.preprocess import constants
@@ -34,6 +35,137 @@
 progress = tqdm
 
 
+class PreprocessError(Exception):
+    pass
+
+
+def convert_fastqs_to_unmapped_bam(
+    fastq_fps: List[str],
+    chemistry: Literal["dropseq", "10xv2", "10xv3", "indropsv3", "slideseq2"],
+    output_directory: str,
+    name: Optional[str] = None,
+    n_threads: int = 1,
+) -> str:
+    """Converts FASTQs into an unmapped BAM based on a chemistry.
+
+    This function converts a set of FASTQs into an unmapped BAM with appropriate
+    BAM tags.
+
+    Args:
+        fastq_fps: List of paths to FASTQ files. Usually, this argument contains
+            two FASTQs, where the first contains the barcode and UMI sequences
+            and the second contains cDNA. The FASTQs may be gzipped.
+        chemistry: Sample-prep/sequencing chemistry used. The following
+            chemistries are supported:
+            * dropseq: Droplet-based scRNA-seq chemistry described in
+                Macosco et al. 2015
+            * 10xv2: 10x Genomics 3' version 2
+            * 10xv3: 10x Genomics 3' version 3
+            * indropsv3: inDrops version 3 by Zilionis et al. 2017
+            * slideseq2: Slide-seq version 2
+        output_directory: The output directory where the unmapped BAM will be
+            written to. This directory must exist prior to calling this function.
+        name: Name of the reads in the FASTQs. This name is set as the read group
+            name for the reads in the output BAM, as well as the filename prefix
+            of the output BAM. If not provided, a short random UUID is used as
+            the read group name, but not as the filename prefix of the BAM.
+        n_threads: Number of threads to use. Defaults to 1.
+
+    Returns:
+        Path to written BAM
+
+    Raises:
+        PreprocessError if the provided chemistry does not exist.
+    """
+    if chemistry not in constants.CHEMISTRY_BAM_TAGS:
+        raise PreprocessError(f"Unknown chemistry {chemistry}")
+
+    logging.info("Converting FASTQs to unmapped BAM...")
+    t0 = time.time()
+
+    tag_map = constants.CHEMISTRY_BAM_TAGS[chemistry]
+    bam_fp = os.path.join(
+        output_directory, f"{name}_unmapped.bam" if name else "unmapped.bam"
+    )
+    ngs.fastq.fastqs_to_bam_with_chemistry(
+        fastq_fps,
+        ngs.chemistry.get_chemistry(chemistry),
+        tag_map,
+        bam_fp,
+        name=name,
+        show_progress=True,
+        n_threads=n_threads,
+    )
+    logging.info(f"Finished writing unmapped BAM in {time.time() - t0} s.")
+    return bam_fp
+
+
+def error_correct_barcodes(
+    bam_fp: str, output_directory: str, whitelist_fp: str, n_threads: int = 1
+) -> str:
+    """Error-correct barcodes in the input BAM.
+
+    The barcode correction procedure used in Cell Ranger by 10X Genomics is used.
+    https://kb.10xgenomics.com/hc/en-us/articles/115003822406-How-does-Cell-Ranger-correct-barcode-sequencing-errors
+
+    Args:
+        bam_fp: Input BAM filepath containing raw barcodes
+        output_directory: The output directory where the corrected BAM will be
+            written to. This directory must exist prior to calling this function.
+        whitelist_fp: Path to plaintext file containing barcode whitelist, one
+            barcode per line.
+        n_threads: Number of threads to use. Defaults to 1.
+
+    Todo:
+        Currently, the user must provide their own whitelist, and Cassiopeia
+        does not use any of the whitelists provided by the ngs-tools library.
+        At some point, we should update the pipeline so that if ngs-tools
+        provides a pre-packaged whitelists, it uses that for those chemistries.
+
+    Returns:
+        Path to corrected BAM
+    """
+    logging.info("Correcting barcodes to whitelist...")
+    t0 = time.time()
+
+    # Read whitelist
+    with open(whitelist_fp, "r") as f:
+        whitelist = [line.strip() for line in f if not line.isspace()]
+
+    # Extract all raw barcodes and their qualities
+    barcodes = []
+    qualities = []
+    with pysam.AlignmentFile(
+        bam_fp, "rb", check_sq=False, threads=n_threads
+    ) as f:
+        for read in f:
+            barcodes.append(read.get_tag(BAM_CONSTANTS["RAW_CELL_BC_TAG"]))
+            qualities.append(
+                read.get_tag(BAM_CONSTANTS["RAW_CELL_BC_QUALITY_TAG"])
+            )
+
+    # Correct
+    corrections = ngs.sequence.correct_sequences_to_whitelist(
+        barcodes, qualities, whitelist, show_progress=True, n_threads=n_threads
+    )
+
+    # Write corrected BAM
+    prefix, ext = os.path.splitext(os.path.basename(bam_fp))
+    corrected_fp = os.path.join(output_directory, f"{prefix}_corrected{ext}")
+    with pysam.AlignmentFile(
+        bam_fp, "rb", check_sq=False, threads=n_threads
+    ) as f_in:
+        with pysam.AlignmentFile(
+            corrected_fp, "wb", template=f_in, threads=n_threads
+        ) as f_out:
+            for i, read in enumerate(f_in):
+                if corrections[i]:
+                    read.set_tag(BAM_CONSTANTS["CELL_BC_TAG"], corrections[i])
+                f_out.write(read)
+    logging.info(f"Finished correcting barcodes in {time.time() - t0} s.")
+    return corrected_fp
+
+
 def collapse_umis(
     bam_fp: str,
     output_directory: str,
@@ -81,9 +213,18 @@ def collapse_umis(
         + "_sorted.bam"
     )
 
+    cell_bc_tag = UMI_utils.detect_cell_bc_tag(bam_fp)
+    logging.info(f"Using BAM tag `{cell_bc_tag}` as cell barcodes")
+
     if not sorted_file_name.exists() and not skip_existing:
-        max_read_length, total_reads_out = UMI_utils.sort_cellranger_bam(
-            bam_fp, str(sorted_file_name)
+        max_read_length, total_reads_out = UMI_utils.sort_bam(
+            bam_fp,
+            str(sorted_file_name),
+            sort_key=lambda al: (
+                al.get_tag(cell_bc_tag),
+                al.get_tag(BAM_CONSTANTS["UMI_TAG"]),
+            ),
+            filter_func=lambda al: al.has_tag(cell_bc_tag),
         )
         logging.info("Sorted bam directory saved to " + str(sorted_file_name))
         logging.info("Max read length of " + str(max_read_length))
@@ -92,7 +233,10 @@ def collapse_umis(
     collapsed_file_name = sorted_file_name.with_suffix(".collapsed.bam")
     if not collapsed_file_name.exists() and not skip_existing:
         UMI_utils.form_collapsed_clusters(
-            str(sorted_file_name), max_hq_mismatches, max_indels
+            str(sorted_file_name),
+            max_hq_mismatches,
+            max_indels,
+            cell_key=lambda al: al.get_tag(cell_bc_tag),
         )
 
     logging.info(f"Finished collapsing UMI sequences in {time.time() - t0} s.")
@@ -426,7 +570,6 @@ def call_alleles(
 
 def error_correct_umis(
     input_df: pd.DataFrame,
-    _id: str,
     max_umi_distance: int = 2,
     verbose: bool = False,
 ) -> pd.DataFrame:
@@ -438,7 +581,6 @@ def error_correct_umis(
 
     Args:
         input_df: Input DataFrame of alignments.
-        _id: Identification of sample.
         max_umi_distance: The threshold specifying the Maximum Hamming distance
             between UMIs for one to be corrected to another.
         verbose: Indicates whether to log every UMI correction.
@@ -489,7 +631,7 @@ def error_correct_umis(
         if verbose:
             logging.info(f"cellBC: {cellBC}, intBC: {intBC}")
         allele_group, num_corr, tot = UMI_utils.correct_umis_in_group(
-            allele_group, _id, max_umi_distance
+            allele_group, max_umi_distance
         )
         num_corrected += num_corr
         total += tot
diff --git a/cassiopeia/preprocess/setup_utilities.py b/cassiopeia/preprocess/setup_utilities.py
old mode 100644
new mode 100755
index bafa4e41..b6d587ef
--- a/cassiopeia/preprocess/setup_utilities.py
+++ b/cassiopeia/preprocess/setup_utilities.py
@@ -67,16 +67,22 @@ def parse_config(config_string: str) -> Dict[str, Dict[str, Any]]:
     minimum_parameters = [
         "output_directory",
         "reference_filepath",
-        "input_file",
+        "input_files",
     ]
     for param in minimum_parameters:
         if param not in parameters["general"]:
             raise UnspecifiedConfigParameterError(
                 "Please specify the following items for analysis: "
-                "output_directory, reference_filepath, and input_file"
+                "output_directory, reference_filepath, and input_files"
             )
 
     # we need to add some extra parameters from the "general" settings
+    parameters["convert"]["output_directory"] = parameters["general"][
+        "output_directory"
+    ]
+    parameters["error_correct_barcodes"]["output_directory"] = parameters[
+        "general"
+    ]["output_directory"]
     parameters["collapse"]["output_directory"] = parameters["general"][
         "output_directory"
     ]
diff --git a/cassiopeia/simulator/BirthDeathFitnessSimulator.py b/cassiopeia/simulator/BirthDeathFitnessSimulator.py
index 4ab3a558..74592e80 100644
--- a/cassiopeia/simulator/BirthDeathFitnessSimulator.py
+++ b/cassiopeia/simulator/BirthDeathFitnessSimulator.py
@@ -146,7 +146,9 @@ def __init__(
         self.collapse_unifurcations = collapse_unifurcations
         self.random_seed = random_seed
 
-    def simulate_tree(self,) -> CassiopeiaTree:
+    def simulate_tree(
+        self,
+    ) -> CassiopeiaTree:
         """Simulates trees from a general birth/death process with fitness.
 
         A forward-time birth/death process is simulated by tracking a series of
@@ -433,4 +435,3 @@ def update_fitness(self, birth_scale: float) -> float:
                     self.fitness_base ** self.fitness_distribution()
                 )
         return birth_scale * base_selection_coefficient
-
diff --git a/cassiopeia/simulator/Cas9LineageTracingDataSimulator.py b/cassiopeia/simulator/Cas9LineageTracingDataSimulator.py
index 9ec27901..129a4e63 100644
--- a/cassiopeia/simulator/Cas9LineageTracingDataSimulator.py
+++ b/cassiopeia/simulator/Cas9LineageTracingDataSimulator.py
@@ -33,7 +33,7 @@ class Cas9LineageTracingDataSimulator(LineageTracingDataSimulator):
     specified mutation rate - specifically, for a lifetime `t` and parameter
     `lambda`, the expected probability of Cas9 mutation, per site, is
     `exp(-lambda * t)`.
-    
+
     Second, the class accepts the architecture of the recorder - described by
     the size of the cassette (by default 3) and the number of cassettes. The
     resulting lineage will have (size_of_cassette * number_of_cassettes)
@@ -58,7 +58,7 @@ class Cas9LineageTracingDataSimulator(LineageTracingDataSimulator):
     heritable silencing rate which is a rare event in which an entire cassette
     is transcriptionally silenced and therefore not observed. The second type
     of silencing is a stochastic dropout rate which simulates the loss of
-    cassettes due to the low sensitivity of the RNA-sequencing assay. 
+    cassettes due to the low sensitivity of the RNA-sequencing assay.
 
     The function `overlay_data` will operate on the tree in place and will
     specifically modify the data stored in the character attributes.
@@ -101,9 +101,9 @@ def __init__(
         number_of_cassettes: int = 10,
         size_of_cassette: int = 3,
         mutation_rate: Union[float, List[float]] = 0.01,
-        state_generating_distribution: Callable[[], float] = lambda: np.random.exponential(
-            1e-5
-        ),
+        state_generating_distribution: Callable[
+            [], float
+        ] = lambda: np.random.exponential(1e-5),
         number_of_states: int = 100,
         state_priors: Optional[Dict[int, float]] = None,
         heritable_silencing_rate: float = 1e-4,
@@ -147,7 +147,6 @@ def __init__(
                 )
             self.mutation_rate_per_character = mutation_rate[:]
 
-
         if state_priors is None:
             self.number_of_states = number_of_states
             self.state_generating_distribution = state_generating_distribution
@@ -184,13 +183,13 @@ def overlay_data(self, tree: CassiopeiaTree):
         if self.mutation_priors is None:
             self.mutation_priors = {}
             probabilites = [
-                self.state_generating_distribution() for _ in range(self.number_of_states)
+                self.state_generating_distribution()
+                for _ in range(self.number_of_states)
             ]
             Z = np.sum(probabilites)
             for i in range(self.number_of_states):
                 self.mutation_priors[i + 1] = probabilites[i] / Z
 
-
         number_of_characters = self.number_of_cassettes * self.size_of_cassette
 
         # initialize character states
@@ -235,9 +234,13 @@ def overlay_data(self, tree: CassiopeiaTree):
             )
 
             # silence cassettes
-            silencing_probability = 1 - (np.exp(-life_time * self.heritable_silencing_rate))
+            silencing_probability = 1 - (
+                np.exp(-life_time * self.heritable_silencing_rate)
+            )
             character_array = self.silence_cassettes(
-                character_array, silencing_probability, self.heritable_missing_data_state
+                character_array,
+                silencing_probability,
+                self.heritable_missing_data_state,
             )
 
             character_matrix[node] = character_array
@@ -245,7 +248,9 @@ def overlay_data(self, tree: CassiopeiaTree):
         # apply stochastic silencing
         for leaf in tree.leaves:
             character_matrix[leaf] = self.silence_cassettes(
-                character_matrix[leaf], self.stochastic_silencing_rate, self.stochastic_missing_data_state
+                character_matrix[leaf],
+                self.stochastic_silencing_rate,
+                self.stochastic_missing_data_state,
             )
 
         tree.initialize_all_character_states(character_matrix)
@@ -282,7 +287,9 @@ def collapse_sites(
                     np.max(sites_to_collapse),
                 )
                 for site in range(left, right + 1):
-                    updated_character_array[site] = self.heritable_missing_data_state
+                    updated_character_array[
+                        site
+                    ] = self.heritable_missing_data_state
             else:
                 cuts_remaining.append(np.array(cuts)[cut_indices[0]])
 
@@ -316,7 +323,10 @@ def introduce_states(
         return updated_character_array
 
     def silence_cassettes(
-        self, character_array: List[int], silencing_rate: float, missing_state: int = -1
+        self,
+        character_array: List[int],
+        silencing_rate: float,
+        missing_state: int = -1,
     ) -> List[int]:
         """Silences cassettes.
 
@@ -352,7 +362,7 @@ def silence_cassettes(
     def get_cassettes(self) -> List[int]:
         """Obtain indices of individual cassettes.
 
-        A helper function that returns the indices that correpspond to the 
+        A helper function that returns the indices that correpspond to the
         independent cassettes in the experiment.
 
         Returns:
diff --git a/cassiopeia/simulator/DataSimulator.py b/cassiopeia/simulator/DataSimulator.py
index a5cec77b..c182e675 100644
--- a/cassiopeia/simulator/DataSimulator.py
+++ b/cassiopeia/simulator/DataSimulator.py
@@ -11,8 +11,10 @@
 
 class DataSimulatorError(Exception):
     """Generic error for the DataSimulator subclasses"""
+
     pass
 
+
 class DataSimulator(abc.ABC):
     """
     DataSimulator is an abstract class that all data overlayers derive from.
diff --git a/cassiopeia/simulator/LeafSubsampler.py b/cassiopeia/simulator/LeafSubsampler.py
index 7787c5e0..6ff7857f 100644
--- a/cassiopeia/simulator/LeafSubsampler.py
+++ b/cassiopeia/simulator/LeafSubsampler.py
@@ -24,9 +24,9 @@ class LeafSubsampler(abc.ABC):
     """
     Abstract base class for all leaf samplers.
 
-    A LeafSubsampler implements a method 'subsample_leaves' which, given a 
-    tree, generates a sample of the observed leaves in that tree and returns a 
-    new tree which is the induced subtree (tree containing only lineages that 
+    A LeafSubsampler implements a method 'subsample_leaves' which, given a
+    tree, generates a sample of the observed leaves in that tree and returns a
+    new tree which is the induced subtree (tree containing only lineages that
     contain a sampled leaf) of the original tree on that sample.
     """
 
@@ -37,13 +37,13 @@ def subsample_leaves(self, tree: CassiopeiaTree) -> CassiopeiaTree:
 
         Returns a new CassiopeiaTree which is the result of subsampling the
         leaves in the original CassiopeiaTree and removing ancestral nodes no
-        longer relevant to the sample. All fields on the original character 
-        matrix persist, but maintains character states, meta data, and the 
+        longer relevant to the sample. All fields on the original character
+        matrix persist, but maintains character states, meta data, and the
         dissimilarity map for the sampled cells only.
 
         Args:
             tree: The CassiopeiaTree for which to subsample leaves
-        
+
         Returns:
             A new CassiopeiaTree that is the induced subtree on a sample of the
             leaves in the given tree.
diff --git a/cassiopeia/simulator/SupercellularSampler.py b/cassiopeia/simulator/SupercellularSampler.py
new file mode 100755
index 00000000..74e94e7e
--- /dev/null
+++ b/cassiopeia/simulator/SupercellularSampler.py
@@ -0,0 +1,144 @@
+"""
+A subclass of LeafSubsampler, the SupercellularSampler.
+
+Iteratively, this subsampler randomly merges two leaves to generate a tree with
+ambiguous character states. The probability that two leaves will be merged is
+proportional to their branch distance.
+"""
+import copy
+from typing import Optional
+
+import numpy as np
+
+from cassiopeia.data import CassiopeiaTree
+from cassiopeia.simulator.LeafSubsampler import (
+    LeafSubsampler,
+    LeafSubsamplerError,
+)
+
+
+class SupercellularSampler(LeafSubsampler):
+    def __init__(
+        self,
+        ratio: Optional[float] = None,
+        number_of_merges: Optional[float] = None,
+    ):
+        """
+        Merge leaves in a tree to generate a new tree with ambiguous states.
+
+        Note that according to this procedure, an already merged (and therefore
+        ambiguous leaf) may be merged again.
+
+        Only one of ``ratio`` or ``number_of_merges`` must be provided.
+
+        Args:
+            ratio: The number of times to merge as a ratio of the total number
+                of leaves. A ratio of 0.5 indicates the number of merges will be
+                approximately half the number of leaves.
+            number_of_merges: Explicit number of merges to perform.
+        """
+        if (ratio is None) == (number_of_merges is None):
+            raise LeafSubsamplerError(
+                "Exactly one of 'ratio' and 'number_of_merges' must be specified."
+            )
+
+        self.__ratio = ratio
+        self.__number_of_merges = number_of_merges
+
+    def subsample_leaves(
+        self,
+        tree: CassiopeiaTree,
+        collapse_source: Optional[str] = None,
+        collapse_duplicates: bool = True,
+    ) -> CassiopeiaTree:
+        """Construct a new CassiopeiaTree by merging leaves.
+
+        Pairs of leaves in the given tree is iteratively merged until a stopping
+        condition is met (specified by ``ratio`` or ``number_of_merges`` when
+        initializing the sampler). Pairs of leaves are selected by the following
+        procedure:
+        1) A random leaf `A` is selected.
+        2) Its pair `B` is randomly selected with probability inversely
+           proportional to the branch distance from the leaf selected in the
+           previous step.
+        3) The pair is merged into a new leaf with name `A-B` and character
+           states merged in to ambiguous states. The new leaf is connected to
+           the LCA of the two leaves and at time max(time of LCA, mean time of
+           the two leaves).
+
+        Args:
+            tree: The CassiopeiaTree for which to subsample leaves
+            collapse_source: The source node from which to collapse unifurcations
+            collapse_duplicates: Whether or not to collapse duplicated character
+                states, so that only unique character states are present in each
+                ambiguous state. Defaults to True.
+
+        Raises:
+            LeafSubsamplerError if the number of merges exceeds the number of
+                leaves in the tree or no merges will be performed.
+        """
+        n_merges = (
+            self.__number_of_merges
+            if self.__number_of_merges is not None
+            else int(tree.n_cell * self.__ratio)
+        )
+        if n_merges >= len(tree.leaves):
+            raise LeafSubsamplerError(
+                "Number of required merges exceeds number of leaves in the tree."
+            )
+        if n_merges == 0:
+            raise LeafSubsamplerError("No merges to be performed.")
+        # Tree needs to have character matrix defined
+        tree.get_current_character_matrix()
+
+        merged_tree = copy.deepcopy(tree)
+        for _ in range(n_merges):
+            # Choose first leaf
+            leaf1 = np.random.choice(merged_tree.leaves)
+            leaf1_state = merged_tree.get_character_states(leaf1)
+
+            # Choose second leaf with weight proportional to inverse distance
+            distances = merged_tree.get_distances(leaf1, leaves_only=True)
+            leaves = []
+            weights = []
+            for leaf in sorted(distances.keys()):
+                if leaf == leaf1:
+                    continue
+                leaves.append(leaf)
+                weights.append(1 / distances[leaf])
+            leaf2 = np.random.choice(
+                leaves, p=np.array(weights) / np.sum(weights)
+            )
+
+            leaf2_state = merged_tree.get_character_states(leaf2)
+
+            # Merge these two leaves at the mean time of the two leaves.
+            # Note that the mean time of the two leaves may never be earlier than
+            # the LCA time, because each of the leaf times must be greater than or
+            # equal to the LCA time.
+            # If the tree is ultrametric, this preserves ultrametricity.
+            new_leaf = f"{leaf1}-{leaf2}"
+            lca = merged_tree.find_lca(leaf1, leaf2)
+            new_time = (
+                merged_tree.get_time(leaf1) + merged_tree.get_time(leaf2)
+            ) / 2
+            new_state = []
+            for char1, char2 in zip(leaf1_state, leaf2_state):
+                new_char = []
+                if not isinstance(char1, tuple):
+                    char1 = (char1,)
+                if not isinstance(char2, tuple):
+                    char2 = (char2,)
+                new_state.append(char1 + char2)
+            merged_tree.add_leaf(lca, new_leaf, states=new_state, time=new_time)
+            merged_tree.remove_leaf_and_prune_lineage(leaf1)
+            merged_tree.remove_leaf_and_prune_lineage(leaf2)
+
+        if collapse_source is None:
+            collapse_source = merged_tree.root
+        merged_tree.collapse_unifurcations(source=collapse_source)
+
+        if collapse_duplicates:
+            merged_tree.collapse_ambiguous_characters()
+
+        return merged_tree
diff --git a/cassiopeia/simulator/UniformLeafSubsampler.py b/cassiopeia/simulator/UniformLeafSubsampler.py
index e4280bb2..5a4c8d1f 100644
--- a/cassiopeia/simulator/UniformLeafSubsampler.py
+++ b/cassiopeia/simulator/UniformLeafSubsampler.py
@@ -12,18 +12,23 @@
 from typing import Optional
 
 from cassiopeia.data import CassiopeiaTree
-from cassiopeia.simulator.LeafSubsampler import LeafSubsampler, LeafSubsamplerError
+from cassiopeia.simulator.LeafSubsampler import (
+    LeafSubsampler,
+    LeafSubsamplerError,
+)
 
 
 class UniformLeafSubsampler(LeafSubsampler):
     def __init__(
-        self, ratio: Optional[float] = None, number_of_leaves: Optional[int] = None
+        self,
+        ratio: Optional[float] = None,
+        number_of_leaves: Optional[int] = None,
     ):
         """
         Uniformly subsample leaf samples of a CassiopeiaTree.
 
         If 'ratio' is provided, samples 'ratio' of the leaves, rounded down,
-        uniformly at random. If instead 'number_of_leaves' is provided, 
+        uniformly at random. If instead 'number_of_leaves' is provided,
         'number_of_leaves' of the leaves are sampled uniformly at random. Only
         one of the two criteria can be provided.
 
@@ -34,30 +39,34 @@ def __init__(
         """
         if ratio is None and number_of_leaves is None:
             raise LeafSubsamplerError(
-                "At least one of 'ratio' and 'number_of_leaves' " "must be specified."
+                "At least one of 'ratio' and 'number_of_leaves' "
+                "must be specified."
             )
         if ratio is not None and number_of_leaves is not None:
             raise LeafSubsamplerError(
-                "Exactly one of 'ratio' and 'number_of_leaves'" "must be specified."
+                "Exactly one of 'ratio' and 'number_of_leaves'"
+                "must be specified."
             )
         self.__ratio = ratio
         self.__number_of_leaves = number_of_leaves
 
-    def subsample_leaves(self, tree: CassiopeiaTree, collapse_source: str = None) -> CassiopeiaTree:
+    def subsample_leaves(
+        self, tree: CassiopeiaTree, collapse_source: str = None
+    ) -> CassiopeiaTree:
         """Uniformly subsample leaf samples of a given tree.
 
         Generates a uniform random sample on the leaves of the given
         CassiopeiaTree and returns a tree pruned to contain lineages relevant
-        to only leaves in the sample (the "induced subtree" on the sample). 
-        All fields on the original character matrix persist, but maintains 
+        to only leaves in the sample (the "induced subtree" on the sample).
+        All fields on the original character matrix persist, but maintains
         character states, meta data, and the dissimilarity map for the sampled
         cells only.
-        
+
         Args:
             tree: The CassiopeiaTree for which to subsample leaves
-            collapse_source: The source node from which to collapse 
+            collapse_source: The source node from which to collapse
                 unifurcations
-            
+
         Returns:
             A new CassiopeiaTree that is the induced subtree on a sample of the
             leaves in the given tree.
@@ -65,7 +74,9 @@ def subsample_leaves(self, tree: CassiopeiaTree, collapse_source: str = None) ->
         ratio = self.__ratio
         number_of_leaves = self.__number_of_leaves
         n_subsample = (
-            number_of_leaves if number_of_leaves is not None else int(tree.n_cell * ratio)
+            number_of_leaves
+            if number_of_leaves is not None
+            else int(tree.n_cell * ratio)
         )
         if n_subsample <= 0:
             raise LeafSubsamplerError(
@@ -79,16 +90,21 @@ def subsample_leaves(self, tree: CassiopeiaTree, collapse_source: str = None) ->
 
         n_remove = len(tree.leaves) - n_subsample
         subsampled_tree = copy.deepcopy(tree)
-        leaf_remove = np.random.choice(subsampled_tree.leaves, n_remove, replace=False)
+        leaf_remove = np.random.choice(
+            subsampled_tree.leaves, n_remove, replace=False
+        )
 
         for i in leaf_remove:
             subsampled_tree.remove_leaf_and_prune_lineage(i)
 
         if collapse_source is None:
             collapse_source = subsampled_tree.root
-        subsampled_tree.collapse_unifurcations(source = collapse_source)
+        subsampled_tree.collapse_unifurcations(source=collapse_source)
 
         # Copy and annotate branch lengths and times
-        subsampled_tree.set_times(dict([(node, tree.get_time(node)) for node in subsampled_tree.nodes]))
+        subsampled_tree.set_times(
+            dict(
+                [(node, tree.get_time(node)) for node in subsampled_tree.nodes]
+            )
+        )
         return subsampled_tree
-
diff --git a/cassiopeia/simulator/__init__.py b/cassiopeia/simulator/__init__.py
old mode 100644
new mode 100755
index 836407d3..b7be7435
--- a/cassiopeia/simulator/__init__.py
+++ b/cassiopeia/simulator/__init__.py
@@ -7,3 +7,5 @@
 from .SimpleFitSubcloneSimulator import SimpleFitSubcloneSimulator
 from .Cas9LineageTracingDataSimulator import Cas9LineageTracingDataSimulator
 from .LeafSubsampler import LeafSubsampler, LeafSubsamplerError
+from .SupercellularSampler import SupercellularSampler
+from .UniformLeafSubsampler import UniformLeafSubsampler
diff --git a/cassiopeia/solver/GreedySolver.py b/cassiopeia/solver/GreedySolver.py
index 08f8a9d1..ba81c27f 100755
--- a/cassiopeia/solver/GreedySolver.py
+++ b/cassiopeia/solver/GreedySolver.py
@@ -1,6 +1,6 @@
 """
 This file stores a subclass of CassiopeiaSolver, the GreedySolver. This class
-represents the structure of top-down algorithms that build the reconstructed 
+represents the structure of top-down algorithms that build the reconstructed
 tree by recursively splitting the set of samples based on some split criterion.
 """
 import logging
@@ -149,8 +149,14 @@ def _solve(
         cassiopeia_tree.populate_tree(tree)
 
         # Collapse 0-mutation edges and append duplicate samples
-        cassiopeia_tree.collapse_mutationless_edges(infer_ancestral_characters = True)
-        duplicates_tree = self.__add_duplicates_to_tree(cassiopeia_tree.get_tree_topology(), character_matrix, node_name_generator)
+        cassiopeia_tree.collapse_mutationless_edges(
+            infer_ancestral_characters=True
+        )
+        duplicates_tree = self.__add_duplicates_to_tree(
+            cassiopeia_tree.get_tree_topology(),
+            character_matrix,
+            node_name_generator,
+        )
         cassiopeia_tree.populate_tree(duplicates_tree)
 
     def compute_mutation_frequencies(
@@ -189,7 +195,10 @@ def compute_mutation_frequencies(
         return freq_dict
 
     def __add_duplicates_to_tree(
-        self, tree: nx.DiGraph, character_matrix: pd.DataFrame, node_name_generator: Generator[str, None, None],
+        self,
+        tree: nx.DiGraph,
+        character_matrix: pd.DataFrame,
+        node_name_generator: Generator[str, None, None],
     ) -> nx.DiGraph:
         """Takes duplicate samples and places them in the tree.
 
@@ -220,4 +229,4 @@ def __add_duplicates_to_tree(
             for duplicate in duplicate_groups[i]:
                 tree.add_edge(new_internal_node, duplicate)
 
-        return tree
\ No newline at end of file
+        return tree
diff --git a/cassiopeia/solver/HybridSolver.py b/cassiopeia/solver/HybridSolver.py
index 7d92ae9c..bd7804f6 100644
--- a/cassiopeia/solver/HybridSolver.py
+++ b/cassiopeia/solver/HybridSolver.py
@@ -37,7 +37,7 @@ class HybridSolver(CassiopeiaSolver.CassiopeiaSolver):
     """
     HybridSolver is an class representing the structure of Cassiopeia Hybrid
     inference algorithms. The solver procedure contains logic for building tree
-    starting with a top-down greedy algorithm until a predetermined criteria is 
+    starting with a top-down greedy algorithm until a predetermined criteria is
     reached at which point a more complex algorithm is used to reconstruct each
     subproblem. The top-down algorithm _must_ be a subclass of a GreedySolver
     as it must have functions `find_split` and `perform_split`. The solver
@@ -172,8 +172,12 @@ def solve(
         cassiopeia_tree.populate_tree(self.__tree)
 
         # Collapse 0-mutation edges and append duplicate samples
-        cassiopeia_tree.collapse_mutationless_edges(infer_ancestral_characters = True)
-        samples_tree = self.__append_sample_names(cassiopeia_tree.get_tree_topology(), character_matrix)
+        cassiopeia_tree.collapse_mutationless_edges(
+            infer_ancestral_characters=True
+        )
+        samples_tree = self.__append_sample_names(
+            cassiopeia_tree.get_tree_topology(), character_matrix
+        )
         cassiopeia_tree.populate_tree(samples_tree)
 
     def apply_top_solver(
@@ -235,7 +239,12 @@ def apply_top_solver(
 
         for clade in new_clades:
             child, new_subproblems = self.apply_top_solver(
-                character_matrix, clade, node_name_generator, weights, missing_state_indicator, root
+                character_matrix,
+                clade,
+                node_name_generator,
+                weights,
+                missing_state_indicator,
+                root,
             )
             self.__tree.add_edge(root, child)
 
@@ -386,4 +395,4 @@ def __append_sample_names(
                 tree.add_edges_from([(node, sample) for sample in samples])
                 states_added.append(node)
 
-        return tree
\ No newline at end of file
+        return tree
diff --git a/cassiopeia/solver/MaxCutGreedySolver.py b/cassiopeia/solver/MaxCutGreedySolver.py
index 4319d1d3..fb0aef0c 100644
--- a/cassiopeia/solver/MaxCutGreedySolver.py
+++ b/cassiopeia/solver/MaxCutGreedySolver.py
@@ -55,7 +55,7 @@ class MaxCutGreedySolver(GreedySolver.GreedySolver):
     def __init__(
         self,
         missing_data_classifier: Callable = missing_data_methods.assign_missing_average,
-        prior_transformation: str = "negative_log"
+        prior_transformation: str = "negative_log",
     ):
 
         super().__init__(prior_transformation)
diff --git a/cassiopeia/solver/NeighborJoiningSolver.py b/cassiopeia/solver/NeighborJoiningSolver.py
old mode 100644
new mode 100755
index 794367c5..eff1fc19
--- a/cassiopeia/solver/NeighborJoiningSolver.py
+++ b/cassiopeia/solver/NeighborJoiningSolver.py
@@ -16,6 +16,7 @@
 from cassiopeia.solver import (
     DistanceSolver,
     dissimilarity_functions,
+    solver_utilities,
 )
 
 
@@ -217,13 +218,12 @@ def __update_dissimilarity_map_numba(
         for v in range(dissimilarity_map.shape[0]):
             if v == cherry_i or v == cherry_j:
                 continue
-            updated_map[v, new_node_index] = updated_map[new_node_index, v] = (
-                0.5
-                * (
-                    dissimilarity_map[v, cherry_i]
-                    + dissimilarity_map[v, cherry_j]
-                    - dissimilarity_map[cherry_i, cherry_j]
-                )
+            updated_map[v, new_node_index] = updated_map[
+                new_node_index, v
+            ] = 0.5 * (
+                dissimilarity_map[v, cherry_i]
+                + dissimilarity_map[v, cherry_j]
+                - dissimilarity_map[cherry_i, cherry_j]
             )
 
         updated_map[new_node_index, new_node_index] = 0
@@ -236,7 +236,8 @@ def setup_root_finder(self, cassiopeia_tree: CassiopeiaTree) -> None:
         By default, the NeighborJoining algorithm returns an unrooted tree.
         To root this tree, an implicit root of all zeros is added to the
         character matrix. Then, the dissimilarity map is recalculated using
-        the updated character matrix.
+        the updated character matrix. If the tree already has a computed
+        dissimilarity map, only the new dissimilarities are calculated.
 
         Args:
             cassiopeia_tree: Input CassiopeiaTree to `solve`
@@ -255,8 +256,25 @@ def setup_root_finder(self, cassiopeia_tree: CassiopeiaTree) -> None:
                 "root, or specify an explicit root"
             )
 
-        cassiopeia_tree.compute_dissimilarity_map(
-            self.dissimilarity_function, self.prior_transformation
-        )
+        dissimilarity_map = cassiopeia_tree.get_dissimilarity_map()
+        if dissimilarity_map is None:
+            cassiopeia_tree.compute_dissimilarity_map(
+                self.dissimilarity_function, self.prior_transformation
+            )
+        else:
+            dissimilarity = {"root": 0}
+            for leaf in character_matrix.index:
+                weights = None
+                if cassiopeia_tree.priors:
+                    weights = solver_utilities.transform_priors(
+                        cassiopeia_tree.priors, self.prior_transformation
+                    )
+                dissimilarity[leaf] = self.dissimilarity_function(
+                    rooted_character_matrix.loc["root"].values,
+                    rooted_character_matrix.loc[leaf].values,
+                    cassiopeia_tree.missing_state_indicator,
+                    weights,
+                )
+            cassiopeia_tree.set_dissimilarity("root", dissimilarity)
 
         cassiopeia_tree.set_character_matrix(character_matrix)
diff --git a/cassiopeia/solver/PercolationSolver.py b/cassiopeia/solver/PercolationSolver.py
index 73b63f47..c3974f01 100644
--- a/cassiopeia/solver/PercolationSolver.py
+++ b/cassiopeia/solver/PercolationSolver.py
@@ -13,6 +13,7 @@
     solver_utilities,
 )
 
+
 class PercolationSolver(CassiopeiaSolver.CassiopeiaSolver):
     """
     The PercolationSolver implements a top-down algorithm that recursively
@@ -67,7 +68,7 @@ def __init__(
                 float,
             ]
         ] = dissimilarity_functions.hamming_similarity_without_missing,
-        threshold: Optional[int] = 0
+        threshold: Optional[int] = 0,
     ):
 
         super().__init__(prior_transformation)
@@ -169,8 +170,14 @@ def _solve(
         cassiopeia_tree.populate_tree(tree)
 
         # Collapse 0-mutation edges and append duplicate samples
-        cassiopeia_tree.collapse_mutationless_edges(infer_ancestral_characters = True)
-        duplicates_tree = self.__add_duplicates_to_tree(cassiopeia_tree.get_tree_topology(), character_matrix, node_name_generator)
+        cassiopeia_tree.collapse_mutationless_edges(
+            infer_ancestral_characters=True
+        )
+        duplicates_tree = self.__add_duplicates_to_tree(
+            cassiopeia_tree.get_tree_topology(),
+            character_matrix,
+            node_name_generator,
+        )
         cassiopeia_tree.populate_tree(duplicates_tree)
 
     def percolate(
@@ -309,7 +316,10 @@ def percolate(
         return partition_named
 
     def __add_duplicates_to_tree(
-        self, tree: nx.DiGraph, character_matrix: pd.DataFrame, node_name_generator: Generator[str, None, None]
+        self,
+        tree: nx.DiGraph,
+        character_matrix: pd.DataFrame,
+        node_name_generator: Generator[str, None, None],
     ) -> nx.DiGraph:
         """Takes duplicate samples and places them in the tree.
         Places samples removed in removing duplicates in the tree as sisters
@@ -337,4 +347,4 @@ def __add_duplicates_to_tree(
             for duplicate in duplicate_groups[i]:
                 tree.add_edge(new_internal_node, duplicate)
 
-        return tree
\ No newline at end of file
+        return tree
diff --git a/cassiopeia/solver/SharedMutationJoiningSolver.py b/cassiopeia/solver/SharedMutationJoiningSolver.py
old mode 100644
new mode 100755
index ba890c31..6addafc8
--- a/cassiopeia/solver/SharedMutationJoiningSolver.py
+++ b/cassiopeia/solver/SharedMutationJoiningSolver.py
@@ -1,12 +1,12 @@
 """
-This file stores the SharedMutationJoiningSolver. The inference procedure is 
+This file stores the SharedMutationJoiningSolver. The inference procedure is
 an agglomerative clustering procedure that joins samples that share the most
 identical character/state mutations.
 """
-
+import abc
+import warnings
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
-import abc
 import networkx as nx
 import numba
 import numpy as np
@@ -15,7 +15,11 @@
 
 from cassiopeia.data import CassiopeiaTree
 from cassiopeia.data import utilities as data_utilities
-from cassiopeia.solver import CassiopeiaSolver, dissimilarity_functions, solver_utilities
+from cassiopeia.solver import (
+    CassiopeiaSolver,
+    dissimilarity_functions,
+    solver_utilities,
+)
 
 
 class SharedMutationJoiningSolverError(Exception):
@@ -24,6 +28,12 @@ class SharedMutationJoiningSolverError(Exception):
     pass
 
 
+class SharedMutationJoiningSolverWarning(UserWarning):
+    """A warning class for SharedMutationJoiningSolver."""
+
+    pass
+
+
 class SharedMutationJoiningSolver(CassiopeiaSolver.CassiopeiaSolver):
     """Shared-Mutation-Joining class for Cassiopeia.
 
@@ -73,8 +83,26 @@ def __init__(
 
         super().__init__(prior_transformation)
 
+        # Attempt to numbaize
         self.similarity_function = similarity_function
+        self.nb_similarity_function = similarity_function
+        numbaize = True
+        try:
+            self.nb_similarity_function = numba.jit(
+                similarity_function, nopython=True
+            )
+        except TypeError:
+            numbaize = False
+            warnings.warn(
+                "Failed to numbaize dissimilarity function. "
+                "Falling back to Python.",
+                SharedMutationJoiningSolverWarning,
+            )
 
+        if numbaize:
+            self.__update_similarity_map = numba.jit(
+                self.__update_similarity_map, nopython=True
+            )
 
     def solve(self, cassiopeia_tree: CassiopeiaTree) -> None:
         """Solves a tree for the SharedMutationJoiningSolver.
@@ -122,9 +150,6 @@ def solve(self, cassiopeia_tree: CassiopeiaTree) -> None:
         N = similarity_map.shape[0]
 
         # Numba-ize the similarity function and weights
-        nb_similarity_function = numba.jit(
-            self.similarity_function, nopython=True
-        )
         nb_weights = numba.typed.Dict.empty(
             numba.types.int64,
             numba.types.DictType(numba.types.int64, numba.types.float64),
@@ -157,7 +182,7 @@ def solve(self, cassiopeia_tree: CassiopeiaTree) -> None:
 
             similarity_map = self.update_similarity_map_and_character_matrix(
                 character_matrix,
-                nb_similarity_function,
+                self.nb_similarity_function,
                 similarity_map,
                 (node_i, node_j),
                 new_node_name,
@@ -191,7 +216,7 @@ def find_cherry(self, similarity_matrix: np.array) -> Tuple[int, int]:
     def update_similarity_map_and_character_matrix(
         self,
         character_matrix: pd.DataFrame,
-        nb_similarity_function: Callable[
+        similarity_function: Callable[
             [np.array, np.array, int, Dict[int, Dict[int, float]]], float
         ],
         similarity_map: pd.DataFrame,
@@ -210,7 +235,7 @@ def update_similarity_map_and_character_matrix(
         Args:
             character_matrix: Contains the character information for all nodes,
                 updated as nodes are joined and new internal LCA nodes are added
-            nb_similarity: A numba compiled similarity function
+            similarity_function: A similarity function
             similarity_map: A similarity map to update
             cherry: A tuple of indices in the similarity map that are joining
             new_node: New node name, to be added to the updated similarity map
@@ -236,11 +261,11 @@ def update_similarity_map_and_character_matrix(
         )
         character_matrix.loc[new_node] = lca
 
-        similarity_array_updated = self.__update_similarity_map_numba(
+        similarity_array_updated = self.__update_similarity_map(
             character_array,
             similarity_array,
             np.array(lca),
-            nb_similarity_function,
+            similarity_function,
             missing_state_indicator,
             weights,
         )
@@ -263,12 +288,11 @@ def update_similarity_map_and_character_matrix(
         return similarity_map
 
     @staticmethod
-    @numba.jit(nopython=True)
-    def __update_similarity_map_numba(
+    def __update_similarity_map(
         character_matrix: np.array,
         similarity_map: np.array,
         lca: np.array,
-        nb_similarity_function: Callable[
+        similarity_function: Callable[
             [np.array, np.array, int, Dict[int, Dict[int, float]]], float
         ],
         missing_state_indicator: int = -1,
@@ -284,7 +308,7 @@ def __update_similarity_map_numba(
             character_matrix: The character information for all nodes
             similarity_map: A matrix of similarities to update
             lca: The character vector of the new LCA node
-            nb_similarity_function: A numba compiled similarity function
+            similarity_function: A similarity function
             missing_state_indicator: Character representing missing data
             weights: Weighting of each (character, state) pair. Typically a
                 transformation of the priors.
@@ -294,11 +318,11 @@ def __update_similarity_map_numba(
         """
 
         C = similarity_map.shape[0]
-        new_row = np.zeros(C, dtype=numba.float64)
+        new_row = np.zeros(C, dtype=np.float64)
         k = 0
         for i in range(C):
             s1 = character_matrix[i, :]
-            new_row[k] = nb_similarity_function(
+            new_row[k] = similarity_function(
                 s1, lca, missing_state_indicator, weights
             )
             k += 1
diff --git a/cassiopeia/solver/VanillaGreedySolver.py b/cassiopeia/solver/VanillaGreedySolver.py
index 61d5d327..f8de0382 100755
--- a/cassiopeia/solver/VanillaGreedySolver.py
+++ b/cassiopeia/solver/VanillaGreedySolver.py
@@ -52,7 +52,7 @@ class VanillaGreedySolver(GreedySolver.GreedySolver):
     def __init__(
         self,
         missing_data_classifier: Callable = missing_data_methods.assign_missing_average,
-        prior_transformation: str = "negative_log"
+        prior_transformation: str = "negative_log",
     ):
 
         super().__init__(prior_transformation)
diff --git a/cassiopeia/solver/__init__.py b/cassiopeia/solver/__init__.py
index 46850ea8..23db77bf 100755
--- a/cassiopeia/solver/__init__.py
+++ b/cassiopeia/solver/__init__.py
@@ -15,4 +15,4 @@
 from .SpectralSolver import SpectralSolver
 from .UPGMASolver import UPGMASolver
 from .VanillaGreedySolver import VanillaGreedySolver
-from . import dissimilarity_functions as dissimilarity
\ No newline at end of file
+from . import dissimilarity_functions as dissimilarity
diff --git a/cassiopeia/solver/dissimilarity_functions.py b/cassiopeia/solver/dissimilarity_functions.py
old mode 100644
new mode 100755
index dea7342c..47e24a10
--- a/cassiopeia/solver/dissimilarity_functions.py
+++ b/cassiopeia/solver/dissimilarity_functions.py
@@ -2,7 +2,8 @@
 A library that contains dissimilarity functions for the purpose of comparing
 phylogenetic samples.
 """
-from typing import Dict, List, Optional
+import itertools
+from typing import Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import numba
 import numpy as np
@@ -110,6 +111,7 @@ def hamming_similarity_without_missing(
 
     return similarity
 
+
 def hamming_similarity_normalized_over_missing(
     s1: List[int],
     s2: List[int],
@@ -134,18 +136,12 @@ def hamming_similarity_normalized_over_missing(
     similarity = 0
     num_present = 0
     for i in range(len(s1)):
-        if (
-            s1[i] == missing_state_indicator
-            or s2[i] == missing_state_indicator
-        ):
+        if s1[i] == missing_state_indicator or s2[i] == missing_state_indicator:
             continue
 
         num_present += 1
 
-        if (
-            s1[i] == 0
-            or s2[i] == 0
-        ):
+        if s1[i] == 0 or s2[i] == 0:
             continue
 
         if s1[i] == s2[i]:
@@ -157,7 +153,7 @@ def hamming_similarity_normalized_over_missing(
     if num_present == 0:
         return 0
 
-    return similarity/num_present
+    return similarity / num_present
 
 
 @numba.jit(nopython=True)
@@ -241,3 +237,96 @@ def weighted_hamming_similarity(
         return 0
 
     return d / num_present
+
+
+def cluster_dissimilarity(
+    dissimilarity_function: Callable[
+        [List[int], List[int], int, Dict[int, Dict[int, float]]], float
+    ],
+    s1: Union[List[int], List[Tuple[int, ...]]],
+    s2: Union[List[int], List[Tuple[int, ...]]],
+    missing_state_indicator: int,
+    weights: Optional[Dict[int, Dict[int, float]]] = None,
+    linkage_function: Callable[[Union[np.array, List[float]]], float] = np.mean,
+    normalize: bool = True,
+) -> float:
+    """Compute the dissimilarity between (possibly) ambiguous character strings.
+
+    An ambiguous character string is a character string in
+    which each character contains an tuple of possible states, and such a
+    character string is represented as a list of tuples of integers.
+
+    A naive implementation is to first disambiguate each of the two ambiguous
+    character strings by generating all possible strings, then computing the
+    dissimilarity between all pairwise combinations, and finally applying the
+    linkage function on the calculated dissimilarities. However, doing so has
+    complexity O(\prod_{i=1}^N |a_i| x |b_i|) where N is the number of target sites,
+    |a_i| is the number of ambiguous characters at target site i of string a,
+    and |b_i| is the number of amiguous characters at target site i of string b.
+    As an example, if we have N=10 and all a_i=b_i=2, then we have to construct
+    1,038,576 * 2 strings and compute over 4 trillion dissimilarities.
+
+    By assuming each target site is independent, simply calculating the sum of
+    the linkages of each target site separately is equivalent to the naive
+    implementation (can be proven by induction). This procedure is implemented
+    in this function. One caveat is that we usually normalize the distance by
+    the number of shared non-missing positions. We approximate this by dividing
+    the absolute distance by the sum of the probability of each site not being
+    a missing site for both strings.
+
+    The idea of linkage is analogous to that in hierarchical clustering, where
+    ``np.min`` can be used for single linkage, ``np.max`` for complete linkage,
+    and ``np.mean`` for average linkage (the default).
+
+    The reason the ``dissimilarity_function`` argument is defined as the first
+    argument is so that this function may be used as input to
+    :func:`cassiopeia.data.CassiopeiaTree.compute_dissimilarity_map`. This can
+    be done by partial application of this function with the desired dissimilarity
+    function.
+
+    Note:
+        If neither character string is ambiguous, then calling this function is
+        equivalent to calling ``dissimilarity_function`` separately.
+
+    Args:
+        s1: The first (possibly) ambiguous sample
+        s2: The second (possibly) ambiguous sample
+        missing_state_indicator: The character representing missing values
+        weights: A set of optional weights to weight the similarity of a mutation
+        dissimilarity_function: The dissimilarity function to use to calculate pairwise
+            dissimilarities.
+        linkage_function: The linkage function to use to aggregate dissimilarities
+            into a single number. Defaults to ``np.mean`` for average linkage.
+        normalize: Whether to normalize to the proportion of sites present in
+            both strings.
+
+    Returns:
+        The dissimilarity between the two ambiguous samples
+    """
+    # Make any unambiguous character strings into pseudo-ambiguous so that we
+    # can easily use itertools.product
+    s1 = [s if isinstance(s, tuple) else (s,) for s in s1]
+    s2 = [s if isinstance(s, tuple) else (s,) for s in s2]
+
+    result = 0
+    num_present = 0
+    for i, (c1, c2) in enumerate(zip(s1, s2)):
+        dissim = []
+        present = []
+        for _c1, _c2 in itertools.product(c1, c2):
+            present.append(
+                _c1 != missing_state_indicator
+                and _c2 != missing_state_indicator
+            )
+            dissim.append(
+                dissimilarity_function(
+                    [_c1],
+                    [_c2],
+                    missing_state_indicator,
+                    {0: weights[i]} if weights else None,
+                )
+            )
+        result += linkage_function(dissim)
+        num_present += np.mean(present)
+
+    return result / num_present if normalize else result
diff --git a/cassiopeia/solver/solver_utilities.py b/cassiopeia/solver/solver_utilities.py
index 056b890b..762b80af 100755
--- a/cassiopeia/solver/solver_utilities.py
+++ b/cassiopeia/solver/solver_utilities.py
@@ -18,7 +18,7 @@ class PriorTransformationError(Exception):
 
 def node_name_generator() -> Generator[str, None, None]:
     """Generates unique node names for building the reconstructed tree.
-    
+
     Creates a generator object that produces unique node names by hashing
     timestamps.
 
@@ -27,7 +27,7 @@ def node_name_generator() -> Generator[str, None, None]:
     """
 
     while True:
-        k = str(time.time()).encode('utf-8')
+        k = str(time.time()).encode("utf-8")
         h = blake2b(key=k, digest_size=12)
         yield "cassiopeia_internal_node" + h.hexdigest()
 
diff --git a/data/preprocess.cfg b/data/preprocess.cfg
old mode 100644
new mode 100755
index 97b16361..11b81ed9
--- a/data/preprocess.cfg
+++ b/data/preprocess.cfg
@@ -1,9 +1,18 @@
 [general]
-output_directory =  "/data/yosef2/users/richardz/projects/Cassiopeia/test_process_pipeline"
-reference_filepath = "/data/yosef2/users/mattjones/projects/scLineages/Cassiopeia/data/PCT48.ref.fasta"
-entry = "collapse"
+output_directory = "/mnt/e/scratch/cassiopeia/pipeline_test"
+reference_filepath = "/mnt/e/scratch/cassiopeia/PCT48-long.ref.fa"
+entry = "convert"
 exit = "call_lineages"
-input_file = "/data/yosef2/users/richardz/projects/Cassiopeia/notebooks/test_possorted_genome_bam.bam"
+input_files = ["/mnt/e/scratch/cassiopeia/smaller_1.fastq.gz", "/mnt/e/scratch/cassiopeia/smaller_2.fastq.gz"]
+
+[convert]
+chemistry = "10xv3"
+name = "test_sample"
+n_threads = 32
+
+[error_correct_barcodes]
+whitelist_fp = "/mnt/e/scratch/cassiopeia/pipeline_test/10x_version3_whitelist.txt"
+n_threads = 32
 
 [collapse]
 max_hq_mismatches = 3
@@ -26,8 +35,7 @@ cutsite_width = 12
 context = True
 context_size = 5
 
-[error_correct]
-_id = "test"
+[error_correct_umis]
 max_umi_distance = 2
 verbose = False
 
@@ -43,8 +51,6 @@ plot = True
 verbose =  False
 
 [call_lineages]
-out_fn = "test_processed_output.txt"
-output_directory =  "/data/yosef2/users/richardz/projects/Cassiopeia/test_process_pipeline"
 min_umi_per_cell = 10
 min_avg_reads_per_umi = 2.0
 min_cluster_prop = 0.005
diff --git a/setup.py b/setup.py
index 046a5b06..27d07d75 100755
--- a/setup.py
+++ b/setup.py
@@ -13,25 +13,26 @@
 
 
 requirements = [
-    "numpy > 1.17",
-    "matplotlib >= 2.2.2",
-    "pandas >= 0.22.0",
-    "networkx >= 2.5",
-    "tqdm >= 4",
-    "ete3 >= 3.1.1",
-    "argparse >= 1.1",
-    "Biopython >= 1.71",
-    "pandas-charm >= 0.1.3",
-    "pysam >= 0.14.1",
-    "bokeh >= 0.12.15",
-    "PyYAML >= 3.12",
-    "cython >= 0.29.2",
-    "scipy >= 1.2.0",
-    "python-Levenshtein",
-    "nbconvert >= 5.4.0",
-    "nbformat >= 4.4.0",
+    "Biopython>=1.71",
+    "bokeh>=0.12.15",
+    "cython>=0.29.2",
+    "ete3>=3.1.1",
     "hits",
-    "scikit-bio >= 0.5.6",
+    "matplotlib>=2.2.2",
+    "nbconvert>=5.4.0",
+    "nbformat>=4.4.0",
+    "networkx>=2.5",
+    "ngs-tools>=1.4.0",
+    "numba>=0.51.0",
+    "numpy>=1.19.5",
+    "pandas>=1.1.4",
+    "pysam>=0.14.1",
+    "python-Levenshtein",
+    "PyYAML>=3.12",
+    "scikit-bio>=0.5.6",
+    "scipy>=1.2.0",
+    "typing-extensions>=3.7.4",
+    "tqdm>=4",
 ]
 
 
@@ -71,7 +72,9 @@
     # ext_modules=to_cythonize,
     setup_requires=["cython", "numpy"],
     cmdclass=cmdclass,
-    entry_points={"console_scripts": ["scLT = cassiopeia.__main__:main"]},
+    entry_points={"console_scripts": [
+        "cassiopeia-preprocess = cassiopeia.preprocess.cassiopeia_preprocess:main",
+    ]},
     author_email="mattjones315@berkeley.edu",
     classifiers=[
         "Development Status :: 4 - Beta",


Test patch:
diff --git a/test/critique_tests/test_tree_compare.py b/test/critique_tests/test_tree_compare.py
index 17de7e40..6c6e39d6 100644
--- a/test/critique_tests/test_tree_compare.py
+++ b/test/critique_tests/test_tree_compare.py
@@ -14,7 +14,9 @@
 class TestTreeComparisons(unittest.TestCase):
     def setUp(self):
 
-        self.ground_truth_tree = cas.data.CassiopeiaTree(tree=nx.balanced_tree(2, 3, create_using=nx.DiGraph))
+        self.ground_truth_tree = cas.data.CassiopeiaTree(
+            tree=nx.balanced_tree(2, 3, create_using=nx.DiGraph)
+        )
 
         tree1 = nx.DiGraph()
         tree1.add_nodes_from([i for i in range(15)])
@@ -63,7 +65,9 @@ def setUp(self):
                 (7, 17),
             ]
         )
-        self.multifurcating_ground_truth = cas.data.CassiopeiaTree(tree=multifurcating_ground_truth)
+        self.multifurcating_ground_truth = cas.data.CassiopeiaTree(
+            tree=multifurcating_ground_truth
+        )
 
         tree2 = nx.DiGraph()
         tree2.add_edges_from(
@@ -98,11 +102,15 @@ def setUp(self):
 
     def test_out_group(self):
 
-        out_group = cas.critique.critique_utilities.get_outgroup(self.tree1, ("11", "14", "9"))
+        out_group = cas.critique.critique_utilities.get_outgroup(
+            self.tree1, ("11", "14", "9")
+        )
 
         self.assertEqual("9", out_group)
 
-        out_group = cas.critique.critique_utilities.get_outgroup(self.ground_truth_rake, ("4", "5", "6"))
+        out_group = cas.critique.critique_utilities.get_outgroup(
+            self.ground_truth_rake, ("4", "5", "6")
+        )
         self.assertEqual("None", out_group)
 
     def test_same_tree_gives_perfect_triplets_correct(self):
diff --git a/test/data_tests/cassiopeia_tree_test.py b/test/data_tests/cassiopeia_tree_test.py
old mode 100644
new mode 100755
index be23b97b..e22fde6f
--- a/test/data_tests/cassiopeia_tree_test.py
+++ b/test/data_tests/cassiopeia_tree_test.py
@@ -2,6 +2,7 @@
 Tests for the CassiopeiaTree object in the data module.
 """
 import unittest
+from functools import partial
 from typing import Dict, Optional
 
 import ete3
@@ -17,6 +18,20 @@
     CassiopeiaTreeError,
     CassiopeiaTreeWarning,
 )
+from cassiopeia.solver import dissimilarity_functions
+
+
+def delta_fn(
+    x: np.array,
+    y: np.array,
+    missing_state: int,
+    priors: Optional[Dict[int, Dict[int, float]]],
+):
+    d = 0
+    for i in range(len(x)):
+        if x[i] != y[i]:
+            d += 1
+    return d
 
 
 class TestCassiopeiaTree(unittest.TestCase):
@@ -64,6 +79,30 @@ def setUp(self):
             },
             orient="index",
         )
+        self.ambiguous_character_matrix = pd.DataFrame.from_dict(
+            {
+                "node3": [1, 0, 0, 0, 0, 0, 0, 0],
+                "node7": [1, 1, 0, 0, 0, 0, 0, 0],
+                "node9": [1, 1, 1, 0, 0, 0, 0, 0],
+                "node11": [1, 1, 1, 1, 0, 0, 0, 0],
+                "node13": [1, 1, 1, 1, 1, 0, 0, 0],
+                "node15": [1, 1, 1, 1, 1, 1, 0, 0],
+                "node17": [1, 1, 1, 1, 1, 1, 1, 0],
+                "node18": [
+                    (1, 1, -1),
+                    (1, -1),
+                    (1, -1),
+                    (1, -1),
+                    (1, -1),
+                    (1, -1),
+                    (1, -1),
+                    (1, -1),
+                ],
+                "node5": [2, 0, 0, 0, 0, 0, 0, 0],
+                "node6": [2, 2, 0, 0, 0, 0, 0, 0],
+            },
+            orient="index",
+        )
 
         # A simple balanced binary tree to test lineage pruning
         complete_binary = nx.balanced_tree(2, 3, create_using=nx.DiGraph())
@@ -216,6 +255,13 @@ def test_construction_without_character_matrix(self):
         for n in obs_nodes:
             self.assertIn(n, expected_nodes)
 
+    def test_n_cell_uses_current_character_matrix(self):
+        tree = cas.data.CassiopeiaTree(
+            character_matrix=self.character_matrix, tree=self.test_network
+        )
+        tree.remove_leaf_and_prune_lineage("node3")
+        self.assertEqual(tree.n_cell, 9)
+
     def test_get_children(self):
 
         tree = cas.data.CassiopeiaTree(
@@ -478,6 +524,18 @@ def test_change_branch_length(self):
             CassiopeiaTreeError, tree.set_branch_length, "node12", "node14", -1
         )
 
+    def test_set_character_states_ambiguous(self):
+        tree = cas.data.CassiopeiaTree(
+            character_matrix=self.character_matrix, tree=self.test_network
+        )
+        tree.set_character_states(
+            "node3", [(0, 0), (2,), (1, 0), (1,), (3,), (0, 1), (0,), (0,)]
+        )
+        self.assertEqual(
+            tree.get_character_states("node3"),
+            [(0, 0), (2,), (1, 0), (1,), (3,), (0, 1), (0,), (0,)],
+        )
+
     def test_set_states_with_pandas_dataframe(self):
 
         tree = cas.data.CassiopeiaTree(tree=self.test_network)
@@ -646,18 +704,6 @@ def test_set_dissimilarity_map(self):
             observed_dissimilarity_map, dissimilarity_map
         )
 
-        def delta_fn(
-            x: np.array,
-            y: np.array,
-            missing_state: int,
-            priors: Optional[Dict[int, Dict[int, float]]],
-        ):
-            d = 0
-            for i in range(len(x)):
-                if x[i] != y[i]:
-                    d += 1
-            return d
-
         tree = cas.data.CassiopeiaTree(self.character_matrix)
         tree.compute_dissimilarity_map(delta_fn)
         observed_dissimilarity_map = tree.get_dissimilarity_map()
@@ -700,6 +746,62 @@ def delta_fn(
             CassiopeiaTreeWarning, tree.set_dissimilarity_map, dissimilarity_map
         )
 
+    def test_compute_dissimilarity_map_cluster_dissimilarity(self):
+        tree = cas.data.CassiopeiaTree(self.ambiguous_character_matrix)
+        with self.assertWarns(data_utilities.CassiopeiaTreeWarning):
+            tree.compute_dissimilarity_map(
+                partial(
+                    dissimilarity_functions.cluster_dissimilarity,
+                    delta_fn,
+                    normalize=False,
+                )
+            )
+        observed_dissimilarity_map = tree.get_dissimilarity_map()
+
+        expected_dissimilarity_map = pd.DataFrame.from_dict(
+            {
+                "node3": [0, 1, 2, 3, 4, 5, 6, 7.33333, 1, 2],
+                "node7": [1, 0, 1, 2, 3, 4, 5, 6.83333, 2, 2],
+                "node9": [2, 1, 0, 1, 2, 3, 4, 6.33333, 3, 3],
+                "node11": [3, 2, 1, 0, 1, 2, 3, 5.83333, 4, 4],
+                "node13": [4, 3, 2, 1, 0, 1, 2, 5.33333, 5, 5],
+                "node15": [5, 4, 3, 2, 1, 0, 1, 4.83333, 6, 6],
+                "node17": [6, 5, 4, 3, 2, 1, 0, 4.33333, 7, 7],
+                "node18": [
+                    7.33333,
+                    6.83333,
+                    6.33333,
+                    5.83333,
+                    5.33333,
+                    4.83333,
+                    4.33333,
+                    0,
+                    8,
+                    8,
+                ],
+                "node5": [1, 2, 3, 4, 5, 6, 7, 8, 0, 1],
+                "node6": [2, 2, 3, 4, 5, 6, 7, 8, 1, 0],
+            },
+            orient="index",
+            columns=[
+                "node3",
+                "node7",
+                "node9",
+                "node11",
+                "node13",
+                "node15",
+                "node17",
+                "node18",
+                "node5",
+                "node6",
+            ],
+            dtype=np.float64,
+        )
+
+        pd.testing.assert_frame_equal(
+            observed_dissimilarity_map, expected_dissimilarity_map
+        )
+
     def test_character_matrix_index_not_strings(self):
         cm = pd.DataFrame.from_dict(
             {
@@ -726,17 +828,15 @@ def test_register_data_with_tree(self):
         small_complete_binary_tree = nx.relabel_nodes(
             complete_binary, str_names
         )
-        
-        cas_tree = cas.data.CassiopeiaTree(
-            tree=small_complete_binary_tree
-        )
+
+        cas_tree = cas.data.CassiopeiaTree(tree=small_complete_binary_tree)
 
         cm = pd.DataFrame.from_dict(
             {
                 "node3": [5, 0, 1, 2, -1],
                 "node4": [0, 0, 3, 2, -1],
                 "node5": [-1, 4, 0, 2, 2],
-                "node6": [4, 4, 1, 2, 0]
+                "node6": [4, 4, 1, 2, 0],
             },
             orient="index",
             columns=["a", "b", "c", "d", "e"],
@@ -756,9 +856,75 @@ def test_register_data_with_tree(self):
         cas_tree.set_dissimilarity_map(delta)
         cas_tree.remove_leaf_and_prune_lineage("node5")
 
-        self.assertEqual(set(cas_tree.get_current_character_matrix().index), set(["node3", "node4", "node6"]))
-        self.assertEqual(set(cas_tree.get_dissimilarity_map().index), set(["node3", "node4", "node6"]))
-        self.assertEqual(set(cas_tree.get_dissimilarity_map().columns), set(["node3", "node4", "node6"]))
+        self.assertEqual(
+            set(cas_tree.get_current_character_matrix().index),
+            set(["node3", "node4", "node6"]),
+        )
+        self.assertEqual(
+            set(cas_tree.get_dissimilarity_map().index),
+            set(["node3", "node4", "node6"]),
+        )
+        self.assertEqual(
+            set(cas_tree.get_dissimilarity_map().columns),
+            set(["node3", "node4", "node6"]),
+        )
+
+    def test_register_data_with_tree_added(self):
+        complete_binary = nx.balanced_tree(2, 2, create_using=nx.DiGraph())
+        str_names = dict(
+            zip(
+                list(complete_binary.nodes),
+                ["node" + str(i) for i in complete_binary.nodes],
+            )
+        )
+        small_complete_binary_tree = nx.relabel_nodes(
+            complete_binary, str_names
+        )
+        cas_tree = cas.data.CassiopeiaTree(tree=small_complete_binary_tree)
+        cm = pd.DataFrame.from_dict(
+            {
+                "node3": [5, 0, 1, 2, -1],
+                "node4": [0, 0, 3, 2, -1],
+                "node5": [-1, 4, 0, 2, 2],
+                "node6": [4, 4, 1, 2, 0],
+            },
+            orient="index",
+            columns=["a", "b", "c", "d", "e"],
+        )
+
+        cas_tree.set_character_matrix(cm)
+        delta = pd.DataFrame.from_dict(
+            {
+                "node3": [0, 15, 21, 17],
+                "node4": [15, 0, 10, 6],
+                "node5": [21, 10, 0, 10],
+                "node6": [17, 6, 10, 0],
+            },
+            orient="index",
+            columns=["node3", "node4", "node5", "node6"],
+        )
+        cas_tree.set_dissimilarity_map(delta)
+        cas_tree.add_leaf("node2", "node7")
+
+        self.assertEqual(
+            set(cas_tree.get_current_character_matrix().index),
+            set(["node3", "node4", "node5", "node6", "node7"]),
+        )
+        self.assertEqual(cas_tree.get_character_states("node7"), [-1] * 5)
+        self.assertEqual(
+            set(cas_tree.get_dissimilarity_map().index),
+            set(["node3", "node4", "node5", "node6", "node7"]),
+        )
+        self.assertEqual(
+            set(cas_tree.get_dissimilarity_map().columns),
+            set(["node3", "node4", "node5", "node6", "node7"]),
+        )
+        self.assertEqual(
+            list(cas_tree.get_dissimilarity_map().loc["node7"]), [np.inf] * 5
+        )
+        self.assertEqual(
+            list(cas_tree.get_dissimilarity_map()["node7"]), [np.inf] * 5
+        )
 
     def test_remove_leaf_and_prune_lineage_all(self):
         """Tests the case where all lineages are removed and pruned."""
@@ -809,12 +975,10 @@ def test_remove_leaf_and_prune_lineage_some(self):
             "node8": 4.5,
             "node9": 4.5,
             "node10": 4.5,
-            "node12": 4.5
+            "node12": 4.5,
         }
         for u in cas_tree.nodes:
-            self.assertEqual(
-                cas_tree.get_time(u), expected_times[u]
-            )
+            self.assertEqual(cas_tree.get_time(u), expected_times[u])
 
     def test_remove_leaf_and_prune_lineage_one_side(self):
         """Tests a case where the entire one side of a tree is removed."""
@@ -865,13 +1029,10 @@ def test_collapse_unifurcations_source(self):
             "node0": 0.0,
             "node1": 1.5,
             "node4": 4.5,
-            "node5": 6.0
+            "node5": 6.0,
         }
         for u in cas_tree.nodes:
-            self.assertEqual(
-                cas_tree.get_time(u), expected_times[u]
-            )
-
+            self.assertEqual(cas_tree.get_time(u), expected_times[u])
 
     def test_collapse_unifurcations(self):
         """Tests a general case with unifurcations throughout the tree."""
@@ -922,12 +1083,10 @@ def test_collapse_unifurcations(self):
             "node4": 4.5,
             "node6": 4.5,
             "node7": 6.0,
-            "node8": 6.0
+            "node8": 6.0,
         }
         for u in cas_tree.nodes:
-            self.assertEqual(
-                cas_tree.get_time(u), expected_times[u]
-            )
+            self.assertEqual(cas_tree.get_time(u), expected_times[u])
 
     def test_collapse_unifurcations_long_root_unifurcation(self):
         """Tests a case where there is a long chain at the root."""
@@ -984,12 +1143,10 @@ def test_collapse_unifurcations_long_root_unifurcation(self):
             "node7": 9.0,
             "node8": 9.0,
             "node11": 7.5,
-            "node14": 12.0
+            "node14": 12.0,
         }
         for u in cas_tree.nodes:
-            self.assertEqual(
-                cas_tree.get_time(u), expected_times[u]
-            )
+            self.assertEqual(cas_tree.get_time(u), expected_times[u])
 
     def test_mutationless_edge_collapse_1(self):
         tree = nx.DiGraph()
@@ -1011,16 +1168,19 @@ def test_mutationless_edge_collapse_1(self):
             orient="index",
             columns=["a", "b", "c", "d", "e"],
         )
-        cas_tree = cas.data.CassiopeiaTree(character_matrix=character_matrix, tree=tree)
+        cas_tree = cas.data.CassiopeiaTree(
+            character_matrix=character_matrix, tree=tree
+        )
         cas_tree.collapse_mutationless_edges(infer_ancestral_characters=True)
 
         new_map = {}
         for i in cas_tree.nodes:
             new_map[i] = (
-                "|".join([str(c) for c in cas_tree.get_character_states(i)]) + f",{i}"
+                "|".join([str(c) for c in cas_tree.get_character_states(i)])
+                + f",{i}"
             )
         cas_tree.relabel_nodes(new_map)
-        
+
         expected_nodes = set(
             [
                 "1|0|3|4|5,0",
@@ -1041,7 +1201,7 @@ def test_mutationless_edge_collapse_1(self):
                 ("1|0|3|0|5,6", "1|0|3|0|-1,5"),
             ]
         )
-        
+
         self.assertEqual(set(cas_tree.nodes), expected_nodes)
         self.assertEqual(set(cas_tree.edges), expected_edges)
 
@@ -1066,16 +1226,19 @@ def test_mutationless_edge_collapse_2(self):
             orient="index",
             columns=["a", "b", "c", "d", "e"],
         )
-        cas_tree = cas.data.CassiopeiaTree(character_matrix=character_matrix, tree=tree)
+        cas_tree = cas.data.CassiopeiaTree(
+            character_matrix=character_matrix, tree=tree
+        )
         cas_tree.collapse_mutationless_edges(infer_ancestral_characters=True)
 
         new_map = {}
         for i in cas_tree.nodes:
             new_map[i] = (
-                "|".join([str(c) for c in cas_tree.get_character_states(i)]) + f",{i}"
+                "|".join([str(c) for c in cas_tree.get_character_states(i)])
+                + f",{i}"
             )
         cas_tree.relabel_nodes(new_map)
-        
+
         expected_nodes = set(
             [
                 "1|0|3|4|5,0",
@@ -1097,7 +1260,7 @@ def test_mutationless_edge_collapse_2(self):
 
         self.assertEqual(set(cas_tree.nodes), expected_nodes)
         self.assertEqual(set(cas_tree.edges), expected_edges)
-    
+
     def test_mutationless_edge_collapse_ground_tree(self):
         tree = nx.DiGraph()
         for i in range(7):
@@ -1124,7 +1287,7 @@ def test_mutationless_edge_collapse_ground_tree(self):
             "5": [-1, 1, 0],
             "6": [0, 1, 0],
             "7": [0, 1, 0],
-            "8": [0, 1, 0]
+            "8": [0, 1, 0],
         }
 
         cas_tree.initialize_all_character_states(character_states)
@@ -1133,7 +1296,8 @@ def test_mutationless_edge_collapse_ground_tree(self):
         new_map = {}
         for i in cas_tree.nodes:
             new_map[i] = (
-                "|".join([str(c) for c in cas_tree.get_character_states(i)]) + f",{i}"
+                "|".join([str(c) for c in cas_tree.get_character_states(i)])
+                + f",{i}"
             )
         cas_tree.relabel_nodes(new_map)
 
@@ -1147,7 +1311,9 @@ def test_mutationless_edge_collapse_ground_tree(self):
         }
         self.assertEqual(set(cas_tree.edges), set(expected_edges.keys()))
         for u, v in cas_tree.edges:
-            self.assertEqual(cas_tree.get_branch_length(u,v), expected_edges[(u,v)])
+            self.assertEqual(
+                cas_tree.get_branch_length(u, v), expected_edges[(u, v)]
+            )
 
         expected_times = {
             "0|1|1,4": 3.0,
@@ -1156,7 +1322,7 @@ def test_mutationless_edge_collapse_ground_tree(self):
             "-1|1|0,5": 3.0,
             "-1|1|0,0": 4.5,
             "-1|-1|0,1": 4.5,
-            "0|1|0,8": 0.0
+            "0|1|0,8": 0.0,
         }
         for u in cas_tree.nodes:
             self.assertEqual(cas_tree.get_time(u), expected_times[u])
@@ -1173,7 +1339,9 @@ def test_set_and_add_attribute(self):
         self.assertEqual(5, tree.get_attribute("node3", "test_attribute"))
         self.assertEqual(10, tree.get_attribute("node5", "test_attribute"))
 
-        self.assertRaises(CassiopeiaTreeError, tree.get_attribute, "node10", "test_attribute")
+        self.assertRaises(
+            CassiopeiaTreeError, tree.get_attribute, "node10", "test_attribute"
+        )
 
     def test_filter_nodes(self):
 
@@ -1183,8 +1351,10 @@ def test_filter_nodes(self):
             if tree.is_root(n):
                 tree.set_attribute(n, "depth", 0)
                 continue
-            tree.set_attribute(n, "depth", tree.get_attribute(tree.parent(n), "depth")+1)
-        
+            tree.set_attribute(
+                n, "depth", tree.get_attribute(tree.parent(n), "depth") + 1
+            )
+
         nodes = tree.filter_nodes(lambda x: tree.get_attribute(x, "depth") == 2)
         expected_nodes = ["node5", "node6", "node3", "node4"]
 
@@ -1196,9 +1366,12 @@ def test_get_all_ancestors(self):
 
         ancestors = tree.get_all_ancestors("node3")
         expected_ancestors = ["node1", "node0"]
+        self.assertEqual(ancestors, expected_ancestors)
+
+        ancestors = tree.get_all_ancestors("node3", include_node=True)
+        expected_ancestors = ["node3", "node1", "node0"]
+        self.assertEqual(ancestors, expected_ancestors)
 
-        self.assertCountEqual(expected_ancestors, ancestors)
-        
         ancestors = tree.get_all_ancestors("node0")
         self.assertEqual(0, len(ancestors))
 
@@ -1208,7 +1381,7 @@ def test_find_lcas_of_pairs(self):
 
         pair = [("node3", "node4")]
         lcas = list(tree.find_lcas_of_pairs(pairs=pair))
-        
+
         self.assertEqual("node1", lcas[0][1])
 
         pairs = [("node3", "node4"), ("node5", "node1")]
@@ -1216,12 +1389,262 @@ def test_find_lcas_of_pairs(self):
 
         expected_lcas = {
             ("node3", "node4"): "node1",
-            ("node5", "node1"): "node0"
+            ("node5", "node1"): "node0",
         }
         for lca in lcas:
-            self.assertEqual(expected_lcas[lca[0]], lca[1])   
+            self.assertEqual(expected_lcas[lca[0]], lca[1])
+
+    def test_is_ambiguous(self):
+        tree = cas.data.CassiopeiaTree(
+            character_matrix=self.ambiguous_character_matrix,
+            tree=self.test_network,
+        )
+        self.assertTrue(tree.is_ambiguous("node18"))
+        self.assertFalse(tree.is_ambiguous("node17"))
+
+    def test_collapse_ambiguous_characters(self):
+        # Trees with no ambiguous characters should not be changed
+        tree = cas.data.CassiopeiaTree(
+            character_matrix=self.character_matrix, tree=self.test_network
+        )
+        tree.collapse_ambiguous_characters()
+        pd.testing.assert_frame_equal(
+            tree.get_current_character_matrix(), self.character_matrix
+        )
+
+        ambiguous_tree = cas.data.CassiopeiaTree(
+            character_matrix=self.ambiguous_character_matrix,
+            tree=self.test_network,
+        )
+        ambiguous_tree.collapse_ambiguous_characters()
+        current_character_matrix = ambiguous_tree.get_current_character_matrix()
+        self.assertEqual(
+            list(ambiguous_tree.get_character_states("node18")),
+            [
+                (1, -1),
+                (1, -1),
+                (1, -1),
+                (1, -1),
+                (1, -1),
+                (1, -1),
+                (1, -1),
+                (1, -1),
+            ],
+        )
+        for leaf in ambiguous_tree.leaves:
+            if leaf != "node18":
+                pd.testing.assert_series_equal(
+                    current_character_matrix.loc[leaf],
+                    self.ambiguous_character_matrix.loc[leaf],
+                )
+
+    def test_resolve_ambiguous_characters(self):
+        # Trees with no ambiguous characters should not be changed
+        tree = cas.data.CassiopeiaTree(
+            character_matrix=self.character_matrix, tree=self.test_network
+        )
+        tree.resolve_ambiguous_characters(lambda state: state[0])
+        pd.testing.assert_frame_equal(
+            tree.get_current_character_matrix(), self.character_matrix
+        )
+
+        ambiguous_tree = cas.data.CassiopeiaTree(
+            character_matrix=self.ambiguous_character_matrix,
+            tree=self.test_network,
+        )
+        ambiguous_tree.resolve_ambiguous_characters(lambda state: state[0])
+        current_character_matrix = ambiguous_tree.get_current_character_matrix()
+        self.assertEqual(
+            list(ambiguous_tree.get_character_states("node18")),
+            [1, 1, 1, 1, 1, 1, 1, 1],
+        )
+        for leaf in ambiguous_tree.leaves:
+            if leaf != "node18":
+                pd.testing.assert_series_equal(
+                    current_character_matrix.loc[leaf],
+                    self.ambiguous_character_matrix.loc[leaf].astype(int),
+                )
+
+    def test_get_newick_raises_on_comma(self):
+        tree = cas.data.CassiopeiaTree(tree=self.test_network)
+        tree.relabel_nodes({"node18": "name,with,comma"})
+        with self.assertRaises(CassiopeiaTreeError):
+            tree.get_newick()
+
+    def test_add_leaf(self):
+        tree = cas.data.CassiopeiaTree(
+            character_matrix=self.character_matrix, tree=self.test_network
+        )
+        with self.assertRaises(CassiopeiaTreeError):
+            tree.add_leaf("node18", "new_leaf")
+
+        tree.add_leaf("node16", "new_leaf")
+        self.assertIn("new_leaf", tree.leaves)
+        self.assertIn(("node16", "new_leaf"), tree.edges)
+        self.assertEqual(tree.get_branch_length("node16", "new_leaf"), 0)
+        self.assertEqual(tree.get_time("new_leaf"), tree.get_time("node16"))
+        self.assertEqual(
+            tree.get_character_states("new_leaf"), [-1] * tree.n_character
+        )
+
+    def test_add_leaf_optional_arguments(self):
+        tree = cas.data.CassiopeiaTree(
+            character_matrix=self.character_matrix,
+            tree=self.test_network,
+        )
+        tree.compute_dissimilarity_map(delta_fn)
+        tree.add_leaf(
+            "node16",
+            "new_leaf",
+            states=[1, 2, 3, 4, 5, 6, 7, 8],
+            branch_length=3,
+            dissimilarity={leaf: 3.0 for leaf in tree.leaves},
+        )
+        self.assertIn("new_leaf", tree.leaves)
+        self.assertIn(("node16", "new_leaf"), tree.edges)
+        self.assertEqual(tree.get_branch_length("node16", "new_leaf"), 3)
+        self.assertEqual(tree.get_time("new_leaf"), tree.get_time("node16") + 3)
+        self.assertEqual(
+            tree.get_character_states("new_leaf"), [1, 2, 3, 4, 5, 6, 7, 8]
+        )
+        expected_dissimilarity = {leaf: 3.0 for leaf in tree.leaves}
+        expected_dissimilarity["new_leaf"] = 0.0
+        self.assertEqual(
+            tree.get_dissimilarity("new_leaf"), expected_dissimilarity
+        )
+
+        with self.assertRaises(CassiopeiaTreeError):
+            tree.add_leaf(
+                "node16",
+                "new_leaf2",
+                states=[1, 2, 3, 4, 5, 6, 7, 8],
+                branch_length=3,
+                time=10,
+                dissimilarity={leaf: 3.0 for leaf in tree.leaves},
+            )
+
+    def test_find_lca(self):
+        tree = cas.data.CassiopeiaTree(tree=self.test_network)
+        self.assertEqual("node1", tree.find_lca("node3", "node4"))
+        self.assertEqual("node0", tree.find_lca("node5", "node1"))
+        self.assertEqual("node14", tree.find_lca("node17", "node18", "node15"))
+        self.assertEqual("node16", tree.find_lca("node16", "node18"))
+
+    def test_get_distance(self):
+        tree = cas.data.CassiopeiaTree(tree=self.test_network)
+        self.assertEqual(tree.get_distance("node18", "node16"), 1)
+        self.assertEqual(tree.get_distance("node18", "node17"), 2)
+        self.assertEqual(tree.get_distance("node16", "node9"), 5)
+
+        # Update a branch length and get distance again
+        tree.set_branch_length("node16", "node18", 1.5)
+        self.assertEqual(tree.get_distance("node18", "node16"), 1.5)
+        self.assertEqual(tree.get_distance("node17", "node18"), 2.5)
+
+    def test_get_distances(self):
+        tree = cas.data.CassiopeiaTree(tree=self.test_network)
+        expected_distances = {
+            "node12": 0,
+            # Nodes below
+            "node13": 1,
+            "node14": 1,
+            "node15": 2,
+            "node16": 2,
+            "node17": 3,
+            "node18": 3,
+            # Nodes above
+            "node11": 2,
+            "node10": 1,
+            "node9": 3,
+            "node8": 2,
+            "node7": 4,
+            "node6": 7,
+            "node5": 7,
+            "node4": 3,
+            "node3": 5,
+            "node2": 6,
+            "node1": 4,
+            "node0": 5,
+        }
+        self.assertEqual(tree.get_distances("node12"), expected_distances)
+        self.assertEqual(
+            tree.get_distances("node12", leaves_only=True),
+            {
+                node: distance
+                for node, distance in expected_distances.items()
+                if tree.is_leaf(node)
+            },
+        )
+
+        # Update a branch length and get distance again
+        tree.set_branch_length("node16", "node18", 1.5)
+        expected_distances["node18"] = 3.5
+        self.assertEqual(tree.get_distances("node12"), expected_distances)
+        self.assertEqual(
+            tree.get_distances("node12", leaves_only=True),
+            {
+                node: distance
+                for node, distance in expected_distances.items()
+                if tree.is_leaf(node)
+            },
+        )
+
+    def test_get_dissimilarity(self):
+        dissimilarity_map = pd.DataFrame.from_dict(
+            {
+                "a": [0, 1, 2],
+                "b": [1, 0, 2],
+                "c": [2, 2, 0],
+            },
+            orient="index",
+            columns=["a", "b", "c"],
+        )
+        tree = cas.data.CassiopeiaTree(dissimilarity_map=dissimilarity_map)
+        tree_no_map = cas.data.CassiopeiaTree(
+            character_matrix=self.character_matrix
+        )
+        self.assertEqual(tree.get_dissimilarity("a"), {"a": 0, "b": 1, "c": 2})
+        with self.assertRaises(CassiopeiaTreeError):
+            tree.get_dissimilarity("d")
+        with self.assertRaises(CassiopeiaTreeError):
+            tree_no_map.get_dissimilarity("node3")
+
+    def test_set_dissimilarity(self):
+        dissimilarity_map = pd.DataFrame.from_dict(
+            {
+                "a": [0, 1, 2],
+                "b": [1, 0, 2],
+                "c": [2, 2, 0],
+            },
+            orient="index",
+            columns=["a", "b", "c"],
+        )
+        tree = cas.data.CassiopeiaTree(dissimilarity_map=dissimilarity_map)
+        tree_no_map = cas.data.CassiopeiaTree(
+            character_matrix=self.character_matrix
+        )
+        new_dissimilarity = {"a": 3, "b": 4, "c": 2, "d": 0}
+        tree.set_dissimilarity("d", new_dissimilarity)
+        expected_dissimilarity_map = dissimilarity_map = pd.DataFrame.from_dict(
+            {
+                "a": [0, 1, 2, 3],
+                "b": [1, 0, 2, 4],
+                "c": [2, 2, 0, 2],
+                "d": [3, 4, 2, 0],
+            },
+            orient="index",
+            columns=["a", "b", "c", "d"],
+        )
+        pd.testing.assert_frame_equal(
+            tree.get_dissimilarity_map(), expected_dissimilarity_map
+        )
+        with self.assertRaises(CassiopeiaTreeError):
+            tree.set_dissimilarity("c", {"a": 3})
+        with self.assertRaises(CassiopeiaTreeError):
+            tree_no_map.set_dissimilarity(
+                "node3", {leaf: 0 for leaf in tree_no_map.leaves}
+            )
 
 
 if __name__ == "__main__":
     unittest.main()
-
diff --git a/test/data_tests/utilities_test.py b/test/data_tests/utilities_test.py
old mode 100644
new mode 100755
index dac870fd..9e59c5ca
--- a/test/data_tests/utilities_test.py
+++ b/test/data_tests/utilities_test.py
@@ -274,11 +274,11 @@ def test_bootstrap_allele_tables_priors(self):
     def test_to_newick_no_branch_lengths(self):
         tree = nx.DiGraph()
         tree.add_nodes_from(["A", "B", "C", "D", "E", "F"])
-        tree.add_edge("F","A",length=0.1)
-        tree.add_edge("F","B",length=0.2)
-        tree.add_edge("F","E",length=0.5)
-        tree.add_edge("E","C",length=0.3)
-        tree.add_edge("E","D",length=0.4)
+        tree.add_edge("F", "A", length=0.1)
+        tree.add_edge("F", "B", length=0.2)
+        tree.add_edge("F", "E", length=0.5)
+        tree.add_edge("E", "C", length=0.3)
+        tree.add_edge("E", "D", length=0.4)
 
         newick_string = data_utilities.to_newick(tree)
         self.assertEqual(newick_string, "(A,B,(C,D));")
@@ -286,13 +286,15 @@ def test_to_newick_no_branch_lengths(self):
     def test_to_newick_branch_lengths(self):
         tree = nx.DiGraph()
         tree.add_nodes_from(["A", "B", "C", "D", "E", "F"])
-        tree.add_edge("F","A",length=0.1)
-        tree.add_edge("F","B",length=0.2)
-        tree.add_edge("F","E",length=0.5)
-        tree.add_edge("E","C",length=0.3)
-        tree.add_edge("E","D",length=0.4)
-
-        newick_string = data_utilities.to_newick(tree, record_branch_lengths = True)
+        tree.add_edge("F", "A", length=0.1)
+        tree.add_edge("F", "B", length=0.2)
+        tree.add_edge("F", "E", length=0.5)
+        tree.add_edge("E", "C", length=0.3)
+        tree.add_edge("E", "D", length=0.4)
+
+        newick_string = data_utilities.to_newick(
+            tree, record_branch_lengths=True
+        )
         self.assertEqual(newick_string, "(A:0.1,B:0.2,(C:0.3,D:0.4):0.5);")
 
     def test_lca_characters(self):
@@ -302,6 +304,25 @@ def test_lca_characters(self):
         )
         self.assertEqual(ret_vec, [1, 0, 3, 0, 5])
 
+    def test_lca_characters_ambiguous(self):
+        vecs = [
+            [(1, 1), (0, 2), (3,), (4,), (5,)],
+            [1, -1, -1, 3, -1],
+            [1, 2, 3, 2, -1],
+        ]
+        ret_vec = data_utilities.get_lca_characters(
+            vecs, missing_state_indicator=-1
+        )
+        self.assertEqual(ret_vec, [1, 0, 3, 0, 5])
+
+    def test_is_ambiguous_state(self):
+        self.assertTrue(data_utilities.is_ambiguous_state((1, 2)))
+        self.assertFalse(data_utilities.is_ambiguous_state(1))
+
+    def test_resolve_most_abundant(self):
+        state = (1, 2, 3, 3)
+        self.assertEqual(data_utilities.resolve_most_abundant(state), 3)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/preprocess_tests/collapse_umi_test.py b/test/preprocess_tests/collapse_umi_test.py
index feff1956..010dc304 100644
--- a/test/preprocess_tests/collapse_umi_test.py
+++ b/test/preprocess_tests/collapse_umi_test.py
@@ -30,14 +30,42 @@ def setUp(self):
             ".collapsed.bam"
         )
 
-        _, _ = UMI_utils.sort_cellranger_bam(
-            self.test_file, str(self.sorted_file_name)
-        )
+        _, _ = UMI_utils.sort_bam(self.test_file, str(self.sorted_file_name))
 
         UMI_utils.form_collapsed_clusters(
             str(self.sorted_file_name), max_hq_mismatches=3, max_indels=2
         )
 
+        self.uncorrected_test_file = (
+            dir_path + "/test_files/test_uncorrected.bam"
+        )
+        uncorrected_sorted_file_name = Path(
+            dir_path
+            + "/test_files/"
+            + ".".join(
+                self.uncorrected_test_file.split("/")[-1].split(".")[:-1]
+            )
+            + "_sorted.bam"
+        )
+        self.uncorrected_sorted_file_name = uncorrected_sorted_file_name
+        self.uncorrected_collapsed_file_name = (
+            uncorrected_sorted_file_name.with_suffix(".collapsed.bam")
+        )
+
+        _, _ = UMI_utils.sort_bam(
+            self.uncorrected_test_file,
+            str(self.uncorrected_sorted_file_name),
+            sort_key=lambda al: (al.get_tag("CR"), al.get_tag("UR")),
+            filter_func=lambda al: al.has_tag("CR"),
+        )
+
+        UMI_utils.form_collapsed_clusters(
+            str(self.uncorrected_sorted_file_name),
+            max_hq_mismatches=3,
+            max_indels=2,
+            cell_key=lambda al: al.get_tag("CR"),
+        )
+
     def test_sort_bam(self):
 
         sorted_bam = pysam.AlignmentFile(
@@ -57,6 +85,25 @@ def test_sort_bam(self):
         self.assertEqual(expected_cellBC, cellBCs[10])
         self.assertEqual(expected_UMI, UMIs[7])
 
+    def test_sort_bam_uncorrected(self):
+
+        sorted_bam = pysam.AlignmentFile(
+            self.uncorrected_sorted_file_name, "rb", check_sq=False
+        )
+
+        cellBCs = []
+        UMIs = []
+        for al in sorted_bam:
+            cellBCs.append(al.get_tag("CR"))
+            UMIs.append(al.get_tag("UR"))
+
+        expected_cellBC = "CCGGATAGAAAGTGGA"
+        expected_UMI = "GATAACATCG"
+
+        self.assertEqual(len(cellBCs), 17)
+        self.assertEqual(expected_cellBC, cellBCs[10])
+        self.assertEqual(expected_UMI, UMIs[7])
+
     def test_collapse_bam(self):
         collapsed_bam = pysam.AlignmentFile(
             self.collapsed_file_name, "rb", check_sq=False
@@ -83,6 +130,32 @@ def test_collapse_bam(self):
         self.assertEqual([7, 1, 2, 3, 3], readCounts)
         self.assertEqual(2, quals[2][0])
 
+    def test_collapse_bam_uncorrected(self):
+        collapsed_bam = pysam.AlignmentFile(
+            self.uncorrected_collapsed_file_name, "rb", check_sq=False
+        )
+
+        cellBCs = []
+        UMIs = []
+        readCounts = []
+        clusterIds = []
+        quals = []
+        for al in collapsed_bam:
+            cellBCs.append(al.get_tag("CB"))
+            UMIs.append(al.get_tag("UR"))
+            readCounts.append(al.get_tag("ZR"))
+            clusterIds.append(al.get_tag("ZC"))
+            quals.append(al.query_qualities)
+
+        expected_cellBC = "CCGGATAGAAAGTGGA"
+        expected_UMI = "GGCAGTAATT"
+
+        self.assertEqual(len(cellBCs), 4)
+        self.assertEqual(expected_cellBC, cellBCs[0])
+        self.assertEqual(expected_UMI, UMIs[1])
+        self.assertEqual([13, 1, 1, 2], readCounts)
+        self.assertEqual(37, quals[2][0])
+
     def test_bam2DF(self):
         collapsed_df_file_name = self.sorted_file_name.with_suffix(
             ".collapsed.txt"
diff --git a/test/preprocess_tests/config_parser_test.py b/test/preprocess_tests/config_parser_test.py
old mode 100644
new mode 100755
index c214a074..cf269bf6
--- a/test/preprocess_tests/config_parser_test.py
+++ b/test/preprocess_tests/config_parser_test.py
@@ -20,7 +20,7 @@ def setUp(self):
             "general": {
                 "output_directory": "'here'",
                 "reference_filepath": "'ref.fa'",
-                "input_file": "'input.txt'",
+                "input_files": ["input1.txt", "input2.txt"],
             },
             "collapse": {"max_hq_mismatches": 5},
             "call_alleles": {"barcode_interval": (10, 25)},
@@ -36,7 +36,7 @@ def setUp(self):
             "general": {
                 "output_directory": "'here'",
                 "reference_filepath": "'ref.fa'",
-                "input_file": "'input.txt'",
+                "input_files": ["input.txt"],
                 "entry": "'align'",
                 "exit": "'filter_molecule_table'",
             },
@@ -66,19 +66,21 @@ def test_read_good_config(self):
         parameters = setup_utilities.parse_config(self.basic_config_string)
 
         # check some default parameters
-        self.assertEqual(parameters["error_correct"]["max_umi_distance"], 2)
+        self.assertEqual(parameters["error_correct_umis"]["max_umi_distance"], 2)
         self.assertEqual(parameters["align"]["gap_open_penalty"], 20)
         self.assertEqual(parameters["call_alleles"]["cutsite_width"], 12)
 
         # check parameters updated correctly
         self.assertEqual(parameters["general"]["output_directory"], "here")
         self.assertEqual(parameters["general"]["reference_filepath"], "ref.fa")
-        self.assertEqual(parameters["general"]["input_file"], "input.txt")
+        self.assertEqual(parameters["general"]["input_files"], ["input1.txt", "input2.txt"])
         self.assertEqual(parameters["collapse"]["max_hq_mismatches"], 5)
         self.assertEqual(
             parameters["call_alleles"]["barcode_interval"], (10, 25)
         )
 
+        self.assertIn("output_directory", parameters["convert"].keys())
+        self.assertIn("output_directory", parameters["error_correct_barcodes"].keys())
         self.assertIn("output_directory", parameters["collapse"].keys())
         self.assertIn("output_directory", parameters["resolve"].keys())
         self.assertIn(
@@ -106,11 +108,13 @@ def test_pipeline_setup_correct(self):
         )
 
         expected_procedures = [
+            "convert",
+            "error_correct_barcodes",
             "collapse",
             "resolve",
             "align",
             "call_alleles",
-            "error_correct",
+            "error_correct_umis",
             "filter_molecule_table",
             "call_lineages",
         ]
@@ -133,7 +137,7 @@ def test_subset_pipeline_setup_correct(self):
         expected_procedures = [
             "align",
             "call_alleles",
-            "error_correct",
+            "error_correct_umis",
             "filter_molecule_table",
         ]
 
diff --git a/test/preprocess_tests/convert_fastqs_to_unmapped_bam_test.py b/test/preprocess_tests/convert_fastqs_to_unmapped_bam_test.py
new file mode 100755
index 00000000..c0ebc4f9
--- /dev/null
+++ b/test/preprocess_tests/convert_fastqs_to_unmapped_bam_test.py
@@ -0,0 +1,303 @@
+"""
+Tests for converting FASTQs to an unmapped BAM in pipeline.py
+"""
+import os
+import unittest
+import tempfile
+
+import pysam
+import ngs_tools as ngs
+
+from cassiopeia.preprocess import pipeline
+
+
+class TestConvertFastqsToUnmappedBam(unittest.TestCase):
+    def setUp(self):
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        test_files_path = os.path.join(dir_path, "test_files")
+
+        self.fastq_10xv3_fps = [
+            os.path.join(test_files_path, "10xv3_1.fastq.gz"),
+            os.path.join(test_files_path, "10xv3_2.fastq.gz"),
+        ]
+        self.fastq_indropsv3_fps = [
+            os.path.join(test_files_path, "indropsv3_1.fastq.gz"),
+            os.path.join(test_files_path, "indropsv3_2.fastq.gz"),
+            os.path.join(test_files_path, "indropsv3_3.fastq.gz"),
+        ]
+        self.fastq_slideseq2_fps = [
+            os.path.join(test_files_path, "slideseq2_1.fastq.gz"),
+            os.path.join(test_files_path, "slideseq2_2.fastq.gz"),
+        ]
+
+    def test_dropseq(self):
+        # NOTE: using 10xv3 fastqs just for testing
+        bam_fp = pipeline.convert_fastqs_to_unmapped_bam(
+            self.fastq_10xv3_fps, "dropseq", tempfile.mkdtemp(), name="test"
+        )
+        with pysam.AlignmentFile(bam_fp, "rb", check_sq=False) as f:
+            alignments = list(f.fetch(until_eof=True))
+        self.assertEqual(2, len(alignments))
+        self.assertEqual(
+            [
+                "M03718:773:000000000-JKHP3:1:1101:18272:1693",
+                "M03718:773:000000000-JKHP3:1:1101:17963:1710",
+            ],
+            [al.query_name for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.sequence
+                for read in ngs.fastq.Fastq(self.fastq_10xv3_fps[1])
+            ],
+            [al.query_sequence for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.qualities.string
+                for read in ngs.fastq.Fastq(self.fastq_10xv3_fps[1])
+            ],
+            [
+                pysam.array_to_qualitystring(al.query_qualities)
+                for al in alignments
+            ],
+        )
+        self.assertEqual(
+            {
+                ("UR", "TACGCCAA"),
+                ("UY", "GGFECEE0"),
+                ("CR", "TACGTCATCTCC"),
+                ("CY", "1111AFAFFFBF"),
+                ("RG", "test"),
+            },
+            set(alignments[0].get_tags()),
+        )
+        self.assertEqual(
+            {
+                ("UR", "AAACATTC"),
+                ("UY", "FFGGBFGF"),
+                ("CR", "TTAGATCGTTAG"),
+                ("CY", "1>>11DFAFAAA"),
+                ("RG", "test"),
+            },
+            set(alignments[1].get_tags()),
+        )
+
+    def test_10xv2(self):
+        # NOTE: using 10xv3 fastqs just for testing
+        bam_fp = pipeline.convert_fastqs_to_unmapped_bam(
+            self.fastq_10xv3_fps, "10xv2", tempfile.mkdtemp(), name="test"
+        )
+        with pysam.AlignmentFile(bam_fp, "rb", check_sq=False) as f:
+            alignments = list(f.fetch(until_eof=True))
+        self.assertEqual(2, len(alignments))
+        self.assertEqual(
+            [
+                "M03718:773:000000000-JKHP3:1:1101:18272:1693",
+                "M03718:773:000000000-JKHP3:1:1101:17963:1710",
+            ],
+            [al.query_name for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.sequence
+                for read in ngs.fastq.Fastq(self.fastq_10xv3_fps[1])
+            ],
+            [al.query_sequence for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.qualities.string
+                for read in ngs.fastq.Fastq(self.fastq_10xv3_fps[1])
+            ],
+            [
+                pysam.array_to_qualitystring(al.query_qualities)
+                for al in alignments
+            ],
+        )
+        self.assertEqual(
+            {
+                ("UR", "CCAAAACAGT"),
+                ("UY", "CEE0C0BA0D"),
+                ("CR", "TACGTCATCTCCTACG"),
+                ("CY", "1111AFAFFFBFGGFE"),
+                ("RG", "test"),
+            },
+            set(alignments[0].get_tags()),
+        )
+        self.assertEqual(
+            {
+                ("UR", "ATTCCTGAGT"),
+                ("UY", "BFGFGFF10F"),
+                ("CR", "TTAGATCGTTAGAAAC"),
+                ("CY", "1>>11DFAFAAAFFGG"),
+                ("RG", "test"),
+            },
+            set(alignments[1].get_tags()),
+        )
+
+    def test_10xv3(self):
+        bam_fp = pipeline.convert_fastqs_to_unmapped_bam(
+            self.fastq_10xv3_fps, "10xv3", tempfile.mkdtemp(), name="test"
+        )
+        with pysam.AlignmentFile(bam_fp, "rb", check_sq=False) as f:
+            alignments = list(f.fetch(until_eof=True))
+        self.assertEqual(2, len(alignments))
+        self.assertEqual(
+            [
+                "M03718:773:000000000-JKHP3:1:1101:18272:1693",
+                "M03718:773:000000000-JKHP3:1:1101:17963:1710",
+            ],
+            [al.query_name for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.sequence
+                for read in ngs.fastq.Fastq(self.fastq_10xv3_fps[1])
+            ],
+            [al.query_sequence for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.qualities.string
+                for read in ngs.fastq.Fastq(self.fastq_10xv3_fps[1])
+            ],
+            [
+                pysam.array_to_qualitystring(al.query_qualities)
+                for al in alignments
+            ],
+        )
+        self.assertEqual(
+            {
+                ("UR", "CCAAAACAGTTT"),
+                ("UY", "CEE0C0BA0DFG"),
+                ("CR", "TACGTCATCTCCTACG"),
+                ("CY", "1111AFAFFFBFGGFE"),
+                ("RG", "test"),
+            },
+            set(alignments[0].get_tags()),
+        )
+        self.assertEqual(
+            {
+                ("UR", "ATTCCTGAGTCA"),
+                ("UY", "BFGFGFF10FG1"),
+                ("CR", "TTAGATCGTTAGAAAC"),
+                ("CY", "1>>11DFAFAAAFFGG"),
+                ("RG", "test"),
+            },
+            set(alignments[1].get_tags()),
+        )
+
+    def test_indropsv3(self):
+        bam_fp = pipeline.convert_fastqs_to_unmapped_bam(
+            self.fastq_indropsv3_fps,
+            "indropsv3",
+            tempfile.mkdtemp(),
+            name="test",
+        )
+        with pysam.AlignmentFile(bam_fp, "rb", check_sq=False) as f:
+            alignments = list(f.fetch(until_eof=True))
+        self.assertEqual(2, len(alignments))
+        self.assertEqual(
+            [
+                "M03718:773:000000000-JKHP3:1:1101:18272:1693",
+                "M03718:773:000000000-JKHP3:1:1101:17963:1710",
+            ],
+            [al.query_name for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.sequence
+                for read in ngs.fastq.Fastq(self.fastq_10xv3_fps[1])
+            ],
+            [al.query_sequence for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.qualities.string
+                for read in ngs.fastq.Fastq(self.fastq_10xv3_fps[1])
+            ],
+            [
+                pysam.array_to_qualitystring(al.query_qualities)
+                for al in alignments
+            ],
+        )
+        self.assertEqual(
+            {
+                ("UR", "CCAAAA"),
+                ("UY", "FFBFGG"),
+                ("CR", "TACGTCATCTCCTACG"),
+                ("CY", "1111AFAF1111AFAF"),
+                ("RG", "test"),
+            },
+            set(alignments[0].get_tags()),
+        )
+        self.assertEqual(
+            {
+                ("UR", "TTAGAA"),
+                ("UY", "FAAAFF"),
+                ("CR", "TTAGATCGTTAGATCG"),
+                ("CY", "1>>11DFA1>>11DFA"),
+                ("RG", "test"),
+            },
+            set(alignments[1].get_tags()),
+        )
+
+    def test_slideseq2(self):
+        bam_fp = pipeline.convert_fastqs_to_unmapped_bam(
+            self.fastq_slideseq2_fps,
+            "slideseq2",
+            tempfile.mkdtemp(),
+            name="test",
+        )
+        with pysam.AlignmentFile(bam_fp, "rb", check_sq=False) as f:
+            alignments = list(f.fetch(until_eof=True))
+        self.assertEqual(2, len(alignments))
+        self.assertEqual(
+            [
+                "NB501583:801:H7JLTBGXH:1:11101:20912:1050",
+                "NB501583:801:H7JLTBGXH:1:11101:8670:1050",
+            ],
+            [al.query_name for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.sequence
+                for read in ngs.fastq.Fastq(self.fastq_slideseq2_fps[1])
+            ],
+            [al.query_sequence for al in alignments],
+        )
+        self.assertEqual(
+            [
+                read.qualities.string
+                for read in ngs.fastq.Fastq(self.fastq_slideseq2_fps[1])
+            ],
+            [
+                pysam.array_to_qualitystring(al.query_qualities)
+                for al in alignments
+            ],
+        )
+        self.assertEqual(
+            {
+                ("UR", "TTTTTTTTT"),
+                ("UY", "EEEEEEEEE"),
+                ("CR", "CTTTGNTCAATGTT"),
+                ("CY", "AAAAA#EEAEEEEE"),
+                ("RG", "test"),
+            },
+            set(alignments[0].get_tags()),
+        )
+        self.assertEqual(
+            {
+                ("UR", "AGTGTCTCA"),
+                ("UY", "EAEAEAEEE"),
+                ("CR", "CTCTTNATCCTCAT"),
+                ("CY", "AAAAA#EEE/EAE/"),
+                ("RG", "test"),
+            },
+            set(alignments[1].get_tags()),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/preprocess_tests/error_correct_barcodes_test.py b/test/preprocess_tests/error_correct_barcodes_test.py
new file mode 100644
index 00000000..a598862b
--- /dev/null
+++ b/test/preprocess_tests/error_correct_barcodes_test.py
@@ -0,0 +1,50 @@
+"""
+Tests for correcting raw barcodes to a whitelist pipeline.py
+"""
+import os
+import unittest
+import tempfile
+
+import pysam
+import ngs_tools as ngs
+
+from cassiopeia.preprocess import pipeline
+
+
+class TestErrorCorrectBarcodes(unittest.TestCase):
+    def setUp(self):
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        test_files_path = os.path.join(dir_path, "test_files")
+
+        self.bam_10xv3_fp = os.path.join(test_files_path, "10xv3_unmapped.bam")
+        self.whitelist_10xv3_fp = os.path.join(test_files_path, "10xv3_whitelist.txt")
+        self.bam_slideseq2_fp = os.path.join(
+            test_files_path, "slideseq2_unmapped.bam"
+        )
+        self.whitelist_slideseq2_fp = os.path.join(test_files_path, "slideseq2_whitelist.txt")
+
+    def test_10xv3(self):
+        bam_fp = pipeline.error_correct_barcodes(
+            self.bam_10xv3_fp, tempfile.mkdtemp(), self.whitelist_10xv3_fp
+        )
+        with pysam.AlignmentFile(bam_fp, "rb", check_sq=False) as f:
+            alignments = list(f.fetch(until_eof=True))
+        self.assertEqual(2, len(alignments))
+        self.assertEqual(
+            ["TACGTCATCTCCTACG", "TTAGATCGTTAGAAAG"],
+            [al.get_tag("CB") for al in alignments],
+        )
+
+    def test_slideseq2(self):
+        bam_fp = pipeline.error_correct_barcodes(
+            self.bam_slideseq2_fp, tempfile.mkdtemp(), self.whitelist_slideseq2_fp
+        )
+        with pysam.AlignmentFile(bam_fp, "rb", check_sq=False) as f:
+            alignments = list(f.fetch(until_eof=True))
+        self.assertEqual(2, len(alignments))
+        self.assertEqual([True, False], [al.has_tag("CB") for al in alignments])
+        self.assertEqual("CTTTGNTCAAAGTT", alignments[0].get_tag("CB"))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/preprocess_tests/error_correct_umi_test.py b/test/preprocess_tests/error_correct_umi_test.py
old mode 100644
new mode 100755
index de2f59d6..b820262f
--- a/test/preprocess_tests/error_correct_umi_test.py
+++ b/test/preprocess_tests/error_correct_umi_test.py
@@ -109,7 +109,7 @@ def setUp(self):
     def test_format(self):
 
         aln_df = cassiopeia.pp.error_correct_umis(
-            self.multi_case, "test", max_umi_distance=1
+            self.multi_case, max_umi_distance=1
         )
 
         expected_columns = [
@@ -131,7 +131,7 @@ def test_format(self):
     def test_zero_dist(self):
 
         aln_df = cassiopeia.pp.error_correct_umis(
-            self.multi_case, "test", max_umi_distance=0
+            self.multi_case, max_umi_distance=0
         )
 
         self.assertEqual(aln_df.shape[0], self.multi_case.shape[0])
@@ -142,7 +142,7 @@ def test_zero_dist(self):
     def test_error_correct_two_dist(self):
 
         aln_df = cassiopeia.pp.error_correct_umis(
-            self.multi_case, "test", max_umi_distance=2
+            self.multi_case, max_umi_distance=2
         )
 
         expected_alignments = {
diff --git a/test/preprocess_tests/test_files/10xv3_1.fastq.gz b/test/preprocess_tests/test_files/10xv3_1.fastq.gz
new file mode 100644
index 00000000..50dc39d1
Binary files /dev/null and b/test/preprocess_tests/test_files/10xv3_1.fastq.gz differ
diff --git a/test/preprocess_tests/test_files/10xv3_2.fastq.gz b/test/preprocess_tests/test_files/10xv3_2.fastq.gz
new file mode 100644
index 00000000..8d9a7602
Binary files /dev/null and b/test/preprocess_tests/test_files/10xv3_2.fastq.gz differ
diff --git a/test/preprocess_tests/test_files/10xv3_unmapped.bam b/test/preprocess_tests/test_files/10xv3_unmapped.bam
new file mode 100644
index 00000000..f85c41e0
Binary files /dev/null and b/test/preprocess_tests/test_files/10xv3_unmapped.bam differ
diff --git a/test/preprocess_tests/test_files/10xv3_whitelist.txt b/test/preprocess_tests/test_files/10xv3_whitelist.txt
new file mode 100755
index 00000000..54df9de6
--- /dev/null
+++ b/test/preprocess_tests/test_files/10xv3_whitelist.txt
@@ -0,0 +1,2 @@
+TACGTCATCTCCTACG
+TTAGATCGTTAGAAAG
diff --git a/test/preprocess_tests/test_files/indropsv3_1.fastq.gz b/test/preprocess_tests/test_files/indropsv3_1.fastq.gz
new file mode 100644
index 00000000..2d6ee2d3
Binary files /dev/null and b/test/preprocess_tests/test_files/indropsv3_1.fastq.gz differ
diff --git a/test/preprocess_tests/test_files/indropsv3_2.fastq.gz b/test/preprocess_tests/test_files/indropsv3_2.fastq.gz
new file mode 100644
index 00000000..9bb0f6ca
Binary files /dev/null and b/test/preprocess_tests/test_files/indropsv3_2.fastq.gz differ
diff --git a/test/preprocess_tests/test_files/indropsv3_3.fastq.gz b/test/preprocess_tests/test_files/indropsv3_3.fastq.gz
new file mode 100644
index 00000000..8d9a7602
Binary files /dev/null and b/test/preprocess_tests/test_files/indropsv3_3.fastq.gz differ
diff --git a/test/preprocess_tests/test_files/slideseq2_1.fastq.gz b/test/preprocess_tests/test_files/slideseq2_1.fastq.gz
new file mode 100644
index 00000000..1de489c2
Binary files /dev/null and b/test/preprocess_tests/test_files/slideseq2_1.fastq.gz differ
diff --git a/test/preprocess_tests/test_files/slideseq2_2.fastq.gz b/test/preprocess_tests/test_files/slideseq2_2.fastq.gz
new file mode 100644
index 00000000..9e08bf57
Binary files /dev/null and b/test/preprocess_tests/test_files/slideseq2_2.fastq.gz differ
diff --git a/test/preprocess_tests/test_files/slideseq2_unmapped.bam b/test/preprocess_tests/test_files/slideseq2_unmapped.bam
new file mode 100644
index 00000000..d2579f00
Binary files /dev/null and b/test/preprocess_tests/test_files/slideseq2_unmapped.bam differ
diff --git a/test/preprocess_tests/test_files/slideseq2_whitelist.txt b/test/preprocess_tests/test_files/slideseq2_whitelist.txt
new file mode 100755
index 00000000..1332df67
--- /dev/null
+++ b/test/preprocess_tests/test_files/slideseq2_whitelist.txt
@@ -0,0 +1,1 @@
+CTTTGNTCAAAGTT
diff --git a/test/preprocess_tests/test_files/test_uncorrected.bam b/test/preprocess_tests/test_files/test_uncorrected.bam
new file mode 100644
index 00000000..9409be3e
Binary files /dev/null and b/test/preprocess_tests/test_files/test_uncorrected.bam differ
diff --git a/test/preprocess_tests/test_files/test_uncorrected_sorted.bam b/test/preprocess_tests/test_files/test_uncorrected_sorted.bam
new file mode 100644
index 00000000..dc56b714
Binary files /dev/null and b/test/preprocess_tests/test_files/test_uncorrected_sorted.bam differ
diff --git a/test/preprocess_tests/test_files/test_uncorrected_sorted.collapsed.bam b/test/preprocess_tests/test_files/test_uncorrected_sorted.collapsed.bam
new file mode 100644
index 00000000..4a323bf8
Binary files /dev/null and b/test/preprocess_tests/test_files/test_uncorrected_sorted.collapsed.bam differ
diff --git a/test/simulator_tests/cas9_lineage_tracing_simulator_test.py b/test/simulator_tests/cas9_lineage_tracing_simulator_test.py
index a3892ca8..29ac0d37 100644
--- a/test/simulator_tests/cas9_lineage_tracing_simulator_test.py
+++ b/test/simulator_tests/cas9_lineage_tracing_simulator_test.py
@@ -58,47 +58,57 @@ def setUp(self):
 
         self.basic_tree = tree
 
-        self.basic_lineage_tracing_data_simulator = cas.sim.Cas9LineageTracingDataSimulator(
-            number_of_cassettes=3,
-            size_of_cassette=3,
-            mutation_rate=0.3,
-            state_priors={1: 0.1, 2: 0.1, 3: 0.75, 4: 0.05},
-            heritable_silencing_rate=1e-5,
-            stochastic_silencing_rate=1e-2,
-            random_seed=123412232,
+        self.basic_lineage_tracing_data_simulator = (
+            cas.sim.Cas9LineageTracingDataSimulator(
+                number_of_cassettes=3,
+                size_of_cassette=3,
+                mutation_rate=0.3,
+                state_priors={1: 0.1, 2: 0.1, 3: 0.75, 4: 0.05},
+                heritable_silencing_rate=1e-5,
+                stochastic_silencing_rate=1e-2,
+                random_seed=123412232,
+            )
         )
 
-        self.basic_lineage_tracing_no_resection = cas.sim.Cas9LineageTracingDataSimulator(
-            number_of_cassettes=9,
-            size_of_cassette=1,
-            mutation_rate=0.3,
-            state_priors={1: 0.1, 2: 0.1, 3: 0.75, 4: 0.05},
-            heritable_silencing_rate=0,
-            stochastic_silencing_rate=0,
-            random_seed=123412232,
+        self.basic_lineage_tracing_no_resection = (
+            cas.sim.Cas9LineageTracingDataSimulator(
+                number_of_cassettes=9,
+                size_of_cassette=1,
+                mutation_rate=0.3,
+                state_priors={1: 0.1, 2: 0.1, 3: 0.75, 4: 0.05},
+                heritable_silencing_rate=0,
+                stochastic_silencing_rate=0,
+                random_seed=123412232,
+            )
         )
 
-        self.lineage_tracing_data_simulator_state_distribution = cas.sim.Cas9LineageTracingDataSimulator(
-            number_of_cassettes=3,
-            size_of_cassette=3,
-            mutation_rate=0.3,
-            state_generating_distribution=lambda: np.random.exponential(1e-5),
-            number_of_states = 10,
-            heritable_silencing_rate=1e-5,
-            stochastic_silencing_rate=1e-2,
-            random_seed=123412232,
+        self.lineage_tracing_data_simulator_state_distribution = (
+            cas.sim.Cas9LineageTracingDataSimulator(
+                number_of_cassettes=3,
+                size_of_cassette=3,
+                mutation_rate=0.3,
+                state_generating_distribution=lambda: np.random.exponential(
+                    1e-5
+                ),
+                number_of_states=10,
+                heritable_silencing_rate=1e-5,
+                stochastic_silencing_rate=1e-2,
+                random_seed=123412232,
+            )
         )
 
-        self.lineage_tracing_data_simulator_missing_data = cas.sim.Cas9LineageTracingDataSimulator(
-            number_of_cassettes=3,
-            size_of_cassette=3,
-            mutation_rate=0.3,
-            state_priors={1: 0.1, 2: 0.1, 3: 0.75, 4: 0.05},
-            heritable_silencing_rate=1e-5,
-            stochastic_silencing_rate=1e-2,
-            heritable_missing_data_state=-2,
-            stochastic_missing_data_state=-1,
-            random_seed=123412232,
+        self.lineage_tracing_data_simulator_missing_data = (
+            cas.sim.Cas9LineageTracingDataSimulator(
+                number_of_cassettes=3,
+                size_of_cassette=3,
+                mutation_rate=0.3,
+                state_priors={1: 0.1, 2: 0.1, 3: 0.75, 4: 0.05},
+                heritable_silencing_rate=1e-5,
+                stochastic_silencing_rate=1e-2,
+                heritable_missing_data_state=-2,
+                stochastic_missing_data_state=-1,
+                random_seed=123412232,
+            )
         )
 
     def test_basic_setup(self):
@@ -213,8 +223,10 @@ def test_introduce_states(self):
         np.random.seed(123412232)
 
         character_array = [0, 0, 0, 0, 0, 0, 0, 0, 0]
-        updated_character_array = self.basic_lineage_tracing_data_simulator.introduce_states(
-            character_array, [0, 3, 5, 6]
+        updated_character_array = (
+            self.basic_lineage_tracing_data_simulator.introduce_states(
+                character_array, [0, 3, 5, 6]
+            )
         )
 
         expected_character_array = [3, 0, 0, 3, 0, 3, 3, 0, 0]
@@ -231,8 +243,10 @@ def test_silence_cassettes(self):
 
         character_array = [0, 0, 0, 0, 0, 0, 0, 0, 0]
 
-        updated_character_array = self.basic_lineage_tracing_data_simulator.silence_cassettes(
-            character_array, 0.1
+        updated_character_array = (
+            self.basic_lineage_tracing_data_simulator.silence_cassettes(
+                character_array, 0.1
+            )
         )
 
         expected_character_array = [0, 0, 0, 0, 0, 0, -1, -1, -1]
@@ -246,7 +260,10 @@ def test_collapse_sites(self):
 
         character_array = [0, 0, 0, 0, 0, 0, 0, 0, 0]
 
-        updated_character_array, remaining_cuts = self.basic_lineage_tracing_data_simulator.collapse_sites(
+        (
+            updated_character_array,
+            remaining_cuts,
+        ) = self.basic_lineage_tracing_data_simulator.collapse_sites(
             character_array, [0, 1]
         )
         self.assertCountEqual([], remaining_cuts)
@@ -256,7 +273,10 @@ def test_collapse_sites(self):
                 expected_character_array[i], updated_character_array[i]
             )
 
-        updated_character_array, remaining_cuts = self.basic_lineage_tracing_data_simulator.collapse_sites(
+        (
+            updated_character_array,
+            remaining_cuts,
+        ) = self.basic_lineage_tracing_data_simulator.collapse_sites(
             character_array, [0, 5, 9]
         )
         self.assertCountEqual([0, 5, 9], remaining_cuts)
@@ -268,8 +288,11 @@ def test_collapse_sites(self):
 
         # test with alternative missing state indicator
         character_array = [0, 0, 0, 0, 0, 0, 0, 0, 0]
-        
-        updated_character_array, remaining_cuts = self.lineage_tracing_data_simulator_missing_data.collapse_sites(
+
+        (
+            updated_character_array,
+            remaining_cuts,
+        ) = self.lineage_tracing_data_simulator_missing_data.collapse_sites(
             character_array, [0, 1]
         )
         self.assertCountEqual([], remaining_cuts)
@@ -379,10 +402,15 @@ def test_no_resection(self):
 
     def test_simulator_with_state_generating_distribution(self):
 
-        self.lineage_tracing_data_simulator_state_distribution.overlay_data(self.basic_tree)
+        self.lineage_tracing_data_simulator_state_distribution.overlay_data(
+            self.basic_tree
+        )
 
         self.assertEqual(
-            10, len(self.lineage_tracing_data_simulator_state_distribution.mutation_priors)
+            10,
+            len(
+                self.lineage_tracing_data_simulator_state_distribution.mutation_priors
+            ),
         )
 
         character_matrix = self.basic_tree.get_original_character_matrix()
diff --git a/test/simulator_tests/supercellular_sampler_test.py b/test/simulator_tests/supercellular_sampler_test.py
new file mode 100755
index 00000000..9b90561b
--- /dev/null
+++ b/test/simulator_tests/supercellular_sampler_test.py
@@ -0,0 +1,192 @@
+import unittest
+
+import networkx as nx
+import numpy as np
+import pandas as pd
+
+from cassiopeia.data.CassiopeiaTree import CassiopeiaTree, CassiopeiaTreeError
+from cassiopeia.simulator.LeafSubsampler import LeafSubsamplerError
+from cassiopeia.simulator.SupercellularSampler import SupercellularSampler
+
+import cassiopeia.data.utilities as utilities
+
+
+class SupercellularSamplerTest(unittest.TestCase):
+    def setUp(self):
+        self.test_network = nx.DiGraph()
+        self.test_network.add_edges_from(
+            [
+                ("node0", "node1"),
+                ("node0", "node2"),
+                ("node1", "node3"),
+                ("node1", "node4"),
+                ("node2", "node5"),
+                ("node2", "node6"),
+                ("node4", "node7"),
+                ("node4", "node8"),
+                ("node8", "node9"),
+                ("node8", "node10"),
+                ("node10", "node11"),
+                ("node10", "node12"),
+                ("node12", "node13"),
+                ("node12", "node14"),
+                ("node14", "node15"),
+                ("node14", "node16"),
+                ("node16", "node17"),
+                ("node16", "node18"),
+            ]
+        )
+
+        # this should obey PP for easy checking of ancestral states
+        self.character_matrix = pd.DataFrame.from_dict(
+            {
+                "node3": [1, 0, 0, 0, 0, 0, 0, 0],
+                "node7": [1, 1, 0, 0, 0, 0, 0, 0],
+                "node9": [1, 1, 1, 0, 0, 0, 0, 0],
+                "node11": [1, 1, 1, 1, 0, 0, 0, 0],
+                "node13": [1, 1, 1, 1, 1, 0, 0, 0],
+                "node15": [1, 1, 1, 1, 1, 1, 0, 0],
+                "node17": [1, 1, 1, 1, 1, 1, 1, 0],
+                "node18": [1, 1, 1, 1, 1, 1, 1, 1],
+                "node5": [2, 0, 0, 0, 0, 0, 0, 0],
+                "node6": [2, 2, 0, 0, 0, 0, 0, 0],
+            },
+            orient="index",
+        )
+
+    def test_bad_parameters(self):
+        with self.assertRaises(LeafSubsamplerError):
+            uniform_sampler = SupercellularSampler(
+                ratio=0.5, number_of_merges=400
+            )
+        with self.assertRaises(LeafSubsamplerError):
+            uniform_sampler = SupercellularSampler()
+
+    def test_bad_number_of_samples(self):
+        tree = CassiopeiaTree(
+            tree=self.test_network, character_matrix=self.character_matrix
+        )
+        tree_no_character_matrix = CassiopeiaTree(tree=self.test_network)
+        with self.assertRaises(LeafSubsamplerError):
+            sampler = SupercellularSampler(number_of_merges=10)
+            sampler.subsample_leaves(tree)
+        with self.assertRaises(LeafSubsamplerError):
+            sampler = SupercellularSampler(number_of_merges=0)
+            sampler.subsample_leaves(tree)
+        with self.assertRaises(CassiopeiaTreeError):
+            sampler = SupercellularSampler(number_of_merges=2)
+            sampler.subsample_leaves(tree_no_character_matrix)
+
+    def test_subsample_balanced_tree(self):
+        tree = CassiopeiaTree(
+            tree=self.test_network, character_matrix=self.character_matrix
+        )
+
+        np.random.seed(10)
+        sampler = SupercellularSampler(number_of_merges=2)
+        res = sampler.subsample_leaves(tree=tree)
+        cm = res.get_current_character_matrix()
+        expected_character_matrix = pd.DataFrame.from_dict(
+            {
+                "node7": [1, 1, 0, 0, 0, 0, 0, 0],
+                "node9": [1, 1, 1, 0, 0, 0, 0, 0],
+                "node11": [1, 1, 1, 1, 0, 0, 0, 0],
+                "node13": [1, 1, 1, 1, 1, 0, 0, 0],
+                "node17": [1, 1, 1, 1, 1, 1, 1, 0],
+                "node6": [2, 2, 0, 0, 0, 0, 0, 0],
+                "node18-node15": [
+                    (1,),
+                    (1,),
+                    (1,),
+                    (1,),
+                    (1,),
+                    (1,),
+                    (0, 1),
+                    (0, 1),
+                ],
+                "node3-node5": [
+                    (1, 2),
+                    (0,),
+                    (0,),
+                    (0,),
+                    (0,),
+                    (0,),
+                    (0,),
+                    (0,),
+                ],
+            },
+            orient="index",
+        )
+        pd.testing.assert_frame_equal(expected_character_matrix, cm)
+        expected_edges = [
+            ("node0", "node3-node5"),
+            ("node0", "node4"),
+            ("node0", "node6"),
+            ("node4", "node7"),
+            ("node4", "node8"),
+            ("node8", "node10"),
+            ("node8", "node9"),
+            ("node10", "node11"),
+            ("node10", "node12"),
+            ("node12", "node13"),
+            ("node14", "node17"),
+            ("node12", "node14"),
+            ("node14", "node18-node15"),
+        ]
+        self.assertEqual(set(res.edges), set(expected_edges))
+
+        np.random.seed(10)
+        res = sampler.subsample_leaves(tree=tree, collapse_duplicates=False)
+        cm = res.get_current_character_matrix()
+        expected_character_matrix = pd.DataFrame.from_dict(
+            {
+                "node7": [1, 1, 0, 0, 0, 0, 0, 0],
+                "node9": [1, 1, 1, 0, 0, 0, 0, 0],
+                "node11": [1, 1, 1, 1, 0, 0, 0, 0],
+                "node13": [1, 1, 1, 1, 1, 0, 0, 0],
+                "node17": [1, 1, 1, 1, 1, 1, 1, 0],
+                "node6": [2, 2, 0, 0, 0, 0, 0, 0],
+                "node18-node15": [
+                    (1, 1),
+                    (1, 1),
+                    (1, 1),
+                    (1, 1),
+                    (1, 1),
+                    (1, 1),
+                    (1, 0),
+                    (1, 0),
+                ],
+                "node3-node5": [
+                    (1, 2),
+                    (0, 0),
+                    (0, 0),
+                    (0, 0),
+                    (0, 0),
+                    (0, 0),
+                    (0, 0),
+                    (0, 0),
+                ],
+            },
+            orient="index",
+        )
+        pd.testing.assert_frame_equal(expected_character_matrix, cm)
+        expected_edges = [
+            ("node4", "node7"),
+            ("node4", "node8"),
+            ("node8", "node9"),
+            ("node12", "node13"),
+            ("node14", "node17"),
+            ("node12", "node14"),
+            ("node14", "node18-node15"),
+            ("node0", "node6"),
+            ("node10", "node12"),
+            ("node10", "node11"),
+            ("node8", "node10"),
+            ("node0", "node4"),
+            ("node0", "node3-node5"),
+        ]
+        self.assertEqual(set(res.edges), set(expected_edges))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/simulator_tests/unifom_leaf_subsampler_test.py b/test/simulator_tests/unifom_leaf_subsampler_test.py
index 3246ed77..ac095acf 100644
--- a/test/simulator_tests/unifom_leaf_subsampler_test.py
+++ b/test/simulator_tests/unifom_leaf_subsampler_test.py
@@ -9,10 +9,13 @@
 
 import cassiopeia.data.utilities as utilities
 
+
 class UniformLeafSubsamplerTest(unittest.TestCase):
     def test_bad_parameters(self):
         with self.assertRaises(LeafSubsamplerError):
-            uniform_sampler = UniformLeafSubsampler(ratio = 0.5, number_of_leaves = 400)
+            uniform_sampler = UniformLeafSubsampler(
+                ratio=0.5, number_of_leaves=400
+            )
         with self.assertRaises(LeafSubsamplerError):
             uniform_sampler = UniformLeafSubsampler()
 
@@ -20,17 +23,20 @@ def test_bad_number_of_samples(self):
         balanced_tree = nx.balanced_tree(2, 3, create_using=nx.DiGraph)
         tree = CassiopeiaTree(tree=balanced_tree)
         with self.assertRaises(LeafSubsamplerError):
-            uniform_sampler = UniformLeafSubsampler(number_of_leaves = 0)
+            uniform_sampler = UniformLeafSubsampler(number_of_leaves=0)
             uniform_sampler.subsample_leaves(tree)
         with self.assertRaises(LeafSubsamplerError):
-            uniform_sampler = UniformLeafSubsampler(ratio = 0.0001)
+            uniform_sampler = UniformLeafSubsampler(ratio=0.0001)
             uniform_sampler.subsample_leaves(tree)
 
     def test_subsample_balanced_tree(self):
         balanced_tree = nx.balanced_tree(2, 3, create_using=nx.DiGraph)
-        balanced_tree = nx.relabel_nodes(balanced_tree, dict([(i, "node" + str(i)) for i in balanced_tree.nodes]))
-        balanced_tree.add_node('node15')
-        balanced_tree.add_edge('node15', 'node0')
+        balanced_tree = nx.relabel_nodes(
+            balanced_tree,
+            dict([(i, "node" + str(i)) for i in balanced_tree.nodes]),
+        )
+        balanced_tree.add_node("node15")
+        balanced_tree.add_edge("node15", "node0")
         tree = CassiopeiaTree(tree=balanced_tree)
 
         np.random.seed(10)
@@ -62,31 +68,33 @@ def test_subsample_balanced_tree(self):
     def test_subsample_custom_tree(self):
         custom_tree = nx.DiGraph()
         custom_tree.add_nodes_from(["node" + str(i) for i in range(17)])
-        custom_tree.add_edges_from([
-            ("node16", "node0"),
-            ("node0", "node1"),
-            ("node0", "node2"),
-            ("node1", "node3"),
-            ("node1", "node4"),
-            ("node2", "node5"),
-            ("node2", "node6"),
-            ("node4", "node7"),
-            ("node4", "node8"),
-            ("node6", "node9"),
-            ("node6", "node10"),
-            ("node7", "node11"),
-            ("node11", "node12"),
-            ("node11", "node13"),
-            ("node9", "node14"),
-            ("node9", "node15"),
-        ])
+        custom_tree.add_edges_from(
+            [
+                ("node16", "node0"),
+                ("node0", "node1"),
+                ("node0", "node2"),
+                ("node1", "node3"),
+                ("node1", "node4"),
+                ("node2", "node5"),
+                ("node2", "node6"),
+                ("node4", "node7"),
+                ("node4", "node8"),
+                ("node6", "node9"),
+                ("node6", "node10"),
+                ("node7", "node11"),
+                ("node11", "node12"),
+                ("node11", "node13"),
+                ("node9", "node14"),
+                ("node9", "node15"),
+            ]
+        )
         tree = CassiopeiaTree(tree=custom_tree)
         for u, v in tree.edges:
             tree.set_branch_length(u, v, 1.5)
 
         np.random.seed(10)
-        uni = UniformLeafSubsampler(ratio = 0.5)
-        res = uni.subsample_leaves(tree=tree, collapse_source = "node0")
+        uni = UniformLeafSubsampler(ratio=0.5)
+        res = uni.subsample_leaves(tree=tree, collapse_source="node0")
 
         expected_edges = {
             ("node16", "node0"): 1.5,
@@ -99,7 +107,9 @@ def test_subsample_custom_tree(self):
         }
         self.assertEqual(set(res.edges), set(expected_edges.keys()))
         for u, v in res.edges:
-            self.assertEqual(res.get_branch_length(u,v), expected_edges[(u,v)])
+            self.assertEqual(
+                res.get_branch_length(u, v), expected_edges[(u, v)]
+            )
 
         expected_times = {
             "node16": 0.0,
@@ -109,16 +119,14 @@ def test_subsample_custom_tree(self):
             "node3": 4.5,
             "node11": 7.5,
             "node12": 9.0,
-            "node13": 9.0
+            "node13": 9.0,
         }
         for u in res.nodes:
-            self.assertEqual(
-                res.get_time(u), expected_times[u]
-            )
+            self.assertEqual(res.get_time(u), expected_times[u])
 
         np.random.seed(11)
-        uni = UniformLeafSubsampler(number_of_leaves = 6)
-        res = uni.subsample_leaves(tree=tree, collapse_source = "node0")
+        uni = UniformLeafSubsampler(number_of_leaves=6)
+        res = uni.subsample_leaves(tree=tree, collapse_source="node0")
 
         expected_edges = [
             ("node16", "node0"),
@@ -132,11 +140,9 @@ def test_subsample_custom_tree(self):
             ("node2", "node6"),
             ("node6", "node10"),
             ("node6", "node15"),
-
         ]
         self.assertEqual(set(res.edges), set(expected_edges))
 
 
 if __name__ == "__main__":
     unittest.main()
-
diff --git a/test/solver_tests/dissimilarity_functions_test.py b/test/solver_tests/dissimilarity_functions_test.py
old mode 100644
new mode 100755
index 9333c64e..e3209c57
--- a/test/solver_tests/dissimilarity_functions_test.py
+++ b/test/solver_tests/dissimilarity_functions_test.py
@@ -3,6 +3,7 @@
 module.
 """
 import unittest
+from unittest import mock
 
 import numpy as np
 
@@ -16,6 +17,7 @@ def setUp(self):
         self.s1 = np.array([0, 1, 0, -1, 1, 2])
         self.s2 = np.array([1, 1, 0, 0, 1, 3])
         self.all_missing = np.array([-1, -1, -1, -1, -1, -1])
+        self.ambiguous = [(0,), (-1, 0), (0,), (-1, 0), (1,), (1,)]
 
         self.priors = {
             0: {1: 0.5, 2: 0.5},
@@ -184,23 +186,29 @@ def test_hamming_similarity_without_missing_all_missing(self):
 
     def test_hamming_similarity_normalized_identical(self):
 
-        similarity = dissimilarity_functions.hamming_similarity_normalized_over_missing(
-            self.s1, self.s1, -1
+        similarity = (
+            dissimilarity_functions.hamming_similarity_normalized_over_missing(
+                self.s1, self.s1, -1
+            )
         )
 
         self.assertEqual(similarity, 3 / 5)
 
     def test_hamming_similarity_normalized_no_priors(self):
 
-        similarity = dissimilarity_functions.hamming_similarity_normalized_over_missing(
-            self.s1, self.s2, -1
+        similarity = (
+            dissimilarity_functions.hamming_similarity_normalized_over_missing(
+                self.s1, self.s2, -1
+            )
         )
 
         self.assertEqual(similarity, 2 / 5)
 
     def test_hamming_similarity_normalized_priors(self):
-        similarity = dissimilarity_functions.hamming_similarity_normalized_over_missing(
-            self.s1, self.s2, -1, weights=self.nlweights
+        similarity = (
+            dissimilarity_functions.hamming_similarity_normalized_over_missing(
+                self.s1, self.s2, -1, weights=self.nlweights
+            )
         )
 
         expected_similarity = np.sum(
@@ -211,8 +219,10 @@ def test_hamming_similarity_normalized_priors(self):
 
     def test_hamming_similarity_normalized_all_missing(self):
 
-        similarity = dissimilarity_functions.hamming_similarity_normalized_over_missing(
-            self.s1, self.all_missing, -1, weights=self.nlweights
+        similarity = (
+            dissimilarity_functions.hamming_similarity_normalized_over_missing(
+                self.s1, self.all_missing, -1, weights=self.nlweights
+            )
         )
 
         self.assertEqual(similarity, 0)
@@ -252,6 +262,22 @@ def test_weighted_hamming_similarity_all_missing(self):
 
         self.assertEqual(similarity, 0)
 
+    def test_cluster_dissimilarity(self):
+        dissimilarity_function = (
+            dissimilarity_functions.weighted_hamming_distance
+        )
+        linkage_function = np.mean
+
+        result = dissimilarity_functions.cluster_dissimilarity(
+            dissimilarity_function,
+            self.s1,
+            self.ambiguous,
+            -1,
+            self.nlweights,
+            linkage_function,
+        )
+        np.testing.assert_almost_equal(result, 1.2544, decimal=4)
+
     def test_hamming_distance(self):
 
         distance = dissimilarity_functions.hamming_distance(self.s1, self.s2)
diff --git a/test/solver_tests/greedy_variants_test.py b/test/solver_tests/greedy_variants_test.py
index 1d98936c..9485e1ec 100644
--- a/test/solver_tests/greedy_variants_test.py
+++ b/test/solver_tests/greedy_variants_test.py
@@ -15,9 +15,9 @@ class GreedyVariantsTest(unittest.TestCase):
     def test_spectral_sparse_case(self):
         cm = pd.DataFrame.from_dict(
             {
-                "c1": [5, 3, 0, 0, 0], 
-                "c2": [0, 3, 4, 2, 1], 
-                "c3": [5, 0, 0, 0, 1], 
+                "c1": [5, 3, 0, 0, 0],
+                "c2": [0, 3, 4, 2, 1],
+                "c3": [5, 0, 0, 0, 1],
                 "c4": [5, 0, 4, 2, 0],
             },
             orient="index",
diff --git a/test/solver_tests/hybrid_solver_test.py b/test/solver_tests/hybrid_solver_test.py
index 08bb66ab..ec106fb3 100644
--- a/test/solver_tests/hybrid_solver_test.py
+++ b/test/solver_tests/hybrid_solver_test.py
@@ -181,10 +181,12 @@ def test_top_down_split_manual(self):
 
         character_matrix = self.pp_tree.get_original_character_matrix()
         # test manually
-        mutation_frequencies = self.hybrid_pp_solver.top_solver.compute_mutation_frequencies(
-            ["a", "b", "c", "d", "e"],
-            character_matrix,
-            self.pp_tree.missing_state_indicator,
+        mutation_frequencies = (
+            self.hybrid_pp_solver.top_solver.compute_mutation_frequencies(
+                ["a", "b", "c", "d", "e"],
+                character_matrix,
+                self.pp_tree.missing_state_indicator,
+            )
         )
 
         expected_dictionary = {
diff --git a/test/solver_tests/neighborjoining_solver_test.py b/test/solver_tests/neighborjoining_solver_test.py
old mode 100644
new mode 100755
index 6d4b26ab..10f8da22
--- a/test/solver_tests/neighborjoining_solver_test.py
+++ b/test/solver_tests/neighborjoining_solver_test.py
@@ -4,6 +4,7 @@
 import os
 import unittest
 from typing import Dict, Optional
+from unittest import mock
 
 import itertools
 import networkx as nx
@@ -363,6 +364,46 @@ def test_duplicate_sample_neighbor_joining(self):
             observed_triplet = find_triplet_structure(triplet, T)
             self.assertEqual(expected_triplet, observed_triplet)
 
+    def test_setup_root_finder_missing_dissimilarity_map(self):
+        tree = cas.data.CassiopeiaTree(character_matrix=self.cm)
+        with mock.patch.object(
+            tree, "compute_dissimilarity_map"
+        ) as compute_dissimilarity_map:
+            self.nj_solver_delta.setup_root_finder(tree)
+            compute_dissimilarity_map.assert_called_once_with(
+                delta_fn, "negative_log"
+            )
+        self.assertEqual(tree.root_sample_name, "root")
+
+    def test_setup_root_finder_existing_dissimilarity_map(self):
+        tree = cas.data.CassiopeiaTree(
+            character_matrix=self.cm,
+            dissimilarity_map=self.basic_dissimilarity_map,
+        )
+        with mock.patch.object(
+            tree, "compute_dissimilarity_map"
+        ) as compute_dissimilarity_map:
+            self.nj_solver_delta.setup_root_finder(tree)
+            compute_dissimilarity_map.assert_not_called()
+        self.assertEqual(tree.root_sample_name, "root")
+        dissimilarity_map = tree.get_dissimilarity_map()
+        self.assertEqual(
+            {"a", "b", "c", "d", "e", "root"}, set(dissimilarity_map.index)
+        )
+        self.assertEqual(
+            {"a", "b", "c", "d", "e", "root"}, set(dissimilarity_map.columns)
+        )
+        for leaf in self.cm.index:
+            delta = delta_fn(
+                [0] * tree.n_character,
+                self.cm.loc[leaf].values,
+                tree.missing_state_indicator,
+                None,
+            )
+            self.assertEqual(dissimilarity_map.loc[leaf, "root"], delta)
+            self.assertEqual(dissimilarity_map.loc["root", leaf], delta)
+        self.assertEqual(dissimilarity_map.loc["root", "root"], 0)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/solver_tests/percolation_test.py b/test/solver_tests/percolation_test.py
index a350f4db..7c46599f 100644
--- a/test/solver_tests/percolation_test.py
+++ b/test/solver_tests/percolation_test.py
@@ -75,7 +75,9 @@ def test_NJ_negative_similarity(self):
         T = p_tree.get_tree_topology()
 
         expected_tree = nx.DiGraph()
-        expected_tree.add_nodes_from(["c1", "c2", "c3", "c4", "c5", "c6", "6", "7", "8", "9"])
+        expected_tree.add_nodes_from(
+            ["c1", "c2", "c3", "c4", "c5", "c6", "6", "7", "8", "9"]
+        )
         expected_tree.add_edges_from(
             [
                 ("6", "7"),
@@ -90,7 +92,9 @@ def test_NJ_negative_similarity(self):
             ]
         )
 
-        triplets = itertools.combinations(["c1", "c2", "c3", "c4", "c5", "c6"], 3)
+        triplets = itertools.combinations(
+            ["c1", "c2", "c3", "c4", "c5", "c6"], 3
+        )
         for triplet in triplets:
             expected_triplet = find_triplet_structure(triplet, expected_tree)
             observed_triplet = find_triplet_structure(triplet, T)
@@ -259,4 +263,3 @@ def test_priors_case(self):
 
 if __name__ == "__main__":
     unittest.main()
-    
\ No newline at end of file
diff --git a/test/solver_tests/sharedmutationjoiner_test.py b/test/solver_tests/sharedmutationjoiner_test.py
old mode 100644
new mode 100755
index b8bc1118..f590d987
--- a/test/solver_tests/sharedmutationjoiner_test.py
+++ b/test/solver_tests/sharedmutationjoiner_test.py
@@ -3,7 +3,10 @@
 """
 import os
 import unittest
+import sys
+from functools import partial
 from typing import Dict, Optional
+from unittest import mock
 
 import itertools
 import networkx as nx
@@ -16,6 +19,7 @@
 from cassiopeia.data import utilities as data_utilities
 from cassiopeia.solver.SharedMutationJoiningSolver import (
     SharedMutationJoiningSolver,
+    SharedMutationJoiningSolverWarning,
 )
 from cassiopeia.solver import dissimilarity_functions
 from cassiopeia.solver import solver_utilities
@@ -75,6 +79,12 @@ def setUp(self):
         self.smj_solver = SharedMutationJoiningSolver(
             similarity_function=dissimilarity_functions.hamming_similarity_without_missing
         )
+        self.smj_solver_no_numba = SharedMutationJoiningSolver(
+            similarity_function=partial(
+                dissimilarity_functions.cluster_dissimilarity,
+                dissimilarity_functions.hamming_similarity_without_missing,
+            )
+        )
 
         # ---------------- Lineage Tracing NJ ----------------
 
@@ -121,6 +131,44 @@ def setUp(self):
             similarity_function=dissimilarity_functions.hamming_similarity_without_missing
         )
 
+    def test_init(self):
+        # This should numbaize
+        solver = SharedMutationJoiningSolver(
+            similarity_function=dissimilarity_functions.hamming_similarity_without_missing
+        )
+        self.assertTrue(
+            isinstance(
+                solver.nb_similarity_function, numba.core.registry.CPUDispatcher
+            )
+        )
+        self.assertTrue(
+            isinstance(
+                solver._SharedMutationJoiningSolver__update_similarity_map,
+                numba.core.registry.CPUDispatcher,
+            )
+        )
+
+        # This shouldn't numbaize
+        with self.assertWarns(SharedMutationJoiningSolverWarning):
+            solver = SharedMutationJoiningSolver(
+                similarity_function=partial(
+                    dissimilarity_functions.cluster_dissimilarity,
+                    dissimilarity_functions.hamming_similarity_without_missing,
+                )
+            )
+            self.assertFalse(
+                isinstance(
+                    solver.nb_similarity_function,
+                    numba.core.registry.CPUDispatcher,
+                )
+            )
+            self.assertFalse(
+                isinstance(
+                    solver._SharedMutationJoiningSolver__update_similarity_map,
+                    numba.core.registry.CPUDispatcher,
+                )
+            )
+
     def test_find_cherry(self):
         cherry = self.smj_solver.find_cherry(self.basic_similarity_map.values)
         delta = self.basic_similarity_map
@@ -309,6 +357,84 @@ def test_basic_solver(self):
                     nx.shortest_path_length(expected_tree, sample1, sample2),
                 )
 
+    def test_solver_no_numba(self):
+
+        self.smj_solver_no_numba.solve(self.basic_tree)
+
+        # test that the dissimilarity map and character matrix were not altered
+        cm = pd.DataFrame.from_dict(
+            {
+                "a": [0, 1, 2],
+                "b": [1, 1, 2],
+                "c": [2, 2, 2],
+                "d": [1, 1, 1],
+                "e": [0, 0, 0],
+            },
+            orient="index",
+            columns=["x1", "x2", "x3"],
+        )
+        for i in self.basic_similarity_map.index:
+            for j in self.basic_similarity_map.columns:
+                self.assertEqual(
+                    self.basic_similarity_map.loc[i, j],
+                    self.basic_tree.get_dissimilarity_map().loc[i, j],
+                )
+        for i in self.basic_tree.get_current_character_matrix().index:
+            for j in self.basic_tree.get_current_character_matrix().columns:
+                self.assertEqual(
+                    cm.loc[i, j],
+                    self.basic_tree.get_current_character_matrix().loc[i, j],
+                )
+
+        # test leaves exist in tree
+        _leaves = self.basic_tree.leaves
+
+        self.assertEqual(len(_leaves), self.basic_similarity_map.shape[0])
+        for _leaf in _leaves:
+            self.assertIn(_leaf, self.basic_similarity_map.index.values)
+
+        # test for expected number of edges
+        edges = list(self.basic_tree.edges)
+        self.assertEqual(len(edges), 8)
+
+        # test relationships between samples
+        expected_tree = nx.DiGraph()
+        expected_tree.add_nodes_from(
+            ["a", "b", "c", "d", "e", "5", "6", "7", "8"]
+        )
+        expected_tree.add_edges_from(
+            [
+                ("5", "a"),
+                ("5", "b"),
+                ("6", "5"),
+                ("6", "c"),
+                ("7", "d"),
+                ("7", "e"),
+                ("8", "6"),
+                ("8", "7"),
+            ]
+        )
+
+        T = self.basic_tree.get_tree_topology()
+        triplets = itertools.combinations(["a", "b", "c", "d", "e"], 3)
+        for triplet in triplets:
+
+            expected_triplet = find_triplet_structure(triplet, expected_tree)
+            observed_triplet = find_triplet_structure(triplet, T)
+            self.assertEqual(expected_triplet, observed_triplet)
+
+        # compare tree distances
+        T = T.to_undirected()
+        expected_tree = expected_tree.to_undirected()
+        for i in range(len(_leaves)):
+            sample1 = _leaves[i]
+            for j in range(i + 1, len(_leaves)):
+                sample2 = _leaves[j]
+                self.assertEqual(
+                    nx.shortest_path_length(T, sample1, sample2),
+                    nx.shortest_path_length(expected_tree, sample1, sample2),
+                )
+
     def test_smj_solver_weights(self):
 
         self.smj_solver_modified_pp.solve(self.pp_tree_priors)
diff --git a/test/solver_tests/spectral_test.py b/test/solver_tests/spectral_test.py
index 4aed0ae8..5712987d 100644
--- a/test/solver_tests/spectral_test.py
+++ b/test/solver_tests/spectral_test.py
@@ -205,9 +205,9 @@ def test_simple_base_case_string(self):
     def test_simple_base_case2(self):
         cm = pd.DataFrame.from_dict(
             {
-                "c1": [5, 3, 0, 0, 0], 
-                "c2": [0, 3, 4, 2, 1], 
-                "c3": [5, 0, 0, 0, 1], 
+                "c1": [5, 3, 0, 0, 0],
+                "c2": [0, 3, 4, 2, 1],
+                "c3": [5, 0, 0, 0, 1],
                 "c4": [5, 0, 4, 2, 0],
             },
             orient="index",
@@ -225,16 +225,15 @@ def test_simple_base_case2(self):
     def test_simple_base_case2_priors(self):
         cm = pd.DataFrame.from_dict(
             {
-                "c1": [5, 3, 0, 0, 0], 
-                "c2": [0, 3, 4, 2, 1], 
-                "c3": [5, 0, 0, 0, 1], 
+                "c1": [5, 3, 0, 0, 0],
+                "c2": [0, 3, 4, 2, 1],
+                "c3": [5, 0, 0, 0, 1],
                 "c4": [5, 0, 4, 2, 0],
             },
             orient="index",
             columns=["a", "b", "c", "d", "e"],
         )
 
-
         priors = {
             0: {5: 0.8},
             1: {3: 0.5},
diff --git a/test/solver_tests/vanillagreedy_test.py b/test/solver_tests/vanillagreedy_test.py
index 40a08f07..f24c4cfb 100644
--- a/test/solver_tests/vanillagreedy_test.py
+++ b/test/solver_tests/vanillagreedy_test.py
@@ -134,12 +134,13 @@ def test_average_missing_data_priors(self):
     def test_all_duplicates_base_case(self):
         cm = pd.DataFrame.from_dict(
             {
-                "c1": [5, 0, 1, 2, 0], 
-                "c2": [5, 0, 1, 2, 0], 
+                "c1": [5, 0, 1, 2, 0],
+                "c2": [5, 0, 1, 2, 0],
                 "c3": [5, 0, 1, 2, 0],
-            },             
+            },
             orient="index",
-            columns=["a", "b", "c", "d", "e"])
+            columns=["a", "b", "c", "d", "e"],
+        )
 
         vg_tree = cas.data.CassiopeiaTree(cm, missing_state_indicator=-1)
         vgsolver = VanillaGreedySolver()