added STDR wrap-around

cassiopeia/solver/STDRSolver.py

@@ -0,0 +1,35 @@
+from typing import Callable
+import dendropy
+import numpy as np
+import ete3
+import spectraltree
+from typing import Callable
+
+from cassiopeia.data import CassiopeiaTree
+from cassiopeia.data import utilities as data_utilities
+from cassiopeia.solver import CassiopeiaSolver
+
+class STDRSolver(CassiopeiaSolver.CassiopeiaSolver):
+
+    def __init__(
+        self,
+        similarity_function: Callable[[np.array], np.array]
+    ):
+
+        self.similarity_function = similarity_function
+
+    def solve(self, cassiopeia_tree: CassiopeiaTree) -> None:
+        character_matrix = cassiopeia_tree.get_current_character_matrix()
+        taxon_namespace = dendropy.TaxonNamespace(list(character_matrix.index))
+        metadata = spectraltree.utils.TaxaMetadata(taxon_namespace, list(character_matrix.index), alphabet=None)
+
+        stdr_nj = spectraltree.STDR(spectraltree.NeighborJoining, self.similarity_function)   
+
+        tree_stdr_nj = stdr_nj(character_matrix.values, 
+                taxa_metadata= metadata, 
+                threshold = 64,
+                min_split = 1,
+                merge_method = "least_square", 
+                verbose=False)
+        tree = data_utilities.ete3_to_networkx(ete3.Tree(tree_stdr_nj.as_string(schema="newick")))
+        cassiopeia_tree.populate_tree(tree)

cassiopeia/solver/__init__.py

@@ -10,6 +10,7 @@
 from .MaxCutSolver import MaxCutSolver
 from .NeighborJoiningSolver import NeighborJoiningSolver
 from .PercolationSolver import PercolationSolver
+from .STDRSolver import STDRSolver
 from .SharedMutationJoiningSolver import SharedMutationJoiningSolver
 from .SpectralGreedySolver import SpectralGreedySolver
 from .SpectralSolver import SpectralSolver

cassiopeia/solver/stdr_similarities.py

@@ -0,0 +1,54 @@
+"""This file stores a variety of ways to calculate the similarity used for STDR with neighborjoining
+"""
+
+import numpy as np
+import scipy
+import spectraltree
+
+def JC_similarity_matrix(vals):
+    return spectraltree.JC_similarity_matrix(vals)
+
+def hamming_similarity_matrix(vals):
+    """Hamming Similarity"""
+    classes = np.unique(vals)
+    if classes.dtype == np.dtype('<U1'):
+        # needed to use hamming distance with string arrays
+        vord = np.vectorize(ord)
+        vals = vord(vals)
+    k = len(classes)
+    hamming_matrix = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(vals, metric='hamming'))
+    return 1 - hamming_matrix
+
+def hamming_sim_ignore_missing_values(vals):
+    missing_val = -1
+    classnames, indices = np.unique(vals, return_inverse=True)
+    num_arr = indices.reshape(vals.shape)
+    hamming_matrix = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(num_arr, metric='hamming'))
+    missing_array = (vals==missing_val)
+    pdist_xor = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(missing_array, lambda u,v: np.sum(np.logical_xor(u,v))))
+
+    return 1 - (hamming_matrix*vals.shape[1] - pdist_xor) / (np.ones_like(hamming_matrix) * vals.shape[1])
+
+def hamming_sim_normalize_missing_values(vals):
+    missing_val = -1
+    classnames, indices = np.unique(vals, return_inverse=True)
+    num_arr = indices.reshape(vals.shape)
+    hamming_matrix = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(num_arr, metric='hamming'))
+    missing_array = (vals==missing_val)
+    pdist_xor = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(missing_array, lambda u,v: np.sum(np.logical_xor(u,v))))
+    pdist_or = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(missing_array, lambda u,v: np.sum(np.logical_or(u,v))))
+
+    return 1 - (hamming_matrix*vals.shape[1] - pdist_xor) / (np.ones_like(hamming_matrix) * vals.shape[1] - pdist_or)
+
+def shared_mutations_normalize_missing_values(vals):
+    missing_val = -1
+    uncut_state = 0
+    classnames, indices = np.unique(vals, return_inverse=True)
+    num_arr = indices.reshape(vals.shape)
+    hamming_matrix = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(num_arr, metric='hamming'))
+    missing_array = (vals==missing_val)
+    uncut_array = (vals==uncut_state)
+    pdist_and = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(missing_array, lambda u,v: np.sum(np.logical_and(u,v))))
+    pdist_or = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(missing_array, lambda u,v: np.sum(np.logical_or(u,v))))
+    uncut_dist_and = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(uncut_array, lambda u,v: np.sum(np.logical_and(u,v))))
+    return 1 - (((1 - hamming_matrix)*vals.shape[1] - pdist_and - uncut_dist_and) / (np.ones_like(hamming_matrix) * vals.shape[1] - pdist_or))