from typing import Tuple, List from rdkit import Chem from rdkit.Chem import AllChem, Descriptors, Draw import networkx as nx import pandas as pd import numpy as np from tqdm import tqdm from protac_splitter.chemoinformatics import get_atom_idx_at_attachment from protac_splitter.display_utils import safe_display, get_mapped_protac_img def bond_capacity(bond: Chem.Bond) -> int: """ Calculate the capacity of a bond based on its type and properties. Parameters: bond (Chem.Bond): The bond object from RDKit. Returns: int: The capacity of the bond, where higher values indicate less preference for cutting. """ # High capacity for aromatic and ring bonds to avoid cutting them if bond.GetIsAromatic() or bond.IsInRing(): return 1000 # very high capacity: avoid cutting aromatic bonds elif bond.GetBondType() == Chem.BondType.SINGLE: return 1 # low capacity: prefer to cut here elif bond.GetBondType() == Chem.BondType.DOUBLE: return 10 # medium penalty elif bond.GetBondType() == Chem.BondType.TRIPLE: return 20 # stronger penalty else: return 50 # fallback for unknown/rare types def smiles_to_nx( smiles: str, use_capacity: bool = False, ) -> nx.Graph: """ Convert a SMILES string to a NetworkX graph. Parameters: smiles (str): The SMILES string to convert. use_capacity (bool): Whether to use bond capacity as edge weights. Returns: nx.Graph: The NetworkX graph representation of the molecule. """ mol = Chem.MolFromSmiles(smiles) if mol is None: raise ValueError(f"Input SMILES could not be parsed: {smiles}") # Canonicalize the SMILES mol = Chem.MolFromSmiles(Chem.MolToSmiles(mol, canonical=True)) if mol is None: raise ValueError(f"Input SMILES could not be canonicalized: {smiles}") # Convert SMILES to NetworkX graph G = nx.Graph() if use_capacity: for bond in mol.GetBonds(): capacity = bond_capacity(bond) G.add_edge(bond.GetBeginAtomIdx(), bond.GetEndAtomIdx(), capacity=capacity) else: for bond in mol.GetBonds(): G.add_edge(bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()) return G def extract_edge_features( protac_smiles: str, e3_split_pair: Tuple[int, int] = None, wh_split_pair: Tuple[int, int] = None, n_bits: int = 512, radius: int = 6, descriptor_names: List[str] = None, fp_as_string: bool = False, ) -> pd.DataFrame: """Extract features from the edges of a PROTAC molecule represented as a SMILES string. Parameters: protac_smiles (str): SMILES representation of the PROTAC molecule. e3_split_pair (Tuple[int, int]): Indices of the E3 split pair. wh_split_pair (Tuple[int, int]): Indices of the warhead split pair. n_bits (int): Number of bits for Morgan fingerprints. radius (int): Radius for Morgan fingerprints. descriptor_names (List[str]): List of RDKit descriptor names to compute. Returns: pd.DataFrame: DataFrame containing edge features. """ mol = Chem.MolFromSmiles(protac_smiles) if mol is None: raise ValueError(f"Input SMILES could not be parsed: {protac_smiles}") # Canonicalize the SMILES mol = Chem.MolFromSmiles(Chem.MolToSmiles(mol, canonical=True)) if mol is None: raise ValueError(f"Input SMILES could not be canonicalized: {protac_smiles}") # Step 1: Convert SMILES to NetworkX G = smiles_to_nx(protac_smiles, use_capacity=False) num_nodes = G.number_of_nodes() num_edges = G.number_of_edges() # Step 2: Create line graph and compute betweenness + degree LG = nx.line_graph(G) line_betweenness = nx.betweenness_centrality(LG, endpoints=True) betweenness = nx.betweenness_centrality(G, endpoints=True) # Compute k-hop degrees (number of nodes within 2, 3 hops) # TODO: Shall I get the degree of the node in the line graph or the original graph? line_degree = dict(LG.degree()) line_degree_r2 = {} line_degree_r3 = {} for node in LG.nodes(): # Nodes within radius 2 and 3 (excluding the center node) neighbors_r2 = nx.single_source_shortest_path_length(LG, node, cutoff=2) neighbors_r3 = nx.single_source_shortest_path_length(LG, node, cutoff=3) line_degree_r2[node] = len([n for n, d in neighbors_r2.items() if d == 2]) line_degree_r3[node] = len([n for n, d in neighbors_r3.items() if d == 3]) degree = dict(G.degree()) degree_r2 = {} degree_r3 = {} for node in G.nodes(): # Nodes within radius 2 and 3 (excluding the center node) neighbors_r2 = nx.single_source_shortest_path_length(G, node, cutoff=2) neighbors_r3 = nx.single_source_shortest_path_length(G, node, cutoff=3) degree_r2[node] = len([n for n, d in neighbors_r2.items() if d == 2]) degree_r3[node] = len([n for n, d in neighbors_r3.items() if d == 3]) if e3_split_pair is not None and wh_split_pair is not None: true_split_edges = {frozenset(e3_split_pair), frozenset(wh_split_pair)} # Get molecular characteristics, i.e., Morgan fingerprints and descriptors # Generate Morgan fingerprint fp_bitvec = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits) fp = np.zeros((n_bits,), dtype=np.float32) AllChem.DataStructs.ConvertToNumpyArray(fp_bitvec, fp) if fp_as_string: fp = {"chem_mol_fp": "".join([str(int(bit)) for bit in fp])} else: fp = {f"chem_mol_fp_{i}": bool(fp[i]) for i in range(n_bits)} # Generate RDKit descriptors descriptor_func_names = descriptor_names or [ "MolWt", "HeavyAtomCount", "NumHAcceptors", "NumHDonors", "TPSA", "NumRotatableBonds", "RingCount", "MolLogP" ] functions = [getattr(Descriptors, name) for name in descriptor_func_names] descriptors = {f"chem_mol_desc_{name}": func(mol) for name, func in zip(descriptor_func_names, functions)} # Step 3: Gather edge features # NOTE: Only consider bridge nodes edge_features = [] for (u, v) in nx.bridges(G): bond = mol.GetBondBetweenAtoms(u, v) # Avoid reporting the same edge twice (i.e., swap u and v if needed) and # ensure to find the node pair in the line graph node = (u, v) if (u, v) in LG else (v, u) node_key = node if node in line_betweenness else (v, u) features = { "graph_num_nodes": num_nodes, "graph_num_edges": num_edges, "graph_betweenness": line_betweenness.get(node_key, 0.0), "graph_degree": line_degree.get(node_key, 0), "graph_degree_r2": line_degree_r2.get(node_key, 0), "graph_degree_r3": line_degree_r3.get(node_key, 0), "graph_node_u_degree": degree.get(u, 0), "graph_node_u_degree_r2": degree_r2.get(u, 0), "graph_node_u_degree_r3": degree_r3.get(u, 0), "graph_node_v_degree": degree.get(v, 0), "graph_node_v_degree_r2": degree_r2.get(v, 0), "graph_node_v_degree_r3": degree_r3.get(v, 0), "graph_node_u_betweenness": betweenness.get(u, 0.0), "graph_node_v_betweenness": betweenness.get(v, 0.0), "chem_bond_idx": bond.GetIdx(), "chem_bond_type": str(bond.GetBondType()), "chem_atom_u": mol.GetAtomWithIdx(u).GetSymbol(), "chem_atom_v": mol.GetAtomWithIdx(v).GetSymbol(), "chem_is_aromatic": bond.GetIsAromatic(), "chem_is_in_ring": bond.IsInRing(), "chem_mol_smiles": protac_smiles, "chem_mol_n_bits": n_bits, "chem_mol_radius": radius, } # Add RDKit descriptors and Morgan fingerprint features.update(fp) features.update(descriptors) # Add E3 and warhead split labels if e3_split_pair is not None and wh_split_pair is not None: features.update({ "label_is_split": frozenset([u, v]) in true_split_edges, "label_e3_split": frozenset([u, v]) == frozenset(e3_split_pair), "label_wh_split": frozenset([u, v]) == frozenset(wh_split_pair), }) # Append the features to the list of edge features edge_features.append(features) df = pd.DataFrame(edge_features) # Identify columns with int64 dtype int64_cols = df.select_dtypes(include=['int64']).columns # Create a dictionary mapping these columns to int32 dtype_mapping = {col: np.int32 for col in int64_cols} # Apply the type conversion df = df.astype(dtype_mapping) return df def get_edge_features( protac_smiles: str | List[str], wh_smiles: str | List[str], lk_smiles: str | List[str], e3_smiles: str | List[str], n_bits: int = 512, radius: int = 6, descriptor_names: List[str] = None, fp_as_string: bool = False, verbose: int = 0, ) -> pd.DataFrame: """Get edge features for a given PROTAC molecule and its components. Parameters: protac_smiles (str | List[str]): SMILES representation of the PROTAC molecule. wh_smiles (str | List[str]): SMILES representation of the warhead. lk_smiles (str | List[str]): SMILES representation of the linker. e3_smiles (str | List[str]): SMILES representation of the E3 binder. n_bits (int): Number of bits for Morgan fingerprints. radius (int): Radius for Morgan fingerprints. descriptor_names (List[str]): List of RDKit descriptor names to compute. Returns: pd.DataFrame: DataFrame containing edge features. """ if isinstance(protac_smiles, str): protac_smiles = [protac_smiles] if isinstance(wh_smiles, str): wh_smiles = [wh_smiles] if isinstance(lk_smiles, str): lk_smiles = [lk_smiles] if isinstance(e3_smiles, str): e3_smiles = [e3_smiles] iterables = zip(protac_smiles, wh_smiles, lk_smiles, e3_smiles) iterables = tqdm(iterables, desc="Extracting edge features", total=len(protac_smiles), disable=verbose == 0) features_list = [] for protac_smi, wh_smi, lk_smi, e3_smi in iterables: if verbose > 1: get_mapped_protac_img(protac_smi, wh_smi, lk_smi, e3_smi, w=1500, h=600, display_image=True, useSVG=True) # Convert SMILES to RDKit molecules protac = Chem.MolFromSmiles(protac_smi) wh = Chem.MolFromSmiles(wh_smi) lk = Chem.MolFromSmiles(lk_smi) e3 = Chem.MolFromSmiles(e3_smi) if protac is None or wh is None or lk is None or e3 is None: raise ValueError(f"Invalid SMILES string: {protac}, {wh}, {lk}, {e3}") # Get the attachment points wh_edge = get_atom_idx_at_attachment(protac, wh, lk) e3_edge = get_atom_idx_at_attachment(protac, e3, lk) # Extract features features = extract_edge_features( protac_smi, e3_split_pair=e3_edge, wh_split_pair=wh_edge, n_bits=n_bits, radius=radius, descriptor_names=descriptor_names, fp_as_string=fp_as_string, ) if verbose > 1: # Randomly sample and display 5 edges sample_edges = features.sample(n=5, random_state=42) # Display the sampled edges for _, row in sample_edges.iterrows(): bond = protac.GetBondWithIdx(row['chem_bond_idx']) u, v = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx() safe_display(Draw.MolToImage( protac, size=(1500, 400), highlightColor=(1, 0, 1, 0.3), # Light purple highlightAtoms=[u, v], # Highlight the two atoms legend=f"Graph nodes: {u}, {v} (Betweenness centrality: {row['graph_betweenness']:.3f})", )) # print(row[[c for c in features.columns if c.startswith('graph_')] + ['chem_atom_u', 'chem_atom_v', 'chem_is_in_ring']]) print(row) # Append the features to the list features_list.append(features) return pd.concat(features_list, ignore_index=True)