torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/utils/graph/features.py

@@ -0,0 +1,178 @@
+# allowable multiple choice node and edge features 
+allowable_features = {
+    # atom types: 1-118, 119 is masked atom, 120 is misc (e.g. * for polymers)
+    # index: 0-117, 118, 119
+    'possible_atomic_num_list' : list(range(1, 120)) + ['misc'],
+    'possible_chirality_list' : [
+        'CHI_UNSPECIFIED',
+        'CHI_TETRAHEDRAL_CW',
+        'CHI_TETRAHEDRAL_CCW',
+        'CHI_OTHER',
+        'misc'
+    ],
+    'possible_degree_list' : [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 'misc'],
+    'possible_formal_charge_list' : [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 'misc'],
+    'possible_numH_list' : [0, 1, 2, 3, 4, 5, 6, 7, 8, 'misc'],
+    'possible_number_radical_e_list': [0, 1, 2, 3, 4, 'misc'],
+    'possible_hybridization_list' : [
+        'SP', 'SP2', 'SP3', 'SP3D', 'SP3D2', 'misc'
+        ],
+    'possible_is_aromatic_list': [False, True],
+    'possible_is_in_ring_list': [False, True],
+    'possible_bond_type_list' : [
+        'SINGLE',
+        'DOUBLE',
+        'TRIPLE',
+        'AROMATIC',
+        'misc'
+    ],
+    'possible_bond_stereo_list': [
+        'STEREONONE',
+        'STEREOZ',
+        'STEREOE',
+        'STEREOCIS',
+        'STEREOTRANS',
+        'STEREOANY',
+    ], 
+    'possible_is_conjugated_list': [False, True],
+}
+
+def safe_index(l, e):
+    """
+    Return index of element e in list l. If e is not present, return the last index
+    """
+    try:
+        return l.index(e)
+    except:
+        return len(l) - 1
+# # miscellaneous case
+# i = safe_index(allowable_features['possible_atomic_num_list'], 'asdf')
+# assert allowable_features['possible_atomic_num_list'][i] == 'misc'
+# # normal case
+# i = safe_index(allowable_features['possible_atomic_num_list'], 2)
+# assert allowable_features['possible_atomic_num_list'][i] == 2
+
+def atom_to_feature_vector(atom):
+    """
+    Converts rdkit atom object to feature list of indices
+    :param mol: rdkit atom object
+    :return: list
+    """
+    atom_feature = [
+            safe_index(allowable_features['possible_atomic_num_list'], atom.GetAtomicNum()),
+            safe_index(allowable_features['possible_chirality_list'], str(atom.GetChiralTag())),
+            safe_index(allowable_features['possible_degree_list'], atom.GetTotalDegree()),
+            safe_index(allowable_features['possible_formal_charge_list'], atom.GetFormalCharge()),
+            safe_index(allowable_features['possible_numH_list'], atom.GetTotalNumHs()),
+            safe_index(allowable_features['possible_number_radical_e_list'], atom.GetNumRadicalElectrons()),
+            safe_index(allowable_features['possible_hybridization_list'], str(atom.GetHybridization())),
+            allowable_features['possible_is_aromatic_list'].index(atom.GetIsAromatic()),
+            allowable_features['possible_is_in_ring_list'].index(atom.IsInRing()),
+            ]
+    return atom_feature
+# from rdkit import Chem
+# mol = Chem.MolFromSmiles('Cl[C@H](/C=C/C)Br')
+# atom = mol.GetAtomWithIdx(1)  # chiral carbon
+# atom_feature = atom_to_feature_vector(atom)
+# assert atom_feature == [5, 2, 4, 5, 1, 0, 2, 0, 0]
+
+
+def get_atom_feature_dims():
+    return list(map(len, [
+        allowable_features['possible_atomic_num_list'],
+        allowable_features['possible_chirality_list'],
+        allowable_features['possible_degree_list'],
+        allowable_features['possible_formal_charge_list'],
+        allowable_features['possible_numH_list'],
+        allowable_features['possible_number_radical_e_list'],
+        allowable_features['possible_hybridization_list'],
+        allowable_features['possible_is_aromatic_list'],
+        allowable_features['possible_is_in_ring_list']
+        ]))
+
+def bond_to_feature_vector(bond):
+    """
+    Converts rdkit bond object to feature list of indices
+    :param mol: rdkit bond object
+    :return: list
+    """
+    bond_feature = [
+                safe_index(allowable_features['possible_bond_type_list'], str(bond.GetBondType())),
+                allowable_features['possible_bond_stereo_list'].index(str(bond.GetStereo())),
+                allowable_features['possible_is_conjugated_list'].index(bond.GetIsConjugated()),
+            ]
+    return bond_feature
+# uses same molecule as atom_to_feature_vector test
+# bond = mol.GetBondWithIdx(2)  # double bond with stereochem
+# bond_feature = bond_to_feature_vector(bond)
+# assert bond_feature == [1, 2, 0]
+
+def get_bond_feature_dims():
+    return list(map(len, [
+        allowable_features['possible_bond_type_list'],
+        allowable_features['possible_bond_stereo_list'],
+        allowable_features['possible_is_conjugated_list']
+        ]))
+
+def atom_feature_vector_to_dict(atom_feature):
+    [atomic_num_idx, 
+    chirality_idx,
+    degree_idx,
+    formal_charge_idx,
+    num_h_idx,
+    number_radical_e_idx,
+    hybridization_idx,
+    is_aromatic_idx,
+    is_in_ring_idx] = atom_feature
+
+    feature_dict = {
+        'atomic_num': allowable_features['possible_atomic_num_list'][atomic_num_idx],
+        'chirality': allowable_features['possible_chirality_list'][chirality_idx],
+        'degree': allowable_features['possible_degree_list'][degree_idx],
+        'formal_charge': allowable_features['possible_formal_charge_list'][formal_charge_idx],
+        'num_h': allowable_features['possible_numH_list'][num_h_idx],
+        'num_rad_e': allowable_features['possible_number_radical_e_list'][number_radical_e_idx],
+        'hybridization': allowable_features['possible_hybridization_list'][hybridization_idx],
+        'is_aromatic': allowable_features['possible_is_aromatic_list'][is_aromatic_idx],
+        'is_in_ring': allowable_features['possible_is_in_ring_list'][is_in_ring_idx]
+    }
+
+    return feature_dict
+# # uses same atom_feature as atom_to_feature_vector test
+# atom_feature_dict = atom_feature_vector_to_dict(atom_feature)
+# assert atom_feature_dict['atomic_num'] == 6
+# assert atom_feature_dict['chirality'] == 'CHI_TETRAHEDRAL_CCW'
+# assert atom_feature_dict['degree'] == 4
+# assert atom_feature_dict['formal_charge'] == 0
+# assert atom_feature_dict['num_h'] == 1
+# assert atom_feature_dict['num_rad_e'] == 0
+# assert atom_feature_dict['hybridization'] == 'SP3'
+# assert atom_feature_dict['is_aromatic'] == False
+# assert atom_feature_dict['is_in_ring'] == False
+
+def bond_feature_vector_to_dict(bond_feature):
+    [bond_type_idx, 
+    bond_stereo_idx,
+    is_conjugated_idx] = bond_feature
+
+    feature_dict = {
+        'bond_type': allowable_features['possible_bond_type_list'][bond_type_idx],
+        'bond_stereo': allowable_features['possible_bond_stereo_list'][bond_stereo_idx],
+        'is_conjugated': allowable_features['possible_is_conjugated_list'][is_conjugated_idx]
+    }
+
+    return feature_dict
+# # uses same bond as bond_to_feature_vector test
+# bond_feature_dict = bond_feature_vector_to_dict(bond_feature)
+# assert bond_feature_dict['bond_type'] == 'DOUBLE'
+# assert bond_feature_dict['bond_stereo'] == 'STEREOE'
+# assert bond_feature_dict['is_conjugated'] == False
+
+from rdkit.Chem import AllChem
+
+def getmorganfingerprint(mol):
+    return list(AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=1024))
+
+def getmaccsfingerprint(mol):
+    fp = AllChem.GetMACCSKeysFingerprint(mol)
+    return [int(b) for b in fp.ToBitString()]

torch_molecule/utils/graph/graph_from_smiles.py

@@ -0,0 +1,134 @@
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem import Crippen
+from .features import atom_to_feature_vector, bond_to_feature_vector
+from .features import getmaccsfingerprint, getmorganfingerprint
+from ..generic.pseudo_tasks import PSEUDOTASK
+
+def add_fingerprint_feature(mol, feature_type, get_fingerprint_fn):
+    if feature_type is None:
+        return None
+    fingerprint = get_fingerprint_fn(mol)
+    return np.expand_dims(np.array(fingerprint, dtype="int8"), axis=0)
+    
+def get_augmented_property(mol, properties):
+    if mol is None:
+        return None
+            
+    supported_properties = set(PSEUDOTASK.keys())
+    unsupported = set(properties) - supported_properties
+    if unsupported:
+        raise ValueError(f"Unsupported properties: {unsupported}. Supported properties are: {supported_properties}")
+    
+    augmented_property = []
+    if 'maccs' in properties:
+        maccs = getmaccsfingerprint(mol)
+        augmented_property.extend(maccs)
+    if 'morgan' in properties:
+        mgf = getmorganfingerprint(mol)
+        augmented_property.extend(mgf)
+    if 'logP' in properties:
+        logp = Crippen.MolLogP(mol)
+        augmented_property.append(logp)
+    return augmented_property
+
+def graph_from_smiles(smiles_or_mol, properties, augmented_features=None, augmented_properties=None):
+    """
+    Converts SMILES string or RDKit molecule to graph Data object
+    
+    Parameters
+    ----------
+    smiles_or_mol : Union[str, rdkit.Chem.rdchem.Mol]
+        SMILES string or RDKit molecule object
+    properties : Any
+        Properties to include in the graph
+    augmented_features : list
+        List of augmented features to include
+    augmented_properties : list, optional
+        List of augmented properties to include
+        
+    Returns
+    -------
+    dict
+        Graph object dictionary
+    """
+    # try:
+    if isinstance(smiles_or_mol, str):
+        mol = Chem.MolFromSmiles(smiles_or_mol)
+    else:
+        mol = smiles_or_mol
+        
+    # atoms
+    atom_features_list = []
+    for atom in mol.GetAtoms():
+        # print(atom.GetSymbol(), atom_to_feature_vector(atom)[0])
+        atom_features_list.append(atom_to_feature_vector(atom))
+
+    x = np.array(atom_features_list, dtype=np.int64)
+
+    # bonds
+    num_bond_features = 3  # bond type, bond stereo, is_conjugated
+    if len(mol.GetBonds()) > 0:  # mol has bonds
+        edges_list = []
+        edge_features_list = []
+        for bond in mol.GetBonds():
+            i = bond.GetBeginAtomIdx()
+            j = bond.GetEndAtomIdx()
+            edge_feature = bond_to_feature_vector(bond)
+            # add edges in both directions
+            edges_list.append((i, j))
+            edge_features_list.append(edge_feature)
+            edges_list.append((j, i))
+            edge_features_list.append(edge_feature)
+
+        # data.edge_index: Graph connectivity in COO format with shape [2, num_edges]
+        edge_index = np.array(edges_list, dtype=np.int64).T
+
+        # data.edge_attr: Edge feature matrix with shape [num_edges, num_edge_features]
+        edge_attr = np.array(edge_features_list, dtype=np.int64)
+
+    else:  # mol has no bonds
+        edge_index = np.empty((2, 0), dtype=np.int64)
+        edge_attr = np.empty((0, num_bond_features), dtype=np.int64)
+
+    graph = dict()
+    graph["edge_index"] = edge_index
+    graph["edge_feat"] = edge_attr
+    graph["node_feat"] = x
+    graph["num_nodes"] = len(x)
+
+    # Handle properties and augmented properties
+    props_list = []            
+    if properties is not None:
+        props_list.append(np.array(properties, dtype=np.float32))
+    if augmented_properties is not None:
+        aug_props = get_augmented_property(mol, augmented_properties)
+        if aug_props:
+            props_list.append(np.array(aug_props, dtype=np.float32))
+    if props_list:
+        combined_props = np.concatenate(props_list)
+        graph['y'] = combined_props.reshape(1, -1)
+    else:
+        graph['y'] = np.full((1, 1), np.nan, dtype=np.float32)
+
+    # Handle augmented features
+    if augmented_features is not None:
+        graph['morgan'] = add_fingerprint_feature(
+            mol,
+            'morgan' if 'morgan' in augmented_features else None,
+            getmorganfingerprint
+        )
+        graph['maccs'] = add_fingerprint_feature(
+            mol,
+            'maccs' if 'maccs' in augmented_features else None,
+            getmaccsfingerprint
+        )
+    else:
+        graph['morgan'] = None
+        graph['maccs'] = None
+
+    return graph    
+    
+    # except Exception as e:
+    #     print(f"Error: {e} during converting {smiles_or_mol} to graph")
+    #     return None

torch_molecule/utils/graph/graph_to_smiles.py

@@ -0,0 +1,312 @@
+from rdkit import Chem, RDLogger
+
+RDLogger.DisableLog("rdApp.*")
+
+import re
+import random
+import logging
+from rdkit import Chem
+from typing import List, Tuple, Optional
+random.seed(0)
+import torch
+
+bond_dict = [
+    None,
+    Chem.rdchem.BondType.SINGLE,
+    Chem.rdchem.BondType.DOUBLE,
+    Chem.rdchem.BondType.TRIPLE,
+    Chem.rdchem.BondType.AROMATIC,
+]
+
+ATOM_VALENCY = {6: 4, 7: 3, 8: 2, 9: 1, 15: 3, 16: 2, 17: 1, 35: 1, 53: 1}
+
+def graph_to_smiles(molecule_list: List[Tuple], atom_decoder: list) -> List[Optional[str]]:
+    smiles_list = []
+    for index, graph in enumerate(molecule_list):
+        try:
+            atom_types, edge_types = graph
+            mol_init = build_molecule_with_partial_charges(atom_types, edge_types, atom_decoder)
+            
+            # Try to correct the molecule with connection=True, then False if needed
+            for connection in (True, False):
+                mol_conn, _ = correct_mol(mol_init, connection=connection)
+                if mol_conn is not None:
+                    break
+            else:
+                mol_conn = mol_init  # Fallback to initial molecule
+
+            # Convert to SMILES
+            smiles = mol2smiles(mol_conn)
+            if not smiles:
+                smiles = Chem.MolToSmiles(mol_conn)
+
+            if smiles:
+                mol = get_mol(smiles)
+                if mol is not None:
+                    # Get the largest fragment
+                    mol_frags = Chem.rdmolops.GetMolFrags(mol, asMols=True, sanitizeFrags=False)
+                    largest_mol = max(mol_frags, key=lambda m: m.GetNumAtoms())
+                    
+                    largest_smiles = mol2smiles(largest_mol)
+                    if largest_smiles and len(largest_smiles) > 1:
+                        smiles_list.append(largest_smiles)
+                    else:
+                        smiles_list.append(None)
+                else:
+                    smiles_list.append(None)
+            else:
+                smiles_list.append(None)
+
+        except Exception as e:
+            try:
+                # Fallback to RDKit's MolToSmiles if everything else fails
+                fallback_smiles = Chem.MolToSmiles(mol_init)
+                if fallback_smiles:
+                    smiles_list.append(fallback_smiles)
+                else:
+                    smiles_list.append(None)
+            except Exception as e2:
+                smiles_list.append(None)
+
+    return smiles_list
+
+def build_molecule_with_partial_charges(atom_types, edge_types, atom_decoder, verbose=False):
+    if verbose:
+        print("\nbuilding new molecule")
+
+    mol = Chem.RWMol()
+    for atom in atom_types:
+        a = Chem.Atom(atom_decoder[atom.item()])
+        mol.AddAtom(a)
+        if verbose:
+            print("Atom added: ", atom.item(), atom_decoder[atom.item()])
+
+    edge_types = torch.triu(edge_types)
+    all_bonds = torch.nonzero(edge_types)
+
+    for i, bond in enumerate(all_bonds):
+        if bond[0].item() != bond[1].item():
+            mol.AddBond(
+                bond[0].item(),
+                bond[1].item(),
+                bond_dict[edge_types[bond[0], bond[1]].item()],
+            )
+            if verbose:
+                print(
+                    "bond added:",
+                    bond[0].item(),
+                    bond[1].item(),
+                    edge_types[bond[0], bond[1]].item(),
+                    bond_dict[edge_types[bond[0], bond[1]].item()],
+                )
+            # add formal charge to atom: e.g. [O+], [N+], [S+]
+            # not support [O-], [N-], [S-], [NH+] etc.
+            flag, atomid_valence = check_valency(mol)
+            if verbose:
+                print("flag, valence", flag, atomid_valence)
+            if flag:
+                continue
+            else:
+                if len(atomid_valence) == 2:
+                    idx = atomid_valence[0]
+                    v = atomid_valence[1]
+                    an = mol.GetAtomWithIdx(idx).GetAtomicNum()
+                    if verbose:
+                        print("atomic num of atom with a large valence", an)
+                    if an in (7, 8, 16) and (v - ATOM_VALENCY[an]) == 1:
+                        mol.GetAtomWithIdx(idx).SetFormalCharge(1)
+                        # print("Formal charge added")
+                else:
+                    continue
+    return mol
+
+
+def correct_mol(mol, connection=False):
+    #####
+    no_correct = False
+    flag, _ = check_valency(mol)
+    if flag:
+        no_correct = True
+
+    while True:
+        if connection:
+            mol_conn = connect_fragments(mol)
+            mol = mol_conn
+            if mol is None:
+                return None, no_correct
+        flag, atomid_valence = check_valency(mol)
+        if flag:
+            break
+        else:
+            try:
+                assert len(atomid_valence) == 2
+                idx = atomid_valence[0]
+                v = atomid_valence[1]
+                queue = []
+                check_idx = 0
+                for b in mol.GetAtomWithIdx(idx).GetBonds():
+                    type = int(b.GetBondType())
+                    queue.append(
+                        (b.GetIdx(), type, b.GetBeginAtomIdx(), b.GetEndAtomIdx())
+                    )
+                    if type == 12:
+                        check_idx += 1
+                queue.sort(key=lambda tup: tup[1], reverse=True)
+
+                if queue[-1][1] == 12:
+                    return None, no_correct
+                elif len(queue) > 0:
+                    start = queue[check_idx][2]
+                    end = queue[check_idx][3]
+                    t = queue[check_idx][1] - 1
+                    mol.RemoveBond(start, end)
+                    if t >= 1:
+                        mol.AddBond(start, end, bond_dict[t])
+            except Exception as e:
+                # print(f"An error occurred in correction: {e}")
+                return None, no_correct
+    return mol, no_correct
+
+
+def check_valid(smiles):
+    mol = get_mol(smiles)
+    if mol is None:
+        return False
+    smiles = mol2smiles(mol)
+    if smiles is None:
+        return False
+    return True
+
+
+def get_mol(smiles_or_mol):
+    """
+    Loads SMILES/molecule into RDKit's object
+    """
+    if isinstance(smiles_or_mol, str):
+        if len(smiles_or_mol) == 0:
+            return None
+        mol = Chem.MolFromSmiles(smiles_or_mol)
+        if mol is None:
+            return None
+        try:
+            Chem.SanitizeMol(mol)
+        except ValueError:
+            return None
+        return mol
+    return smiles_or_mol
+
+
+def mol2smiles(mol):
+    if mol is None:
+        return None
+    try:
+        Chem.SanitizeMol(mol)
+    except ValueError:
+        return None
+    return Chem.MolToSmiles(mol)
+
+
+def check_valency(mol):
+    try:
+        # First attempt to sanitize with specific properties
+        Chem.SanitizeMol(mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_PROPERTIES)
+        return True, None
+    except ValueError as e:
+        e = str(e)
+        p = e.find("#")
+        e_sub = e[p:]
+        atomid_valence = list(map(int, re.findall(r"\d+", e_sub)))
+        return False, atomid_valence
+    except Exception as e:
+        # print(f"An unexpected error occurred: {e}")
+        return False, []
+
+
+##### connect fragements
+def select_atom_with_available_valency(frag):
+    atoms = list(frag.GetAtoms())
+    random.shuffle(atoms)
+    for atom in atoms:
+        if atom.GetAtomicNum() > 1 and atom.GetImplicitValence() > 0:
+            return atom
+    return None
+
+
+def select_atoms_with_available_valency(frag):
+    return [
+        atom
+        for atom in frag.GetAtoms()
+        if atom.GetAtomicNum() > 1 and atom.GetImplicitValence() > 0
+    ]
+
+
+def try_to_connect_fragments(combined_mol, frag, atom1, atom2):
+    # Make copies of the molecules to try the connection
+    trial_combined_mol = Chem.RWMol(combined_mol)
+    trial_frag = Chem.RWMol(frag)
+
+    # Add the new fragment to the combined molecule with new indices
+    new_indices = {
+        atom.GetIdx(): trial_combined_mol.AddAtom(atom)
+        for atom in trial_frag.GetAtoms()
+    }
+
+    # Add the bond between the suitable atoms from each fragment
+    trial_combined_mol.AddBond(
+        atom1.GetIdx(), new_indices[atom2.GetIdx()], Chem.BondType.SINGLE
+    )
+
+    # Adjust the hydrogen count of the connected atoms
+    for atom_idx in [atom1.GetIdx(), new_indices[atom2.GetIdx()]]:
+        atom = trial_combined_mol.GetAtomWithIdx(atom_idx)
+        num_h = atom.GetTotalNumHs()
+        atom.SetNumExplicitHs(max(0, num_h - 1))
+
+    # Add bonds for the new fragment
+    for bond in trial_frag.GetBonds():
+        trial_combined_mol.AddBond(
+            new_indices[bond.GetBeginAtomIdx()],
+            new_indices[bond.GetEndAtomIdx()],
+            bond.GetBondType(),
+        )
+
+    # Convert to a Mol object and try to sanitize it
+    new_mol = Chem.Mol(trial_combined_mol)
+    try:
+        Chem.SanitizeMol(new_mol)
+        return new_mol  # Return the new valid molecule
+    except Chem.MolSanitizeException:
+        return None  # If the molecule is not valid, return None
+
+
+def connect_fragments(mol):
+    # Get the separate fragments
+    frags = Chem.GetMolFrags(mol, asMols=True, sanitizeFrags=False)
+    if len(frags) < 2:
+        return mol
+
+    combined_mol = Chem.RWMol(frags[0])
+
+    for frag in frags[1:]:
+        # Select all atoms with available valency from both molecules
+        atoms1 = select_atoms_with_available_valency(combined_mol)
+        atoms2 = select_atoms_with_available_valency(frag)
+
+        # Try to connect using all combinations of available valency atoms
+        for atom1 in atoms1:
+            for atom2 in atoms2:
+                new_mol = try_to_connect_fragments(combined_mol, frag, atom1, atom2)
+                if new_mol is not None:
+                    # If a valid connection is made, update the combined molecule and break
+                    combined_mol = new_mol
+                    break
+            else:
+                # Continue if the inner loop didn't break (no valid connection found for atom1)
+                continue
+            # Break if the inner loop did break (valid connection found)
+            break
+        else:
+            # If no valid connections could be made with any of the atoms, return None
+            return None
+
+    return combined_mol