dmlc--dgl
Mufei Li
b76ac11c90
* Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * Update * add docs * Fix style * Fix lint * Bug fix * Fix test * Update * Update * Update * Update Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>
50 行
1.5 KiB
Python
50 行
1.5 KiB
Python
import numpy as np
|
|
import os
|
|
import shutil
|
|
|
|
from dgl.data.utils import download, _get_dgl_url, extract_archive
|
|
from dgllife.utils.rdkit_utils import get_mol_3d_coordinates, load_molecule
|
|
from rdkit import Chem
|
|
from rdkit.Chem import AllChem
|
|
|
|
def test_get_mol_3D_coordinates():
|
|
mol = Chem.MolFromSmiles('CCO')
|
|
# Test the case when conformation does not exist
|
|
assert get_mol_3d_coordinates(mol) is None
|
|
|
|
# Test the case when conformation exists
|
|
AllChem.EmbedMolecule(mol)
|
|
AllChem.MMFFOptimizeMolecule(mol)
|
|
coords = get_mol_3d_coordinates(mol)
|
|
assert isinstance(coords, np.ndarray)
|
|
assert coords.shape == (mol.GetNumAtoms(), 3)
|
|
|
|
def remove_dir(dir):
|
|
if os.path.isdir(dir):
|
|
try:
|
|
shutil.rmtree(dir)
|
|
except OSError:
|
|
pass
|
|
|
|
def test_load_molecule():
|
|
remove_dir('tmp1')
|
|
remove_dir('tmp2')
|
|
|
|
url = _get_dgl_url('dgllife/example_mols.tar.gz')
|
|
local_path = 'tmp1/example_mols.tar.gz'
|
|
download(url, path=local_path)
|
|
extract_archive(local_path, 'tmp2')
|
|
|
|
load_molecule('tmp2/example_mols/example.sdf')
|
|
load_molecule('tmp2/example_mols/example.mol2', use_conformation=False, sanitize=True)
|
|
load_molecule('tmp2/example_mols/example.pdbqt', calc_charges=True)
|
|
mol, _ = load_molecule('tmp2/example_mols/example.pdb', remove_hs=True)
|
|
assert mol.GetNumAtoms() == mol.GetNumHeavyAtoms()
|
|
|
|
remove_dir('tmp1')
|
|
remove_dir('tmp2')
|
|
|
|
if __name__ == '__main__':
|
|
test_get_mol_3D_coordinates()
|
|
test_load_molecule()
|