def __getitem__(self, index: int):
if not 0 <= index < self._num_examples:
raise IndexError(index)
# Read biopython structure
file_path = self._file_list[index]
structure = fo.read_sdf(str(file_path),
sanitize=False,
add_hs=False,
remove_hs=False)
# assemble the item (no bonds)
item = {
'atoms': fo.bp_to_df(structure),
'id': structure.id,
'file_path': str(file_path),
}
# Add bonds if included
if self._read_bonds:
mol = fo.read_sdf_to_mol(str(file_path),
sanitize=False,
add_hs=False,
remove_hs=False)
bonds_df = fo.get_bonds_list_from_mol(mol[0])
item['bonds'] = bonds_df
if self._transform:
item = self._transform(item)
return item
def _get_rdkit_data(smiles):
mol = Chem.MolFromSmiles(smiles)
mol = Chem.AddHs(mol)
bonds_df = fo.get_bonds_list_from_mol(mol)
type_mapping = {'H': 0, 'C': 1, 'N': 2, 'O': 3, 'F': 4}
type_idx = []
atomic_number = []
acceptor = []
donor = []
aromatic = []
sp = []
sp2 = []
sp3 = []
num_hs = []
for atom in mol.GetAtoms():
type_idx.append(type_mapping[atom.GetSymbol()])
atomic_number.append(atom.GetAtomicNum())
donor.append(0)
acceptor.append(0)
aromatic.append(1 if atom.GetIsAromatic() else 0)
hybridization = atom.GetHybridization()
sp.append(1 if hybridization == HybridizationType.SP else 0)
sp2.append(1 if hybridization == HybridizationType.SP2 else 0)
sp3.append(1 if hybridization == HybridizationType.SP3 else 0)
num_hs.append(atom.GetTotalNumHs(includeNeighbors=True))
feats = factory.GetFeaturesForMol(mol)
for j in range(0, len(feats)):
if feats[j].GetFamily() == 'Donor':
node_list = feats[j].GetAtomIds()
for k in node_list:
donor[k] = 1
elif feats[j].GetFamily() == 'Acceptor':
node_list = feats[j].GetAtomIds()
for k in node_list:
acceptor[k] = 1
atom_feats = [
atomic_number, acceptor, donor, aromatic, sp, sp2, sp3, num_hs
]
return bonds_df, atom_feats