Source code for graph.StringLoader

from __future__ import annotations

from typing import Any, Optional, List
import requests
import pandas as pd
import json
from html import escape

from graph.Loader import Loader


[docs] class StringLoader(Loader): """ Retrieve protein-protein interaction data from the STRING database. This loader queries the STRING API with a list of protein identifiers, maps them to STRING identifiers, and retrieves the corresponding PPI network. Attributes: base_url (str): The base URL for the STRING API. protein_query (List[str]): List of input protein identifiers. species (int): NCBI taxonomy ID for the organism. required_score (int): Minimum confidence score for interactions. network_type (str): Type of network ("functional" or "physical"). add_nodes (int): Number of additional nodes to include. raw_mapping (List[dict]): Raw mapping response from STRING. raw_network (List[dict]): Raw network response from STRING. Example: >>> loader = StringLoader( ... protein_query=["TP53", "BRCA1", "EGFR"], ... species=9606, ... add_nodes=10, ... ) >>> data = ( ... loader ... .retrieve_data() ... .standardize_data_format() ... .get_data() ... ) >>> print(data["nodes"].head()) >>> print(data["edges"].head()) """ base_url = "https://version-12-0.string-db.org/api" def __init__( self, protein_query: List[str], species: int = 9606, caller_identity: str = "StringLoaderClass", required_score: int = 400, network_type: str = "functional", add_nodes: int = 10, mapping_limit: int = 1, timeout: int = 30, ): """ Initialize the StringLoader with configuration parameters. Args: protein_query: List of protein identifiers to query (gene names, UniProt IDs, etc.). species: NCBI taxonomy ID for the species. Defaults to 9606 (human). caller_identity: Identifier for the caller in STRING API requests. Defaults to "StringLoaderClass". required_score: Minimum confidence score for interactions (0-1000). Defaults to 400. network_type: Type of network to retrieve ("functional" or "physical"). Defaults to "functional". add_nodes: Number of additional nodes to include in the retrieved network. Defaults to 10. mapping_limit: Maximum number of STRING identifiers to return per input protein. Defaults to 1. timeout: Timeout for HTTP requests in seconds. Defaults to 30. """ super().__init__() self.protein_query = protein_query self.species = species self.caller_identity = caller_identity self.required_score = required_score self.network_type = network_type self.add_nodes = add_nodes self.mapping_limit = mapping_limit self.timeout = timeout self.raw_mapping: list[dict[str, Any]] = [] self.raw_network: list[dict[str, Any]] = [] def _post_json(self, method: str, params: dict[str, Any]) -> list[dict[str, Any]]: """ Send a POST request to the STRING API and return the JSON response. Args: method: The STRING API method name (e.g., "get_string_ids", "network"). params: Dictionary of parameters to send in the POST request. Returns: List[dict]: The JSON response from the STRING API. Raises: ValueError: If the API returns a 404 status code. requests.HTTPError: If the API returns any other error status code. """ url = f"{self.base_url}/json/{method}" response = requests.post( url, data=params, timeout=self.timeout, ) if response.status_code == 404: raise ValueError( f"No result has been found in STRING for: {self.protein_query!r}" ) response.raise_for_status() return response.json()
[docs] def retrieve_data(self) -> "StringLoader": """ Retrieve and map protein identifiers from STRING, then fetch the network. This method performs two API calls: 1. Maps input proteins to STRING identifiers (get_string_ids) 2. Retrieves the PPI network for the mapped identifiers (network) Returns: StringLoader: Self for method chaining. Raises: ValueError: If no identifiers are found for the input proteins. """ # Build parameters for the STRING identifier mapping request mapping_params = { "identifiers": "\r".join(self.protein_query), "species": self.species, "echo_query": 1, "limit": self.mapping_limit, "caller_identity": self.caller_identity, } # Retrieve STRING identifiers for the input proteins self.raw_mapping = self._post_json( method="get_string_ids", params=mapping_params, ) # Stop if STRING returned no mappings if not self.raw_mapping: raise ValueError( f"No identifier has been found in STRING for: {self.protein_query!r}" ) # Select one STRING identifier for each input protein selected_mapping_df = self._get_selected_identifier_rows() # Extract STRING IDs as a list selected_string_ids = selected_mapping_df["stringId"].tolist() # Build parameters for the STRING network request network_params = { "identifiers": "\r".join(selected_string_ids), "species": self.species, "required_score": self.required_score, "network_type": self.network_type, "add_nodes": self.add_nodes, "caller_identity": self.caller_identity, } # Retrieve the network using all selected STRING identifiers self.raw_network = self._post_json( method="network", params=network_params, ) return self
def _get_selected_identifier_rows(self) -> pd.DataFrame: """ Select one STRING identifier for each input protein. This method ensures that exactly one STRING identifier is returned for each input protein in the same order as the input list. Returns: pd.DataFrame: DataFrame with one row per input protein containing the selected STRING identifier and metadata. Raises: RuntimeError: If retrieve_data() has not been called first. RuntimeError: If the STRING response is missing the queryIndex column. ValueError: If one or more input proteins could not be mapped to STRING identifiers. """ # Check that STRING mapping data exists if not self.raw_mapping: raise RuntimeError("No raw_mapping presente, call retrieve_data() first") # Convert raw mapping response to DataFrame identifiers_df = pd.DataFrame(self.raw_mapping) # Check that STRING returned the queryIndex column if "queryIndex" not in identifiers_df.columns: raise RuntimeError("STRING mapping response does not contain queryIndex") # Sort mappings according to the original input order identifiers_df = identifiers_df.sort_values("queryIndex") # Keep one mapping row for each input protein selected_df = ( identifiers_df .groupby("queryIndex", as_index=False) .first() .sort_values("queryIndex") .reset_index(drop=True) ) # Check which input proteins were successfully mapped found_indexes = set(selected_df["queryIndex"].astype(int)) expected_indexes = set(range(len(self.protein_query))) # Detect proteins without a STRING identifier missing_indexes = sorted(expected_indexes - found_indexes) # Raise an error if one or more proteins were not mapped if missing_indexes: missing_queries = [ self.protein_query[i] for i in missing_indexes ] raise ValueError( f"No STRING identifier has been found for: {missing_queries!r}" ) return selected_df
[docs] def standardize_data_format(self) -> "StringLoader": """ Standardize and organize retrieved PPI data into a structured format. Converts raw STRING API responses into a dictionary containing standardized DataFrames for nodes, edges, and metadata. Returns: StringLoader: Self for method chaining. Raises: RuntimeError: If retrieve_data() has not been called first. Note: After calling this method, use get_data() to retrieve the standardized data. """ # Check that the mapping step has already been executed if not self.raw_mapping: raise RuntimeError("No raw_mapping presente, call retrieve_data() first") # Convert raw mapping and network responses to DataFrames identifiers_df = pd.DataFrame(self.raw_mapping) edges_df = pd.DataFrame(self.raw_network) # Select one STRING identifier for each input protein selected_mapping_df = self._get_selected_identifier_rows() # Extract selected STRING IDs selected_string_ids = selected_mapping_df["stringId"].tolist() # Build the node table from the edge table nodes_df = self._build_nodes_dataframe(edges_df) # Store standardized data self.data = { "query": self.protein_query, "species": self.species, "selected_string_ids": selected_string_ids, "proteins": selected_mapping_df, "identifiers": identifiers_df, "edges": edges_df, "nodes": nodes_df, "raw": { "mapping": self.raw_mapping, "network": self.raw_network, }, } return self
@staticmethod def _build_nodes_dataframe(edges_df: pd.DataFrame) -> pd.DataFrame: """ Build a nodes DataFrame from an edges DataFrame. Extracts unique nodes from both endpoints of edges and removes duplicates. Args: edges_df: DataFrame containing edges with columns stringId_A, stringId_B, preferredName_A, preferredName_B, and ncbiTaxonId. Returns: pd.DataFrame: DataFrame with columns [stringId, preferredName, ncbiTaxonId], sorted by preferredName and deduplicated. """ if edges_df.empty: return pd.DataFrame( columns=["stringId", "preferredName", "ncbiTaxonId"] ) nodes_a = edges_df[ ["stringId_A", "preferredName_A", "ncbiTaxonId"] ].rename( columns={ "stringId_A": "stringId", "preferredName_A": "preferredName", } ) nodes_b = edges_df[ ["stringId_B", "preferredName_B", "ncbiTaxonId"] ].rename( columns={ "stringId_B": "stringId", "preferredName_B": "preferredName", } ) nodes = pd.concat([nodes_a, nodes_b], ignore_index=True) nodes = nodes.drop_duplicates(subset=["stringId"]) nodes = nodes.sort_values("preferredName").reset_index(drop=True) return nodes
[docs] def get_network_html_div( self, width: str = "100%", height: str = "650px", network_flavor: str = "confidence", add_color_nodes: Optional[int] = None, add_white_nodes: Optional[int] = None, hide_node_labels: bool = False, hide_disconnected_nodes: bool = False, show_query_node_labels: bool = True, block_structure_pics_in_bubbles: bool = False, include_script: bool = True, use_selected_string_ids: bool = True, ) -> str: """ Generate an interactive STRING network HTML block for embedding in reports. The returned HTML can be inserted directly into an HTML report and will render an interactive Cytoscape visualization of the protein network. Args: width: CSS width for the network container. Defaults to "100%". height: CSS height for the network container. Defaults to "650px". network_flavor: Type of interaction to display ("confidence", "evidence", or "actions"). Defaults to "confidence". add_color_nodes: Number of additional nodes with color highlighting. Defaults to self.add_nodes. add_white_nodes: Number of additional nodes without color highlighting. Defaults to self.add_nodes. hide_node_labels: If True, hide labels on nodes. Defaults to False. hide_disconnected_nodes: If True, hide nodes without edges. Defaults to False. show_query_node_labels: If True, show labels only on query proteins. Defaults to True. block_structure_pics_in_bubbles: If True, disable structure pictures in node bubbles. Defaults to False. include_script: If True, include the STRING JavaScript library. Defaults to True. use_selected_string_ids: If True, use mapped STRING IDs; otherwise use original protein_query. Defaults to True. Returns: str: HTML string containing the embedded STRING network with required scripts. Note: - STRING expects the embedded network container to have id="stringEmbedded". - This method is intended for one STRING embedded network per HTML page. - If retrieve_data() has already been called, the method uses mapped STRING IDs. Otherwise, it falls back to the original protein_query values. """ if add_color_nodes is None: add_color_nodes = self.add_nodes elif add_white_nodes is None: add_white_nodes = self.add_nodes identifiers = self.protein_query if use_selected_string_ids and self.raw_mapping: selected_mapping_df = self._get_selected_identifier_rows() identifiers = selected_mapping_df["stringId"].tolist() string_web_url = self.base_url.rsplit("/api", 1)[0] string_params = { "species": self.species, "identifiers": identifiers, "required_score": self.required_score, "network_type": self.network_type, "network_flavor": network_flavor, "add_color_nodes": add_color_nodes, "add_white_nodes": add_white_nodes, "hide_node_labels": int(hide_node_labels), "hide_disconnected_nodes": int(hide_disconnected_nodes), "show_query_node_labels": int(show_query_node_labels), "block_structure_pics_in_bubbles": int(block_structure_pics_in_bubbles), "caller_identity": self.caller_identity, } params_json = json.dumps(string_params) script_tag = "" if include_script: script_tag = """ <script type="text/javascript" src="https://string-db.org/javascript/combined_embedded_network_v2.0.6.js"> </script> """ return f""" <div class="string-network-report" style="width: {escape(width)}; min-height: {escape(height)};"> {script_tag} <div id="stringEmbedded" style="width: 100%; min-height: {escape(height)};"> </div> <script type="text/javascript"> (function () {{ function renderStringNetwork() {{ getSTRING({json.dumps(string_web_url)}, {params_json}); }} if (document.readyState === "loading") {{ document.addEventListener("DOMContentLoaded", renderStringNetwork); }} else {{ renderStringNetwork(); }} }})(); </script> </div> """