Correct assignment of distances using model in partial query mode

bacpop · Nov 29, 2024 · f2c3179 · f2c3179
1 parent 3a1332b
commit f2c3179
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/PopPUNK/network.py b/PopPUNK/network.py
@@ -2026,6 +2026,7 @@ def generate_network_from_distances(mode = None,
                                     combined_seq = None,
                                     rlist = None,
                                     old_rlist = None,
+                                    model = None,
                                     distance_type = 'core',
                                     threads = 1,
                                     gpu_graph = False):
@@ -2050,6 +2051,7 @@ def generate_network_from_distances(mode = None,
             List of reference sequence labels
         old_rlist (list)
             List of reference sequence labels for previous MST
+        model ()
         distance_type (str)
             Whether to use core or accessory distances for MST calculation
             or dense network weighting
@@ -2075,12 +2077,12 @@ def generate_network_from_distances(mode = None,
         complete_distMat = \
             np.hstack((pp_sketchlib.squareToLong(core_distMat, threads).reshape(-1, 1),
                     pp_sketchlib.squareToLong(acc_distMat, threads).reshape(-1, 1)))
-        # Dense network may be slow
-        sys.stderr.write("Generating MST from dense distances (may be slow)\n")
+        # Identify short distances and use these to extend the model
+        indivAssignments = model.assign(complete_distMat)
         G = construct_network_from_assignments(combined_seq,
                                                 combined_seq,
-                                                [0]*complete_distMat.shape[0],
-                                                within_label = 0,
+                                                indivAssignments,
+                                                model.within_label,
                                                 distMat = complete_distMat,
                                                 weights_type = distance_type,
                                                 use_gpu = gpu_graph,

diff --git a/PopPUNK/visualise.py b/PopPUNK/visualise.py
@@ -219,7 +219,6 @@ def generate_visualisations(query_db,
 
     from .models import loadClusterFit
 
-    from .network import construct_network_from_assignments
     from .network import generate_minimum_spanning_tree
     from .network import load_network_file
     from .network import cugraph_to_graph_tool
@@ -706,6 +705,7 @@ def generate_visualisations(query_db,
                                                                 core_distMat = core_distMat,
                                                                 acc_distMat = acc_distMat,
                                                                 combined_seq = combined_seq,
+                                                                model = model,
                                                                 distance_type = mst_distances,
                                                                 threads = threads,
                                                                 gpu_graph = gpu_graph)