Add calculation of minimum spanning trees

PopPUNK/__main__.py

@@ -16,6 +16,9 @@
 # import poppunk package
 from .__init__ import __version__
 
+# globals
+accepted_weights_types = ["core", "accessory", "euclidean"]
+
 #******************************#
 #*                            *#
 #* Command line parsing       *#
@@ -238,14 +241,15 @@ def main():
     # check if working with lineages
     if args.fit_model == 'lineage':
         rank_list = sorted([int(x) for x in args.ranks.split(',')])
-        if min(rank_list) == 0 or max(rank_list) > 100:
-            sys.stderr.write('Ranks should be small non-zero integers for sensible results\n')
-            sys.exit(1)
+        if int(min(rank_list)) == 0:
+            sys.stderr.write("Ranks must be >= 1\n")
+        if max(rank_list) > 100:
+            sys.stderr.write("WARNING: Ranks should be small non-zero integers for sensible lineage results\n")
 
     # run according to mode
     sys.stderr.write("PopPUNK (POPulation Partitioning Using Nucleotide Kmers)\n")
     sys.stderr.write("\t(with backend: " + dbFuncs['backend'] + " v" + dbFuncs['backend_version'] + "\n")
-    sys.stderr.write('\t sketchlib: ' + checkSketchlibLibrary() + ')\n')
+    sys.stderr.write("\t sketchlib: " + checkSketchlibLibrary() + ")\n")
 
     # Check on parallelisation of graph-tools
     setGtThreads(args.threads)

PopPUNK/models.py

@@ -90,12 +90,13 @@ def loadClusterFit(pkl_file, npz_file, outPrefix = "", max_samples = 100000):
         sys.stderr.write("Loading previously refined model\n")
         load_obj = RefineFit(outPrefix)
     elif fit_type == "lineage":
-        sys.stderr.write("Loading previously lineage cluster model\n")
+        sys.stderr.write("Loading lineage cluster model\n")
         load_obj = LineageFit(outPrefix, fit_object[0])
     else:
         raise RuntimeError("Undefined model type: " + str(fit_type))
 
     load_obj.load(fit_data, fit_object)
+    sys.stderr.write("Completed model loading\n")
     return load_obj
 
 class ClusterFit:

PopPUNK/network.py

@@ -20,6 +20,9 @@
 from functools import partial
 from multiprocessing import Pool
 import graph_tool.all as gt
+import dendropy
+
+from .__main__ import accepted_weights_types
 
 from .sketchlib import addRandom
 
@@ -46,7 +49,6 @@ def fetchNetwork(network_dir, model, refList, ref_graph = False,
                 Names of references that should be in the network
             ref_graph (bool)
                 Use ref only graph, if available
-
                 [default = False]
             core_only (bool)
                 Return the network created using only core distances
@@ -263,7 +265,8 @@ def writeReferences(refList, outPrefix):
     return refFileName
 
 def constructNetwork(rlist, qlist, assignments, within_label,
-                     summarise = True, edge_list = False, weights = None):
+                     summarise = True, edge_list = False, weights = None,
+                     weights_type = 'euclidean', sparse_input = None):
     """Construct an unweighted, undirected network without self-loops.
     Nodes are samples and edges where samples are within the same cluster
 
@@ -281,13 +284,17 @@ def constructNetwork(rlist, qlist, assignments, within_label,
             from :func:`~PopPUNK.bgmm.findWithinLabel`
         summarise (bool)
             Whether to calculate and print network summaries with :func:`~networkSummary`
-
             (default = True)
         edge_list (bool)
             Whether input is edges, tuples of (v1, v2). Used with lineage assignment
         weights (numpy.array)
             If passed, the core,accessory distances for each assignment, which will
             be annotated as an edge attribute
+        weights_type (str)
+            Specifies the type of weight to be annotated on the graph - options are core,
+            accessory or euclidean distance
+        sparse_input (numpy.array)
+            Sparse distance matrix from lineage fit
 
     Returns:
         G (graph)
@@ -303,21 +310,37 @@ def constructNetwork(rlist, qlist, assignments, within_label,
         self_comparison = False
         vertex_labels.append(qlist)
 
+    # Check weights type is valid
+    if weights_type not in accepted_weights_types:
+        sys.stderr.write("Unable to calculate distance type " + str(weights_type) + "; "
+                         "accepted types are " + str(accepted_weights_types) + "\n")
+        sys.exit(1)
+    if edge_list and sparse_input:
+        raise RuntimeError("Cannot construct network from edge list and sparse matrix")
+
     # identify edges
+    connections = []
     if edge_list:
         if weights is not None:
-            connections = []
             for weight, (ref, query) in zip(weights, assignments):
                 connections.append((ref, query, weight))
         else:
             connections = assignments
+    elif sparse_input is not None:
+        for ref, query, weight in zip(sparse_input.row, sparse_input.col, sparse_input.data):
+            connections.append((ref, query, weight))
     else:
         for row_idx, (assignment, (ref, query)) in enumerate(zip(assignments,
                                                                  listDistInts(rlist, qlist,
                                                                               self = self_comparison))):
             if assignment == within_label:
                 if weights is not None:
-                    dist = np.linalg.norm(weights[row_idx, :])
+                    if weights_type == 'euclidean':
+                        dist = np.linalg.norm(weights[row_idx, :])
+                    elif weights_type == 'core':
+                        dist = weights[row_idx, 0]
+                    elif weights_type == 'accessory':
+                        dist = weights[row_idx, 1]
                     edge_tuple = (ref, query, dist)
                 else:
                     edge_tuple = (ref, query)
@@ -327,7 +350,7 @@ def constructNetwork(rlist, qlist, assignments, within_label,
     G = gt.Graph(directed = False)
     G.add_vertex(len(vertex_labels))
 
-    if weights is not None:
+    if weights is not None or sparse_input is not None:
         eweight = G.new_ep("float")
         G.add_edge_list(connections, eprops = [eweight])
         G.edge_properties["weight"] = eweight
@@ -704,3 +727,77 @@ def printExternalClusters(newClusters, extClusterFile, outPrefix,
     pd.DataFrame(data=d).to_csv(outPrefix + "_external_clusters.csv",
                                 columns = ["sample"] + list(extClusters.keys()),
                                 index = False)
+
+def generate_minimum_spanning_tree(G, from_cugraph = False):
+    """Generate a minimum spanning tree from a network
+
+    Args:
+       G (network)
+           Graph tool network
+       from_cugraph (bool)
+            If a pre-calculated MST from cugraph
+            [default = False]
+
+    Returns:
+       mst_network (str)
+           Minimum spanning tree (as graph-tool graph)
+    """
+    #
+    # Create MST
+    #
+    if not from_cugraph:
+        sys.stderr.write("Starting calculation of minimum-spanning tree\n")
+
+        # Test if weighted network and calculate minimum spanning tree
+        if "weight" in G.edge_properties:
+            mst_edge_prop_map = gt.min_spanning_tree(G, weights = G.ep["weight"])
+            mst_network = gt.GraphView(G, efilt = mst_edge_prop_map)
+        else:
+            sys.stderr.write("generate_minimum_spanning_tree requires a weighted graph\n")
+            raise RuntimeError("MST passed unweighted graph")
+    else:
+        mst_network = G
+
+    # Find seed nodes as those with greatest outdegree in each component
+    seed_vertices = set()
+    component_assignments, component_frequencies = gt.label_components(mst_network)
+    for component_index in range(len(component_frequencies)):
+        component_members = component_assignments.a == component_index
+        component = gt.GraphView(mst_network, vfilt = component_members)
+        component_vertices = component.get_vertices()
+        out_degrees = component.get_out_degrees(component_vertices)
+        seed_vertex = list(component_vertices[np.where(out_degrees == np.amax(out_degrees))])
+        seed_vertices.add(seed_vertex[0]) # Can only add one otherwise not MST
+
+    # If multiple components, calculate distances between seed nodes
+    if len(component_frequencies) > 1:
+        # Extract distances
+        connections = []
+        max_weight = float(np.max(G.edge_properties["weight"].a))
+        for ref in seed_vertices:
+            seed_edges = G.get_all_edges(ref, [G.ep['weight']])
+            found = False  # Not all edges may be in graph
+            for seed_edge in seed_edges:
+                if seed_edge[1] in seed_vertices:
+                    found = True
+                    connections.append((seed_edge))
+            # TODO: alternative would be to requery the DB (likely quick)
+            if found == False:
+                for query in seed_vertices:
+                    if query != ref:
+                        connections.append((ref, query, max_weight))
+
+        # Construct graph
+        seed_G = gt.Graph(directed = False)
+        seed_G.add_vertex(len(seed_vertex))
+        eweight = seed_G.new_ep("float")
+        seed_G.add_edge_list(connections, eprops = [eweight])
+        seed_G.edge_properties["weight"] = eweight
+        seed_mst_edge_prop_map = gt.min_spanning_tree(seed_G, weights = seed_G.ep["weight"])
+        seed_mst_network = gt.GraphView(seed_G, efilt = seed_mst_edge_prop_map)
+        # Insert seed MST into original MST - may be possible to use graph_union with include=True & intersection
+        deep_edges = seed_mst_network.get_edges([seed_mst_network.ep["weight"]])
+        mst_network.add_edge_list(deep_edges)
+
+    sys.stderr.write("Completed calculation of minimum-spanning tree\n")
+    return mst_network