Spaces:

HUBioDataLab
/

ProtHGT

Paused

ProtHGT / visualize_kg.py

Erva Ulusoy

replace pyvis with eccharts kg view

614205a 4 days ago

17.3 kB

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 <script> block."""
    return json.dumps(value).replace('</', '<\\/')

	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 <script> block."""
	return json.dumps(value).replace('</', '<\\/')