working app

ProtHGT_app.py

@@ -9,18 +9,301 @@ from run_prothgt_app import *
 def convert_df(df):
    return df.to_csv(index=False).encode('utf-8')
 
+# Initialize session state variables
+if 'predictions_df' not in st.session_state:
+    st.session_state.predictions_df = None
+if 'submitted' not in st.session_state:
+    st.session_state.submitted = False
+
+
 with st.sidebar:
-    st.title("ProtHGT: Heterogeneous Graph Transformers for Automated Protein Function Prediction Using Knowledge Graphs and Language Models")
-    st.write("[![publication](https://img.shields.io/badge/DOI-10.1002/pro.4988-b31b1b.svg)]() [![github-repository](https://img.shields.io/badge/GitHub-black?logo=github)](https://github.com/HUBioDataLab/ProtHGT)")
-
-    # Add protein selection
-    # You'll need to replace this with your actual data loading
-    available_proteins = get_available_proteins()  # Function to get list of proteins from your data
-    selected_protein = st.selectbox(
-        "Select or search for a protein (UniProt ID)",
-        options=available_proteins,
-        placeholder="Start typing to search...",
+    st.markdown("""
+        <style>
+        .title {
+            font-size: 35px;
+            font-weight: bold;
+            color: #424242;
+            margin-bottom: 0px;
+        }
+        .subtitle {
+            font-size: 20px;
+            color: #424242;
+            margin-bottom: 20px;
+            line-height: 1.5;
+        }
+        .badges {
+            margin-top: 10px;
+            margin-bottom: 20px;
+        }
+        </style>
+        
+        <div class="title">ProtHGT</div>
+        <div class="subtitle">Heterogeneous Graph Transformers for Automated Protein Function Prediction Using Knowledge Graphs and Language Models</div>
+        <div class="badges">
+            <a href="">
+                <img src="https://img.shields.io/badge/DOI-10.1002/pro.4988-b31b1b.svg" alt="publication">
+            </a>
+            <a href="https://github.com/HUBioDataLab/ProtHGT">
+                <img src="https://img.shields.io/badge/GitHub-black?logo=github" alt="github-repository">
+            </a>
+        </div>
+    """, unsafe_allow_html=True)
+
+    available_proteins = get_available_proteins() 
+
+    selected_proteins = []
+    
+    # Add protein selection methods
+    selection_method = st.radio(
+        "Choose input method:",
+        ["Search proteins", "Upload protein ID file"]
     )
+
+    if selection_method == "Search proteins":
+        # Add custom CSS to make multiselect scrollable
+        st.markdown("""
+            <style>
+            [data-testid="stMultiSelect"] div:nth-child(2) {
+                max-height: 200px;
+                overflow-y: auto;
+            }
+            </style>
+            """, unsafe_allow_html=True)
+            
+        selected_proteins = st.multiselect(
+            "Select or search for proteins (UniProt IDs)",
+            options=available_proteins,
+            placeholder="Start typing to search...",
+            max_selections=100
+        )
+
+        if selected_proteins:
+            st.write(f"Selected {len(selected_proteins)} proteins")
+            
+    else:
+        uploaded_file = st.file_uploader(
+            "Upload a text file with UniProt IDs (one per line, max 100)*",
+            type=['txt']
+        )        
+
+        if uploaded_file:
+            protein_list = [line.decode('utf-8').strip() for line in uploaded_file]
+            # Remove empty lines and duplicates
+            protein_list = list(filter(None, protein_list))
+            protein_list = list(dict.fromkeys(protein_list))
+            
+            # filter out proteins that are not in available_proteins
+            protein_list = [p for p in protein_list if p in available_proteins]
+            proteins_not_found = [p for p in protein_list if p not in available_proteins]
+
+            if len(protein_list) > 100:
+                st.error("Please upload a file with maximum 100 protein IDs.")
+                selected_proteins = []
+            else:
+                selected_proteins = protein_list
+                st.write(f"Loaded {len(selected_proteins)} proteins")
+                if proteins_not_found:
+                    st.error(f"Proteins not found in input knowledge graph: {', '.join(proteins_not_found)}")
+                    st.warning("Currently, our system can generate predictions only for proteins included in our input knowledge graph. Real-time retrieval of relationship data from external source databases is not yet supported. However, we are actively working on integrating this capability in future updates.")
     
-    if selected_protein:
-        st.write(f"Selected protein: {selected_protein}")
+    if selected_proteins:
+        # Option 1: Collapsible expander
+        with st.expander("View Selected Proteins"):
+            st.write(f"Total proteins selected: {len(selected_proteins)}")
+            
+            # Create scrollable container with fixed height
+            st.markdown(
+                f"""
+                <div style="
+                    height: 150px; 
+                    overflow-y: scroll;
+                    border: 1px solid #ccc;
+                    border-radius: 4px;
+                    padding: 8px;
+                    background-color: white;">
+                    {'<br>'.join(selected_proteins)}
+                </div>
+                """, 
+                unsafe_allow_html=True
+            )
+
+            st.markdown("<div style='padding-top: 10px;'></div>", unsafe_allow_html=True)
+
+            # Add download button
+            proteins_text = '\n'.join(selected_proteins)
+            st.download_button(
+                label="Download List",
+                data=proteins_text,
+                file_name="selected_proteins.txt",
+                mime="text/plain",
+                key="download_button"
+            )
+
+        # Add GO category selection
+        go_category_options = {
+            'All Categories': None,
+            'Molecular Function': 'GO_term_F',
+            'Biological Process': 'GO_term_P',
+            'Cellular Component': 'GO_term_C'
+        }
+        selected_go_category = st.selectbox(
+            "Select GO Category for predictions",
+            options=list(go_category_options.keys()),
+            help="Choose which GO category to generate predictions for. Selecting 'All Categories' will generate predictions for all three categories."
+        )
+
+    st.warning("⚠️ Due to memory and computational constraints, the maximum number of proteins that can be processed at once is limited to 100 proteins. For larger datasets, please consider running the model locally using our GitHub repository.")
+
+    if selected_proteins and selected_go_category:
+        # Add a button to trigger predictions
+        if st.button("Generate Predictions"):
+            st.session_state.submitted = True
+
+if st.session_state.submitted:
+    with st.spinner("Generating predictions..."):
+        # Generate predictions only if not already in session state
+        if st.session_state.predictions_df is None:
+
+            # Load model config from JSON file
+            import json
+            import os
+
+            # Define data directory path
+            data_dir = "data"
+            models_dir = os.path.join(data_dir, "models")
+
+            # Load model configuration
+            model_config_paths = {
+                'GO_term_F': os.path.join(models_dir, "prothgt-config-molecular-function.yaml"),
+                'GO_term_P': os.path.join(models_dir, "prothgt-config-biological-process.yaml"),
+                'GO_term_C': os.path.join(models_dir, "prothgt-config-cellular-component.yaml")
+            }
+
+            # Paths for model and data
+            model_paths = {
+                'GO_term_F': os.path.join(models_dir, "prothgt-model-molecular-function.pt"),
+                'GO_term_P': os.path.join(models_dir, "prothgt-model-biological-process.pt"),
+                'GO_term_C': os.path.join(models_dir, "prothgt-model-cellular-component.pt")
+            }
+
+            # Get the selected GO category
+            go_category = go_category_options[selected_go_category]
+
+            # If a specific category is selected, use that model path
+            if go_category:
+                model_config_paths = [model_config_paths[go_category]]
+                model_paths = [model_paths[go_category]]
+                go_categories = [go_category]
+            else:
+                model_config_paths = [model_config_paths[cat] for cat in ['GO_term_F', 'GO_term_P', 'GO_term_C']]
+                model_paths = [model_paths[cat] for cat in ['GO_term_F', 'GO_term_P', 'GO_term_C']]
+                go_categories = ['GO_term_F', 'GO_term_P', 'GO_term_C']
+
+            # Generate predictions
+            predictions_df = generate_prediction_df(
+                protein_ids=selected_proteins,
+                model_paths=model_paths,
+                model_config_paths=model_config_paths,
+                go_category=go_categories
+            )
+
+            st.session_state.predictions_df = predictions_df
+        
+        # Display and filter predictions
+        st.success("Predictions generated successfully!")
+        st.markdown("### Filter and View Predictions")
+        
+        # Create filters
+        st.markdown("### Filter Predictions")
+        col1, col2, col3 = st.columns(3)
+        
+        with col1:
+            # Protein filter
+            selected_protein = st.selectbox(
+                "Filter by Protein",
+                options=['All'] + sorted(st.session_state.predictions_df['Protein'].unique().tolist())
+            )
+            
+        with col2:
+            # GO category filter
+            selected_category = st.selectbox(
+                "Filter by GO Category",
+                options=['All'] + sorted(st.session_state.predictions_df['GO_category'].unique().tolist())
+            )
+            
+        with col3:
+            # Probability threshold
+            min_probability_threshold = st.slider(
+                "Minimum Probability",
+                min_value=0.0,
+                max_value=1.0,
+                value=0.5,
+                step=0.05
+            )
+
+            max_probability_threshold = st.slider(
+                "Maximum Probability",
+                min_value=0.0,
+                max_value=1.0,
+                value=1.0,
+                step=0.05
+            )
+
+        # Filter the dataframe using session state data
+        filtered_df = st.session_state.predictions_df.copy()
+
+        if selected_protein != 'All':
+            filtered_df = filtered_df[filtered_df['Protein'] == selected_protein]
+            
+        if selected_category != 'All':
+            filtered_df = filtered_df[filtered_df['GO_category'] == selected_category]
+            
+        filtered_df = filtered_df[(filtered_df['Probability'] >= min_probability_threshold) & 
+                                (filtered_df['Probability'] <= max_probability_threshold)]
+
+        # Sort by probability
+        filtered_df = filtered_df.sort_values('Probability', ascending=False)
+
+
+        # Display the filtered dataframe
+        st.dataframe(
+            filtered_df,
+            hide_index=True,
+            column_config={
+                "Probability": st.column_config.ProgressColumn(
+                    "Probability",
+                    format="%.2f",
+                    min_value=0,
+                    max_value=1,
+                ),
+                "Protein": st.column_config.TextColumn(
+                    "Protein",
+                    help="UniProt ID",
+                ),
+                "GO_category": st.column_config.TextColumn(
+                    "GO Category",
+                    help="Gene Ontology Category",
+                ),
+                "GO_term": st.column_config.TextColumn(
+                    "GO Term",
+                    help="Gene Ontology Term ID",
+                ),
+            }
+        )
+
+        # Download filtered results
+        st.download_button(
+            label="Download Filtered Results",
+            data=convert_df(filtered_df),
+            file_name="filtered_predictions.csv",
+            mime="text/csv",
+            key="download_filtered_predictions"
+        )
+
+# Add a reset button in the sidebar
+with st.sidebar:
+    if st.session_state.submitted:
+        if st.button("Reset"):
+            st.session_state.predictions_df = None
+            st.session_state.submitted = False
+            st.experimental_rerun()

run_prothgt_app.py

@@ -1,33 +1,29 @@
-from datasets import load_dataset
-from torch_geometric.transforms import ToUndirected
 import torch
 from torch.nn import Linear
 from torch_geometric.nn import HGTConv, MLP
 import pandas as pd
+import yaml
+import os
+from datasets import load_dataset
 
 class ProtHGT(torch.nn.Module):
     def __init__(self, data,hidden_channels, num_heads, num_layers, mlp_hidden_layers, mlp_dropout):
         super().__init__()
 
-        self.lin_dict = torch.nn.ModuleDict({
-            node_type: Linear(-1, hidden_channels)
-            for node_type in data.node_types
-        })
+        self.lin_dict = torch.nn.ModuleDict()
+        for node_type in data.node_types:
+            input_dim = data[node_type].x.size(1)  # Get actual input dimension from data
+            self.lin_dict[node_type] = Linear(input_dim, hidden_channels)
 
         self.convs = torch.nn.ModuleList()
         for _ in range(num_layers):
             conv = HGTConv(hidden_channels, hidden_channels, data.metadata(), num_heads, group='sum')
             self.convs.append(conv)
         
-        # self.left_linear = Linear(hidden_channels, hidden_channels)
-        # self.right_linear = Linear(hidden_channels, hidden_channels)
-        # self.sqrt_hd  = hidden_channels**1/2
-
-        # self.mlp =MLP([2*hidden_channels, 128, 1], dropout=0.5, norm=None)
         self.mlp = MLP(mlp_hidden_layers , dropout=mlp_dropout, norm=None)
 
     def generate_embeddings(self, x_dict, edge_index_dict):
-        # Generate updated embeddings through the GNN layers
+        # Generate updated embeddings through the HGT layers
         x_dict = {
             node_type: self.lin_dict[node_type](x).relu_()
             for node_type, x in x_dict.items()
@@ -48,9 +44,11 @@ class ProtHGT(torch.nn.Module):
 
         return self.mlp(z).view(-1), x_dict
 
-def _load_data(protein_id, go_category=None, heterodata_path=''):
-    heterodata = load_dataset(heterodata_path)
+def _load_data(protein_ids, go_category=None):
     
+    # heterodata = load_dataset('HUBioDataLab/ProtHGT-KG', data_files="prothgt-kg.pt")
+    heterodata = torch.load('data/prothgt-kg.pt')
+    print('Loading data...')
     # Remove unnecessary edge types in one go
     edge_types_to_remove = [
         ('Protein', 'protein_function', 'GO_term_F'),
@@ -62,68 +60,123 @@ def _load_data(protein_id, go_category=None, heterodata_path=''):
     ]
     
     for edge_type in edge_types_to_remove:
-        if edge_type in heterodata:
-            del heterodata[edge_type]
+        if edge_type in heterodata.edge_index_dict:
+            del heterodata.edge_index_dict[edge_type]
 
-    # Remove reverse edges
-    heterodata = {k: v for k, v in heterodata.items() if not isinstance(k, tuple) or 'rev' not in k[1]}
-
-    protein_index = heterodata['Protein']['id_mapping'][protein_id]
+    # Get protein indices for all input proteins
+    protein_indices = [heterodata['Protein']['id_mapping'][pid] for pid in protein_ids]
     
-    # Create edge indices more efficiently
+    # Create edge indices for prediction
     categories = [go_category] if go_category else ['GO_term_F', 'GO_term_P', 'GO_term_C']
     
     for category in categories:
-        pairs = [(protein_index, i) for i in range(len(heterodata[category]))]
-        heterodata['Protein', 'protein_function', category] = {'edge_index': pairs}
-    
-    return ToUndirected(merge=False)(heterodata)
+        # Create pairs for all proteins with all GO terms
+        n_terms = len(heterodata[category]['id_mapping'])
+        protein_indices_repeated = torch.tensor(protein_indices).repeat_interleave(n_terms)
+        term_indices = torch.arange(n_terms).repeat(len(protein_indices))
+        
+        edge_index = torch.stack([protein_indices_repeated, term_indices])
+        heterodata.edge_index_dict[('Protein', 'protein_function', category)] = edge_index
+
+    return heterodata
 
 def get_available_proteins(protein_list_file='data/available_proteins.txt'):
     with open(protein_list_file, 'r') as file:
         return [line.strip() for line in file.readlines()]
 
-def _generate_predictions(heterodata, model_path, model_config, target_type):
-
+def _generate_predictions(heterodata, model, target_type):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = ProtHGT(heterodata, model_config['hidden_channels'], model_config['num_heads'], model_config['num_layers'], model_config['mlp_hidden_layers'], model_config['mlp_dropout'])
-    print('Loading model from', model_path)
-    model.load_state_dict(torch.load(model_path, map_location=device))
-
+    
     model.to(device)
     model.eval()
-    heterodata.to(device)
+    heterodata = heterodata.to(device)
 
     with torch.no_grad():
-        predictions, _ = model(heterodata.x_dict, heterodata.edge_index_dict, heterodata[("Protein", "protein_function", target_type)].edge_label_index, target_type)
-    return predictions
-
-def _create_prediction_df(predictions, heterodata, protein_id, go_category):
-    prediction_df = pd.DataFrame({
-        'Protein': protein_id,
-        'GO_category': go_category,
-        'GO_term': heterodata[go_category]['id_mapping'].keys(),
-        'Probability': predictions.tolist()
-    })
-    prediction_df.sort_values(by='Probability', ascending=False, inplace=True)
-    prediction_df.reset_index(drop=True, inplace=True)
-    return prediction_df
-
-
-def generate_prediction_df(protein_id, heterodata_path, model_path, model_config, go_category=None):
-    heterodata = _load_data(protein_id, go_category, heterodata_path)
-
-    if go_category:
-        predictions = _generate_predictions(heterodata, model_path, model_config, go_category)
-        prediction_df = _create_prediction_df(predictions, heterodata, protein_id, go_category)
-        return prediction_df
+        edge_label_index = heterodata.edge_index_dict[('Protein', 'protein_function', target_type)]
+        predictions, _ = model(heterodata.x_dict, heterodata.edge_index_dict, edge_label_index, target_type)
+        predictions = torch.sigmoid(predictions)
+    
+    return predictions.cpu()
+
+def _create_prediction_df(predictions, heterodata, protein_ids, go_category):
+    go_category_dict = {
+        'GO_term_F': 'Molecular Function',
+        'GO_term_P': 'Biological Process',
+        'GO_term_C': 'Cellular Component'
+    }
+    # Create a list to store individual protein predictions
+    all_predictions = []
+    
+    # Number of GO terms for this category
+    n_go_terms = len(heterodata[go_category]['id_mapping'])
+    
+    # Process predictions for each protein
+    for i, protein_id in enumerate(protein_ids):
+        # Get the slice of predictions for this protein
+        protein_predictions = predictions[i * n_go_terms:(i + 1) * n_go_terms]
+        
+        prediction_df = pd.DataFrame({
+            'Protein': protein_id,
+            'GO_category': go_category_dict[go_category],
+            'GO_term': list(heterodata[go_category]['id_mapping'].keys()),
+            'Probability': protein_predictions.numpy()
+        })
+        all_predictions.append(prediction_df)
+    
+    # Combine all predictions
+    combined_df = pd.concat(all_predictions, ignore_index=True)
+    combined_df.sort_values(by=['Protein', 'Probability'], ascending=[True, False], inplace=True)
+    combined_df.reset_index(drop=True, inplace=True)
+    return combined_df
+
+def generate_prediction_df(protein_ids, model_paths, model_config_paths, go_category):
+    all_predictions = []
+    
+    # Convert single protein ID to list if necessary
+    if isinstance(protein_ids, str):
+        protein_ids = [protein_ids]
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    for go_cat, model_config_path, model_path in zip(go_category, model_config_paths, model_paths):
+        print(f'Generating predictions for {go_cat}...')
+            
+        # Load data
+        heterodata = _load_data(protein_ids, go_cat)
+        
+        # Load model configuration
+        with open(model_config_path, 'r') as file:
+            model_config = yaml.safe_load(file)
+        
+        # Initialize model with configuration
+        model = ProtHGT(
+            heterodata,
+            hidden_channels=model_config['hidden_channels'][0],
+            num_heads=model_config['num_heads'],
+            num_layers=model_config['num_layers'],
+            mlp_hidden_layers=model_config['hidden_channels'][1],
+            mlp_dropout=model_config['mlp_dropout']
+        )
+        
+        # Load model weights
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        print(f'Loaded model weights from {model_path}')
+        
+        # Generate predictions
+        predictions = _generate_predictions(heterodata, model, go_cat)
+        prediction_df = _create_prediction_df(predictions, heterodata, protein_ids, go_cat)
+        all_predictions.append(prediction_df)
+        
+        # Clean up memory
+        del heterodata
+        del model
+        del predictions
+        torch.cuda.empty_cache()  # Clear CUDA cache if using GPU
+
+    # Combine all predictions
+    final_df = pd.concat(all_predictions, ignore_index=True)
     
-    else:
-        all_predictions = []
-        for go_category in ['GO_term_F', 'GO_term_P', 'GO_term_C']:
-            predictions = _generate_predictions(heterodata, model_path, model_config, go_category)
-            category_df = _create_prediction_df(predictions, heterodata, protein_id, go_category)
-            all_predictions.append(category_df)
-
-        return pd.concat(all_predictions, ignore_index=True)
+    # Clean up
+    del all_predictions
+    torch.cuda.empty_cache()
     
+    return final_df