Spaces:
Sleeping
Sleeping
File size: 3,043 Bytes
dbbf75a 136190c dbbf75a 09e6d6a dbbf75a 09e6d6a 136190c 09e6d6a 136190c 09e6d6a 136190c 09e6d6a 136190c dbbf75a 09e6d6a dbbf75a 136190c dbbf75a 136190c dbbf75a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import torch
from rdkit import Chem
from rdkit.Chem import AllChem
from torch_geometric.data import Data
ATOM_TYPES = [6, 7, 8, 16, 15, 9, 17, 35, 53]
HYBRIDIZATION_TYPES = [
Chem.rdchem.HybridizationType.SP,
Chem.rdchem.HybridizationType.SP2,
Chem.rdchem.HybridizationType.SP3,
Chem.rdchem.HybridizationType.SP3D,
Chem.rdchem.HybridizationType.SP3D2,
Chem.rdchem.HybridizationType.UNSPECIFIED,
]
BOND_TYPES = [
Chem.rdchem.BondType.SINGLE,
Chem.rdchem.BondType.DOUBLE,
Chem.rdchem.BondType.TRIPLE,
Chem.rdchem.BondType.AROMATIC,
]
STEREO_TYPES = [
Chem.rdchem.BondStereo.STEREONONE,
Chem.rdchem.BondStereo.STEREOANY,
Chem.rdchem.BondStereo.STEREOZ,
Chem.rdchem.BondStereo.STEREOE,
]
def one_hot(val, choices):
return [1 if val == c else 0 for c in choices]
def get_atom_features(atom):
feat = one_hot(atom.GetAtomicNum(), ATOM_TYPES)
feat += one_hot(atom.GetDegree(), list(range(7)))
feat += one_hot(atom.GetTotalValence(), list(range(7)))
feat += one_hot(atom.GetFormalCharge(), list(range(-3, 4)))
feat += one_hot(atom.GetHybridization(), HYBRIDIZATION_TYPES)
feat.append(1.0 if atom.GetIsAromatic() else 0.0)
feat += one_hot(atom.GetTotalNumHs(), list(range(5)))
feat.append(float(atom.GetNumRadicalElectrons()))
feat.append(1.0 if atom.IsInRing() else 0.0)
feat.append(float(atom.GetMass()) / 200.0)
return feat
def get_bond_features(bond, mol):
feat = one_hot(bond.GetBondType(), BOND_TYPES)
feat.append(1.0 if bond.GetIsConjugated() else 0.0)
feat.append(1.0 if bond.IsInRing() else 0.0)
feat += one_hot(bond.GetStereo(), STEREO_TYPES)
conf = mol.GetConformer()
i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
dist = conf.GetAtomPosition(i).Distance(conf.GetAtomPosition(j))
feat.append(min(dist / 2.0, 1.0))
return feat
def smiles_to_graph(smiles: str):
mol = Chem.MolFromSmiles(smiles)
if mol is None:
raise ValueError(f"Invalid SMILES: {smiles}")
mol = Chem.AddHs(mol)
AllChem.EmbedMolecule(mol, randomSeed=42)
AllChem.MMFFOptimizeMolecule(mol)
mol = Chem.RemoveHs(mol)
atom_features = []
for atom in mol.GetAtoms():
atom_features.append(get_atom_features(atom))
x = torch.tensor(atom_features, dtype=torch.float32)
edge_indices = []
edge_features = []
for bond in mol.GetBonds():
i = bond.GetBeginAtomIdx()
j = bond.GetEndAtomIdx()
edge_indices.append([i, j])
edge_indices.append([j, i])
bf = get_bond_features(bond, mol)
edge_features.append(bf)
edge_features.append(bf)
if len(edge_indices) == 0:
edge_index = torch.zeros((2, 0), dtype=torch.long)
edge_attr = torch.zeros((0, 11), dtype=torch.float32)
else:
edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()
edge_attr = torch.tensor(edge_features, dtype=torch.float32)
return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
|