handled issues with kpm results

tasks/keypathwayminer_task.py

@@ -198,10 +198,16 @@ def kpm_task(task_hook: TaskHook):
             network = {'nodes': nodes, 'edges': edges}
 
     # Remapping everything from UniProt Accession numbers to internal IDs
-    result_nodes = Protein.objects.filter(uniprot_code__in=network["nodes"])
+    flat_map = lambda f, xs: (y for ys in xs for y in f(ys))
+    uniprote_nodes = []
+    uniprote_nodes.extend(network["nodes"])
+    uniprote_nodes.extend(set(flat_map(lambda l: [l['from'], l['to']], network['edges'])))
+
+    result_nodes = Protein.objects.filter(uniprot_code__in=uniprote_nodes)
     node_map = {}
     node_map_for_edges = {}
 
+
     for node in result_nodes:
         node_map_for_edges[node.uniprot_code] = node.id
         if id_space == 'symbol':
@@ -213,19 +219,21 @@ def kpm_task(task_hook: TaskHook):
         if id_space == 'ensembl':
             node_map[node.uniprot_code] = [ensg.name for ensg in EnsemblGene.objects.filter(protein_id=node.id)]
 
-    flat_map = lambda f, xs: (y for ys in xs for y in f(ys))
 
-    network["nodes"] = flat_map(lambda uniprot: node_map[uniprot], network["nodes"])
-    network["edges"] = list(map(
-        lambda uniprot_edge: {"from": "p" + str(node_map_for_edges[uniprot_edge["from"]]),
-                              "to": "p" + str(node_map_for_edges[uniprot_edge["to"]])},
-        network["edges"]))
+
+    network["nodes"] = list(flat_map(lambda uniprot: node_map[uniprot], network["nodes"]))
+    drugstone_edges = []
+    for uniprot_edge in network['edges']:
+        from_node = f'p{node_map_for_edges[uniprot_edge["from"]]}' if uniprot_edge['from'] in node_map_for_edges else uniprot_edge['from']
+        to_node = f'p{node_map_for_edges[uniprot_edge["to"]]}' if uniprot_edge['to'] in node_map_for_edges else uniprot_edge['to']
+        drugstone_edges.append({"from": from_node,"to": to_node})
+    network['edges']=drugstone_edges
 
     node_types = {node: "protein" for node in network["nodes"]}
-    is_seed = {node: node in set(map(lambda p: "p"+str(p),protein_backend_ids)) for node in network["nodes"]}
+    is_seed = {node: node in set(map(lambda p: "p" + str(p), protein_backend_ids)) for node in network["nodes"]}
     result_dict = {
         "network": network,
-        "target_nodes":[node for node in network["nodes"] if node not in task_hook.seeds],
+        "target_nodes": [node for node in network["nodes"] if node not in task_hook.seeds],
         "node_attributes": {"node_types": node_types, "is_seed": is_seed}
     }
     task_hook.set_results(results=result_dict)