Protein_ligand graph support & junction tree feature

.requirements/extras.in

@@ -1,6 +1,7 @@
 biovec
 propy3
 pyaaisc
+prody
 mpl_chord_diagram==0.3.2
 rdkit
 selfies

graphein/molecule/graphs.py

@@ -26,8 +26,7 @@
 
 from .chembl import get_smiles_from_chembl
 from .config import MoleculeGraphConfig
-from .edges.atomic import add_atom_bonds
-from .utils import compute_fragments
+from .utils import compute_fragments, get_clique_mol, get_smiles, tree_decomp
 from .zinc import get_smiles_from_zinc
 
 try:
@@ -136,6 +135,7 @@ def construct_graph(
     smiles: Optional[str] = None,
     zinc_id: Optional[str] = None,
     chembl_id: Optional[str] = None,
+    junction_tree: bool = False,
     generate_conformer: Optional[bool] = False,
     edge_construction_funcs: Optional[str] = None,
     edge_annotation_funcs: Optional[List[Callable]] = None,
@@ -166,6 +166,9 @@ def construct_graph(
     :type zinc_id: str, optional
     :param chembl_id: ChEMBL ID to build graph from. Default is ``None``.
     :type chembl_id: str, optional
+    :param junction_tree: boolean to indicate whether to use a junction tree or not. Default is ``False``.
+    :type junction_tree: bool
+    :param generate_conformer: Whether to generate a conformer for the molecule. Defaults to ``False``.
     :param generate_conformer: Whether to generate a conformer for the molecule.
         Defaults to ``False``.
     :type generate_conformer: bool, optional
@@ -271,30 +274,21 @@ def construct_graph(
     if config.add_hs:
         rdmol = Chem.AddHs(rdmol)
 
-    if coords is None:
-        # If no coords are provided, add edges by bonds
-        config.edge_construction_functions = [add_atom_bonds]
-        g = initialise_graph_with_metadata(
-            name=name,
-            rdmol=rdmol,
-            coords=None,
-        )
-    else:
-        # If config params are provided, overwrite them
-        config.edge_construction_functions = (
-            edge_construction_funcs
-            if config.edge_construction_functions is None
-            else config.edge_construction_functions
-        )
-
-        g = initialise_graph_with_metadata(
-            name=name,
-            rdmol=rdmol,
-            coords=np.asarray(coords),
-        )
+    # If no coords are provided, add edges by bonds
+    # config.edge_construction_functions = [add_atom_bonds]
+    g = initialise_graph_with_metadata(
+        name=name,
+        rdmol=rdmol,
+        coords=coords,
+    )
 
     # Add nodes to graph
-    g = add_nodes_to_graph(g)
+    if junction_tree:
+        jt = construct_junction_tree(rdmol)
+        jt.graph = g.graph
+        g = jt
+    else:
+        g = add_nodes_to_graph(g)
 
     # Add config to graph
     g.graph["config"] = config
@@ -320,6 +314,34 @@ def construct_graph(
     return g
 
 
+def construct_junction_tree(mol=Chem.Mol) -> nx.Graph:
+    """
+    Constructs molecule structure junction tree graph from a ``smiles``.
+
+    :param mol: smiles string to build graph from. Default is ``None``.
+    :type mol: str, optional # TODO docstring
+    :return: Molecule Structure Junction Tree Graph
+    :type: nx.Graph
+    """
+    g = nx.Graph()
+
+    cliques, edges = tree_decomp(mol)
+    clique_map = {}
+    for i, c in enumerate(cliques):
+        cmol = get_clique_mol(mol, c)
+        id = f"{get_smiles(cmol)}:{str(i)}"
+        g.add_node(id)
+        clique_map[i] = id
+
+    for n1, n2 in edges:
+        n1, n2 = clique_map[n1], clique_map[n2]
+        if g.has_edge(n1, n2):
+            g.edges[n1, n2]["kind"].add("junction_tree")
+        else:
+            g.add_edge(n1, n2, kind={"junction_tree"})
+    return g
+
+
 def compute_fragment_graphs(
     g: Union[nx.Graph, Chem.Mol], config: Optional[MoleculeGraphConfig] = None
 ) -> List[nx.Graph]:

graphein/molecule/utils.py

@@ -213,7 +213,7 @@ def copy_edit_mol(mol: rdkit.Chem.rdchem.Mol) -> rdkit.Chem.rdchem.Mol:
     """
     new_mol = Chem.RWMol(Chem.MolFromSmiles(""))
     for atom in mol.GetAtoms():
-        new_atom = copy_atom(atom)
+        new_atom = copy_rdmol_atom(atom)
         new_mol.AddAtom(new_atom)
     for bond in mol.GetBonds():
         a1 = bond.GetBeginAtom().GetIdx()

graphein/protein_ligand/__init__.py

@@ -0,0 +1,22 @@
+"""Protein graph construction module."""
+from .config import *
+from .edges import *
+from .graphs import *
+from .utils import *
+from .visualisation import (
+    plot_distance_landscape,
+    plot_distance_matrix,
+    plot_protein_structure_graph,
+    plotly_protein_structure_graph,
+)
+
+try:
+    from .visualisation import plot_chord_diagram
+except ImportError:
+    pass
+
+try:
+    from .meshes import *
+    from .visualisation import plot_pointcloud
+except ImportError:
+    pass

graphein/protein_ligand/config.py

@@ -0,0 +1,237 @@
+"""Base Config object for use with Protein Graph Construction."""
+# Graphein
+# Author: Arian Jamasb <[email protected]>
+# License: MIT
+# Project Website: https://github.com/a-r-j/graphein
+# Code Repository: https://github.com/a-r-j/graphein
+from __future__ import annotations
+
+from functools import partial
+from pathlib import Path
+from typing import Any, Callable, List, Optional, Union
+
+from deepdiff import DeepDiff
+from pydantic import BaseModel
+from typing_extensions import Literal
+
+from graphein.protein_ligand.edges.atomic import (
+    add_atom_bonds,
+    add_atomic_edges,
+)
+from graphein.protein_ligand.edges.distance import (
+    add_distance_threshold_ligand,
+    add_distance_threshold_protein_ligand,
+    add_fully_connected_edges_ligand,
+    add_fully_connected_edges_protein_ligand,
+    add_k_nn_edges_ligand,
+    add_k_nn_edges_protein_ligand,
+    add_peptide_bonds,
+)
+from graphein.protein_ligand.features.nodes.amino_acid import meiler_embedding
+from graphein.protein_ligand.features.nodes.atom_type import atom_type_one_hot
+from graphein.utils.config import PartialMatchOperator, PathMatchOperator
+
+
+class DSSPConfig(BaseModel):
+    executable: str = "mkdssp"
+
+
+class GetContactsConfig(BaseModel):
+    """
+    Config object for parameters relating to running ``GetContacts``.
+    ``GetContacts`` is an optional dependency from which intramolecular interactions can be computed and used as edges in the graph.
+
+    More information about ``GetContacts`` can be found at https://getcontacts.github.io/
+
+    :param get_contacts_path: Path to ``GetContacts`` installation
+    :type get_contacts_path: pathlib.Path
+    :param contacts_dir: Path to store output of ``GetContacts``
+    :type contacts_dir: pathlib.Path
+    :param pdb_dir: Path to PDB files to be used to compute intramolecular interactions.
+    :type pdb_dir: pathlib.Path
+    :param granularity: Specifies the node types of the graph, defaults to ``"CA"`` for alpha-carbons as nodes. Other options are ``"CB"`` (beta-carbon), ``"atom"`` for all-atom graphs, and ``"centroid"`` for nodes positioned as residue centroids.
+    :type granularity: str
+    """
+
+    get_contacts_path: Path = Path(
+        "/Users/arianjamasb/github/getcontacts/"
+    ).resolve()
+    contacts_dir: Path = Path("../examples/contacts/").resolve()
+    pdb_dir: Path = Path("../examples/pdbs/").resolve()
+    granularity: str = "CA"
+
+
+GraphAtoms = Literal[
+    "N",
+    "CA",
+    "C",
+    "O",
+    "CB",
+    "OG",
+    "CG",
+    "CD1",
+    "CD2",
+    "CE1",
+    "CE2",
+    "CZ",
+    "OD1",
+    "ND2",
+    "CG1",
+    "CG2",
+    "CD",
+    "CE",
+    "NZ",
+    "OD2",
+    "OE1",
+    "NE2",
+    "OE2",
+    "OH",
+    "NE",
+    "NH1",
+    "NH2",
+    "OG1",
+    "SD",
+    "ND1",
+    "SG",
+    "NE1",
+    "CE3",
+    "CZ2",
+    "CZ3",
+    "CH2",
+    "OXT",
+]
+"""Allowable atom types for nodes in the graph."""
+
+GranularityOpts = Literal["atom", "centroids"]
+"""Allowable granularity options for nodes in the graph."""
+
+
+class ProteinLigandGraphConfig(BaseModel):
+    """
+    Config Object for Protein-ligand Structure Graph Construction.
+
+    If you encounter a problematic structure, perusing https://www.umass.edu/microbio/chime/pe_beta/pe/protexpl/badpdbs.htm may provide some additional insight.
+    PDBs are notoriously troublesome and this is an excellent overview.
+
+    :param granularity: Controls the granularity of the graph construction. ``"atom"`` builds an atomic-scale graph where
+        nodes are constituent atoms. Residue-level graphs can be build by specifying which constituent atom should
+        represent node positions (see :const:`~graphein.protein.config.GraphAtoms`). Additionally, ``"centroids"`` can be specified to
+        compute the centre of gravity for a given atom (Specified in :const:`~graphein.protein.config.GranularityOpts`).
+        Defaults to ``"CA"`` (alpha-Carbon).
+    :type granularity: str (Union[graphein.protein.config.GraphAtoms, graphein.protein.config.GranularityOpts])
+    :param keep_hets: Controls whether or not heteroatoms are removed from the PDB file. These are typically modified
+        residues, bound ligands, crystallographic adjuvants, ions or water molecules.
+
+        For more information, see: https://proteopedia.org/wiki/index.php/Hetero_atoms
+    :type keep_hets: bool
+    :param insertions: Controls whether or not insertions are allowed.
+    :type insertions: bool
+    :param pdb_dir: Specifies path to download protein structures into.
+    :type pdb_dir: pathlib.Path
+    :param verbose: Specifies verbosity of graph creation process.
+    :type verbose: bool
+    :param exclude_waters: Specifies whether or not water molecules are excluded from the structure
+    :type excluded_waters: bool
+    :param deprotonate: Specifies whether or not to remove ``H`` atoms from the graph.
+    :type deprotonate: bool
+    :param protein_df_processing_functions: List of functions that take a ``pd.DataFrame`` and return a ``pd.DataFrame``.
+        This allows users to define their own series of processing functions for the protein structure DataFrame and
+        override the default sequencing of processing steps provided by Graphein. We refer users to our low-level API
+        tutorial for more details.
+    :type protein_df_processing_functions: Optional[List[Callable]]
+    :param protein_edge_construction_functions: List of functions that take an ``nx.Graph`` and return an ``nx.Graph`` with desired
+        edges added. Prepared edge constructions can be found in :ref:`graphein.protein_ligand.edges`
+    :type edge_construction_functions: List[Callable]
+    :param protein_node_metadata_functions: List of functions that take an ``nx.Graph``
+    :type node_metadata_functions: List[Callable], optional
+    :param protein_edge_metadata_functions: List of functions that take an
+    :type edge_metadata_functions: List[Callable], optional
+    :param ligand_edge_construction_functions: List of functions that take an ``nx.Graph`` and return an ``nx.Graph`` with desired
+    edges added. Prepared edge constructions can be found in :ref:`graphein.protein_ligand.edges`
+    :type edge_construction_functions: List[Callable]
+    :param ligand_node_metadata_functions: List of functions that take an ``nx.Graph``
+    :type node_metadata_functions: List[Callable], optional
+    :param ligand_edge_metadata_functions: List of functions that take an
+    :type edge_metadata_functions: List[Callable], optional
+    :param protein_ligand_edge_construction_functions: List of functions that take an ``nx.Graph`` and return an ``nx.Graph`` with desired
+    edges added. Prepared edge constructions can be found in :ref:`graphein.protein_ligand.edges`
+    :type edge_construction_functions: List[Callable]
+    :param protein_ligand_edge_metadata_functions: List of functions that take an
+    :type edge_metadata_functions: List[Callable], optional
+    :param graph_metadata_functions: List of functions that take an ``nx.Graph`` and return an ``nx.Graph`` with added
+        graph-level features and metadata.
+    :type graph_metadata_functions: List[Callable], optional
+    :param get_contacts_config: Config object containing parameters for running ``GetContacts`` for computing intramolecular
+        contact-based edges. Defaults to None.
+    :type get_contacts_config: GetContactsConfig, optional
+    :param dssp_config: Config Object containing reference to ``DSSP`` executable. Defaults to None.
+        **NB** DSSP must be installed. See installation instructions: https://graphein.ai/getting_started/installation.html#optional-dependencies
+    :type dssp_config: DSSPConfig, optional
+    """
+
+    granularity: Union[GraphAtoms, GranularityOpts] = "CA"
+    keep_hets: bool = False
+    insertions: bool = False
+    pdb_dir: Path = Path(
+        "../examples/pdbs/"
+    )  # Also suggest to avoid hard-coding paths if possible!
+    verbose: bool = False
+    exclude_waters: bool = True
+    deprotonate: bool = False
+    add_hs: bool = False
+
+    # Graph construction functions for protein
+    protein_df_processing_functions: Optional[List[Callable]] = None
+    protein_edge_construction_functions: List[Union[Callable, str]] = [
+        add_peptide_bonds,
+        add_atomic_edges,
+    ]
+    protein_node_metadata_functions: Optional[List[Union[Callable, str]]] = [
+        meiler_embedding
+    ]
+    protein_edge_metadata_functions: Optional[
+        List[Union[Callable, str]]
+    ] = None
+    graph_metadata_functions: Optional[List[Callable]] = None
+
+    # Graph construction functions for ligand
+    ligand_edge_construction_functions: List[Union[Callable, str]] = [
+        add_atom_bonds,
+        add_fully_connected_edges_ligand,
+        add_k_nn_edges_ligand,
+        add_distance_threshold_ligand,
+    ]
+    ligand_node_metadata_functions: Optional[List[Union[Callable, str]]] = [
+        atom_type_one_hot
+    ]
+    ligand_edge_metadata_functions: Optional[List[Union[Callable, str]]] = None
+
+    # Graph construction functions for protein-ligand
+    protein_ligand_edge_construction_functions: List[Union[Callable, str]] = [
+        add_fully_connected_edges_protein_ligand,
+        add_k_nn_edges_protein_ligand,
+        add_distance_threshold_protein_ligand,
+    ]
+    protein_ligand_edge_metadata_functions: Optional[
+        List[Union[Callable, str]]
+    ] = None
+
+    # External Dependency configs
+    get_contacts_config: Optional[GetContactsConfig] = None
+    dssp_config: Optional[DSSPConfig] = None
+
+    def __eq__(self, other: Any) -> bool:
+        """Overwrites the BaseModel __eq__ function in order to check more specific cases (like partial functions)."""
+        if isinstance(other, ProteinLigandGraphConfig):
+            return (
+                DeepDiff(
+                    self,
+                    other,
+                    custom_operators=[
+                        PartialMatchOperator(types=[partial]),
+                        PathMatchOperator(types=[Path]),
+                    ],
+                )
+                == {}
+            )
+        return self.dict() == other

graphein/protein_ligand/edges/__init__.py

@@ -0,0 +1,3 @@
+from .atomic import *
+from .distance import *
+from .intramolecular import *