import os import json import gzip import uuid # ECharts page the graph is rendered into (placeholders __ELEMENTS__, __KINDS__, __QUERY__, __PREDICTED__) with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'kg_graph_template.html'), encoding='utf-8') as _f: _GRAPH_TEMPLATE = _f.read() NODE_TYPE_COLORS = { 'Disease': '#079dbb', 'HPO': '#58d0e8', 'Drug': '#815ac0', 'Compound': '#d2b7e5', 'Domain': '#6bbf59', 'GO_term_P': '#ff8800', 'GO_term_F': '#ffaa00', 'GO_term_C': '#ffc300', 'Pathway': '#720026', 'kegg_Pathway': '#720026', 'EC_number': '#ce4257', 'Protein': '#3aa6a4' } EDGE_KIND_COLORS = { 'confirmed': '#8338ec', # predicted and annotated in the KG 'novel': '#c1121f', # predicted, not annotated in the KG 'annotation': '#219ebc', # annotated in the KG, not among the shown predictions 'other': '#666666', } EDGE_KIND_LABELS = { 'confirmed': 'Confirmed prediction (in KG)', 'novel': 'Novel prediction', 'annotation': 'Known GO annotation', 'other': 'Other relationship', } EDGE_LABEL_TRANSLATION = { 'Orthology': 'is ortholog to', 'Pathway': 'takes part in', 'kegg_path_prot': 'takes part in', ('domain_function', 'GO_term_F'): 'enables', ('domain_function', 'GO_term_P'): 'is involved in', ('domain_function', 'GO_term_C'): 'localizes to', 'function_function': 'ontological relationship', 'protein_domain': 'has', 'PPI': 'interacts with', 'HPO': 'is associated with', 'kegg_dis_prot': 'is related to', 'Disease': 'is related to', 'Drug': 'targets', 'kegg_dis_path': 'modulates', 'protein_ec': 'catalyzes', 'hpodis': 'is associated with', 'kegg_dis_drug': 'treats', 'Chembl': 'targets', ('protein_function', 'GO_term_F'): 'enables', ('protein_function', 'GO_term_P'): 'is involved in', ('protein_function', 'GO_term_C'): 'localizes to', } NODE_LABEL_TRANSLATION = { 'HPO': 'Phenotype', 'GO_term_P': 'Biological Process', 'GO_term_F': 'Molecular Function', 'GO_term_C': 'Cellular Component', 'kegg_Pathway': 'Pathway', 'EC_number': 'EC Number', } GO_CATEGORY_MAPPING = { 'Biological Process': 'GO_term_P', 'Molecular Function': 'GO_term_F', 'Cellular Component': 'GO_term_C' } def get_node_url(node_type, node_id): """Get the URL for a node based on its type and ID""" if node_type.startswith('GO_term'): return f"https://www.ebi.ac.uk/QuickGO/term/{node_id}" elif node_type == 'Protein': return f"https://www.uniprot.org/uniprotkb/{node_id}/entry" elif node_type == 'Disease': if ':' in node_id: ontology = node_id.split(':')[0] if ontology == 'EFO': return f"http://www.ebi.ac.uk/efo/EFO_{node_id.split(':')[1]}" elif ontology == 'MONDO': return f'http://purl.obolibrary.org/obo/MONDO_{node_id.split(":")[1]}' elif ontology == 'Orphanet': return f"http://www.orpha.net/ORDO/Orphanet_{node_id.split(':')[1]}" else: return f"https://www.genome.jp/entry/{node_id}" elif node_type == 'HPO': return f"https://hpo.jax.org/browse/term/{node_id}" elif node_type == 'Drug': return f"https://go.drugbank.com/drugs/{node_id}" elif node_type == 'Compound': return f"https://www.ebi.ac.uk/chembl/explore/compound/{node_id}" elif node_type == 'Domain': return f"https://www.ebi.ac.uk/interpro/entry/InterPro/{node_id}" elif node_type == 'Pathway': return f"https://reactome.org/content/detail/{node_id}" elif node_type == 'kegg_Pathway': return f"https://www.genome.jp/pathway/{node_id}" elif node_type == 'EC_number': return f"https://enzyme.expasy.org/EC/{node_id}" else: return None def _gather_protein_edges(data, protein_id): protein_idx = data['Protein']['id_mapping'][protein_id] reverse_id_mapping = {} for node_type in data.node_types: reverse_id_mapping[node_type] = {v:k for k, v in data[node_type]['id_mapping'].items()} protein_edges = {} print(f'Gathering edges for {protein_id}...') for edge_type in data.edge_types: if 'rev' not in edge_type[1]: if edge_type not in protein_edges: protein_edges[edge_type] = [] if edge_type[0] == 'Protein': print(f'Gathering edges for {edge_type}...') # append the edges with protein_idx as source node edges = data[edge_type].edge_index[:, data[edge_type].edge_index[0] == protein_idx] protein_edges[edge_type].extend(edges.T.tolist()) elif edge_type[2] == 'Protein': print(f'Gathering edges for {edge_type}...') # append the edges with protein_idx as target node edges = data[edge_type].edge_index[:, data[edge_type].edge_index[1] == protein_idx] protein_edges[edge_type].extend(edges.T.tolist()) for edge_type in protein_edges.keys(): if protein_edges[edge_type]: mapped_edges = set() for edge in protein_edges[edge_type]: # Get source and target node types from edge_type source_type, _, target_type = edge_type # Map indices back to original IDs source_id = reverse_id_mapping[source_type][edge[0]] target_id = reverse_id_mapping[target_type][edge[1]] mapped_edges.add((source_id, target_id)) protein_edges[edge_type] = mapped_edges return protein_edges def _filter_edges(protein_id, protein_edges, prediction_df, limit=10, is_second_degree=False, second_degree_limit=3, prediction_limit=5): """ Select the edges to draw for one node. Returns {edge_type: [((source, target), probability, kind), ...]}, where kind is 'confirmed' - predicted GO term that is also annotated in the KG (purple); includes annotated terms predicted below the top `prediction_limit` terms 'novel' - one of the top `prediction_limit` predicted GO terms, not annotated in the KG (red) 'annotation' - protein-GO annotation in the KG that was not predicted (blue); probability is None 'other' - any other relationship, incl. domain-GO and GO-GO edges (gray) Args: protein_id: ID of the node whose edges are filtered protein_edges: {edge_type: set of (source, target)} from the KG prediction_df: DataFrame containing predictions limit: Maximum number of KG edges per edge type for first-degree connections is_second_degree: Whether these are second-degree edges (KG only, no predictions) second_degree_limit: Maximum number of edges per edge type for second-degree connections prediction_limit: Maximum number of predicted GO terms per predicted category """ filtered_edges = {} # Use appropriate limit based on edge degree current_limit = second_degree_limit if is_second_degree else limit # Predictions are drawn for the query protein only; second-degree edges show the KG as is. predicted = {} if not is_second_degree: rows = prediction_df[prediction_df['UniProt_ID'] == protein_id] for category, group in rows.groupby('GO_category'): predicted[GO_CATEGORY_MAPPING[category]] = group.sort_values( by=['Probability', 'GO_ID'], ascending=[False, True]) # Include predicted categories even when the protein has no KG annotation in them edge_types = {edge_type for edge_type, edges in protein_edges.items() if edges} edge_types |= {('Protein', 'protein_function', category) for category in predicted} for edge_type in sorted(edge_types): # Sorted so the same subset is shown on every run (set order varies between processes) edges = sorted(protein_edges.get(edge_type) or []) if edge_type[1] == 'protein_function' and edge_type[2] in predicted: kg_edges = set(edges) scores = dict(zip(predicted[edge_type[2]]['GO_ID'], predicted[edge_type[2]]['Probability'])) top_terms = list(scores)[:prediction_limit] shown = [((protein_id, term), scores[term], 'confirmed' if (protein_id, term) in kg_edges else 'novel') for term in top_terms] # KG annotations outside the top predictions, under the regular neighbour limit: purple if the model # also predicted them (any rank above the threshold, highest scores first), blue if it did not rest = [edge for edge in edges if edge[1] not in top_terms] rest.sort(key=lambda edge: -scores.get(edge[1], -1)) shown += [(edge, scores[edge[1]], 'confirmed') if edge[1] in scores else (edge, None, 'annotation') for edge in rest[:current_limit]] filtered_edges[edge_type] = shown elif edge_type[1] == 'protein_function': filtered_edges[edge_type] = [(edge, None, 'annotation') for edge in edges[:current_limit]] else: filtered_edges[edge_type] = [(edge, None, 'other') for edge in edges[:current_limit]] return filtered_edges def _gather_neighbor_edges(data, node_id, node_type, exclude_node_id): """Gather edges for a neighbor node, excluding edges back to the original query protein""" node_idx = data[node_type]['id_mapping'][node_id] reverse_id_mapping = {} for ntype in data.node_types: reverse_id_mapping[ntype] = {v:k for k, v in data[ntype]['id_mapping'].items()} neighbor_edges = {} for edge_type in data.edge_types: if 'rev' not in edge_type[1]: if edge_type not in neighbor_edges: neighbor_edges[edge_type] = [] if edge_type[0] == node_type: # Get edges where neighbor is source edges = data[edge_type].edge_index[:, data[edge_type].edge_index[0] == node_idx] edges = edges.T.tolist() # Filter out edges going back to the query protein edges = [edge for edge in edges if reverse_id_mapping[edge_type[2]][edge[1]] != exclude_node_id] neighbor_edges[edge_type].extend(edges) elif edge_type[2] == node_type: # Get edges where neighbor is target edges = data[edge_type].edge_index[:, data[edge_type].edge_index[1] == node_idx] edges = edges.T.tolist() # Filter out edges coming from the query protein edges = [edge for edge in edges if reverse_id_mapping[edge_type[0]][edge[0]] != exclude_node_id] neighbor_edges[edge_type].extend(edges) # Map indices back to IDs for edge_type in neighbor_edges.keys(): if neighbor_edges[edge_type]: mapped_edges = set() for edge in neighbor_edges[edge_type]: source_type, _, target_type = edge_type source_id = reverse_id_mapping[source_type][edge[0]] target_id = reverse_id_mapping[target_type][edge[1]] mapped_edges.add((source_id, target_id)) neighbor_edges[edge_type] = mapped_edges return neighbor_edges def visualize_protein_subgraph(data, protein_id, prediction_df, limit=10, second_degree_limit=3, include_second_degree=False, prediction_limit=5, predicted_categories=None): """ Build the knowledge-graph view around one protein and save it as an HTML file. predicted_categories: GO categories the run predicted (e.g. ['GO_term_F']); used to tell apart, in the tooltips, annotations that were not predicted from those whose category was not run at all. Defaults to the categories present in prediction_df. Returns (file_path, visualized_edges). """ from run_prothgt_app import load_name_info # cached loader; avoids re-reading 24 MB per call name_info = load_name_info() # Get the first-degree edges and filter them protein_edges = _gather_protein_edges(data, protein_id) first_degree_edges = _filter_edges(protein_id, protein_edges, prediction_df, limit=limit, is_second_degree=False, prediction_limit=prediction_limit) # Initialize all_edges with first degree edges all_edges = first_degree_edges.copy() if include_second_degree: # Collect neighbor nodes from first-degree edges neighbor_nodes = set() for edge_type, edges in first_degree_edges.items(): source_type, _, target_type = edge_type for edge_info in edges: edge = edge_info[0] source, target = edge if source != protein_id: neighbor_nodes.add((source, source_type)) if target != protein_id: neighbor_nodes.add((target, target_type)) # Gather and filter second-degree edges with the smaller limit second_degree_edges = {} for neighbor_id, neighbor_type in sorted(neighbor_nodes): neighbor_edges = _gather_neighbor_edges(data, neighbor_id, neighbor_type, protein_id) filtered_neighbor_edges = _filter_edges(neighbor_id, neighbor_edges, prediction_df, limit=limit, is_second_degree=True, second_degree_limit=second_degree_limit) # Merge filtered neighbor edges into second_degree_edges for edge_type, edges in filtered_neighbor_edges.items(): if edge_type not in second_degree_edges: second_degree_edges[edge_type] = [] second_degree_edges[edge_type].extend(edges) # Merge first and second degree edges for edge_type, edges in second_degree_edges.items(): if edge_type in all_edges: all_edges[edge_type].extend(edges) else: all_edges[edge_type] = edges # Update visualized_edges with all edges visualized_edges = all_edges print(f'Edges to be visualized: {visualized_edges}') # Render with the ECharts template (radial layout, see kg_graph_template.html) elements = _build_graph_elements(protein_id, visualized_edges, name_info) html = _GRAPH_TEMPLATE \ .replace('__ELEMENTS__', _json_for_script(elements)) \ .replace('__KINDS__', _json_for_script(EDGE_KIND_LABELS)) \ .replace('__QUERY__', _json_for_script(protein_id)) \ .replace('__PREDICTED__', _json_for_script( sorted(NODE_LABEL_TRANSLATION[c] for c in predicted_categories) if predicted_categories is not None else sorted(prediction_df['GO_category'].unique().tolist()))) # Save graph to a protein-specific file in a temporary directory os.makedirs('temp_viz', exist_ok=True) suffix = "_with_2nd_degree" if include_second_degree else "_1st_degree" # Unique per call: the app caches one file per setting combination and the Space is shared file_path = os.path.join('temp_viz', f'{protein_id}_graph{suffix}_{uuid.uuid4().hex[:8]}.html') with open(file_path, 'w', encoding='utf-8') as f: f.write(html) return file_path, visualized_edges def _build_graph_elements(protein_id, visualized_edges, name_info): """Nodes and edges for the graph template, as [{'data': {...}}, ...] (nodes first).""" nodes, edges = {}, [] def add_node(node_id, node_type): if node_id in nodes: return names = name_info['GO_term' if node_type.startswith('GO_term') else node_type] nodes[node_id] = {'data': { 'id': node_id, 'name': names.get(node_id, node_id), 'type': NODE_LABEL_TRANSLATION.get(node_type, node_type), 'color': NODE_TYPE_COLORS[node_type], 'url': get_node_url(node_type, node_id) or '', 'role': 'query' if node_id == protein_id else '', }} add_node(protein_id, 'Protein') for (source_type, relation, target_type), items in visualized_edges.items(): key = (relation, target_type) if relation in ('protein_function', 'domain_function') else relation for (source, target), probability, kind in items: source, target = str(source), str(target) add_node(source, source_type) add_node(target, target_type) edges.append({'data': { 'source': source, 'target': target, 'type': f'{source_type}_{relation}_{target_type}', # key used in the edges JSON download 'relation': EDGE_LABEL_TRANSLATION[key], 'kind': kind, 'color': EDGE_KIND_COLORS[kind], 'prob': None if probability is None else float(probability), }}) return list(nodes.values()) + edges def _json_for_script(value): """JSON that is safe to inline in a