Upload proteins_viz.py

proteins_viz.py

@@ -0,0 +1,136 @@
+import pandas as pd
+from biopandas.pdb import PandasPdb
+from prody import parsePDBHeader
+
+
+
+
+def read_pdb_to_dataframe(
+    pdb_path,
+    model_index: int = 1,
+    parse_header: bool = True,
+    ) -> pd.DataFrame:
+    """
+    Read a PDB file, and return a Pandas DataFrame containing the atomic coordinates and metadata.
+
+    Args:
+        pdb_path (str, optional): Path to a local PDB file to read. Defaults to None.
+        model_index (int, optional): Index of the model to extract from the PDB file, in case
+            it contains multiple models. Defaults to 1.
+        parse_header (bool, optional): Whether to parse the PDB header and extract metadata.
+            Defaults to True.
+
+    Returns:
+        pd.DataFrame: A DataFrame containing the atomic coordinates and metadata, with one row
+            per atom
+    """
+    atomic_df = PandasPdb().read_pdb(pdb_path)
+    if parse_header:
+        header = parsePDBHeader(pdb_path)
+    else:
+        header = None
+    atomic_df = atomic_df.get_model(model_index)
+    if len(atomic_df.df["ATOM"]) == 0:
+        raise ValueError(f"No model found for index: {model_index}")
+
+    return pd.concat([atomic_df.df["ATOM"], atomic_df.df["HETATM"]]), header
+
+from graphein.protein.graphs import label_node_id
+
+def process_dataframe(df: pd.DataFrame, granularity='CA') -> pd.DataFrame:
+    """
+    Process a DataFrame of protein structure data to reduce ambiguity and simplify analysis.
+
+    This function performs the following steps:
+    1. Handles alternate locations for an atom, defaulting to keep the first one if multiple exist.
+    2. Assigns a unique node_id to each residue in the DataFrame, using a helper function label_node_id.
+    3. Filters the DataFrame based on specified granularity (defaults to 'CA' for alpha carbon).
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        The DataFrame containing protein structure data to process. It is expected to contain columns 'alt_loc' and 'atom_name'.
+        
+    granularity : str, optional
+        The level of detail or perspective at which the DataFrame should be analyzed. Defaults to 'CA' (alpha carbon).
+    """
+    # handle the case of alternative locations,
+    # if so default to the 1st one = A
+    if 'alt_loc' in df.columns:
+      df['alt_loc'] = df['alt_loc'].replace('', 'A')
+      df = df.loc[(df['alt_loc']=='A')]
+    df = label_node_id(df, granularity)
+    df = df.loc[(df['atom_name']==granularity)]
+    return df
+
+
+from graphein.protein.graphs import initialise_graph_with_metadata
+from graphein.protein.graphs import add_nodes_to_graph
+from graphein.protein.visualisation import plotly_protein_structure_graph
+from PIL import Image
+import networkx as nx
+
+def take_care(pdb_path):
+    
+
+    df, header = read_pdb_to_dataframe(pdb_path)
+    process_df = process_dataframe(df)
+
+    g = initialise_graph_with_metadata(protein_df=process_df, # from above cell
+                                        raw_pdb_df=df, # Store this for traceability
+                                        pdb_code = '3nir', #and again
+                                        granularity = 'CA' # Store this so we know what kind of graph we have
+                                        )
+    g = add_nodes_to_graph(g)
+
+    
+    def add_backbone_edges(G: nx.Graph) -> nx.Graph:
+        # Iterate over every chain
+        for chain_id in G.graph["chain_ids"]:
+            # Find chain residues
+            chain_residues = [
+                (n, v) for n, v in G.nodes(data=True) if v["chain_id"] == chain_id
+            ]
+            # Iterate over every residue in chain
+            for i, residue in enumerate(chain_residues):
+                try:
+                    # Checks not at chain terminus
+                    if i == len(chain_residues) - 1:
+                        continue
+                    # Asserts residues are on the same chain
+                    cond_1 = ( residue[1]["chain_id"] == chain_residues[i + 1][1]["chain_id"])
+                    # Asserts residue numbers are adjacent
+                    cond_2 = (abs(residue[1]["residue_number"] - chain_residues[i + 1][1]["residue_number"])== 1)
+
+                    # If this checks out, we add a peptide bond
+                    if (cond_1) and (cond_2):
+                        # Adds "peptide bond" between current residue and the next
+                        if G.has_edge(i, i + 1):
+                            G.edges[i, i + 1]["kind"].add('backbone_bond')
+                        else:
+                            G.add_edge(residue[0],chain_residues[i + 1][0],kind={'backbone_bond'},)
+                except IndexError as e:
+                    print(e)
+        return G
+
+    g = add_backbone_edges(g)
+
+    
+
+    p = plotly_protein_structure_graph(
+        g,
+        colour_edges_by="kind",
+        colour_nodes_by="seq_position",
+        label_node_ids=False,
+        plot_title="Backbone Protein Graph",
+        node_size_multiplier=1,
+    )
+    image_file = "protein_graph.png"
+    p.write_image(image_file, format='png')
+
+
+    # Load the PNG image into a PIL image
+    image = Image.open(image_file)
+
+    
+    return image