def test_combined(self):
ecfp_power = 5
splif_power = 5
# test voxel features
featurizer = rgf.RdkitGridFeaturizer(
voxel_width=1.0,
box_width=20.0,
feature_types=['voxel_combined'],
ecfp_power=ecfp_power,
splif_power=splif_power,
flatten=False,
sanitize=True)
feature_tensor, _ = featurizer.featurize_complexes([self.ligand_file],
[self.protein_file])
self.assertIsInstance(feature_tensor, np.ndarray)
voxel_total_len = (
2**ecfp_power +
len(featurizer.cutoffs['splif_contact_bins']) * 2**splif_power + len(
featurizer.cutoffs['hbond_dist_bins']) + 5)
self.assertEqual(feature_tensor.shape, (1, 20, 20, 20, voxel_total_len))
# test flat features
featurizer = rgf.RdkitGridFeaturizer(
voxel_width=1.0,
feature_types=['flat_combined'],
ecfp_power=ecfp_power,
splif_power=splif_power,
sanitize=True)
feature_tensor, _ = featurizer.featurize_complexes([self.ligand_file],
[self.protein_file])
self.assertIsInstance(feature_tensor, np.ndarray)
flat_total_len = (
3 * 2**ecfp_power +
len(featurizer.cutoffs['splif_contact_bins']) * 2**splif_power + len(
featurizer.cutoffs['hbond_dist_bins']))
self.assertEqual(feature_tensor.shape, (1, flat_total_len))
# check if aromatic features are ignores if sanitize=False
featurizer = rgf.RdkitGridFeaturizer(
voxel_width=16.0,
feature_types=['all_combined'],
ecfp_power=ecfp_power,
splif_power=splif_power,
flatten=True,
sanitize=False)
self.assertTrue('pi_stack' not in featurizer.feature_types)
self.assertTrue('cation_pi' not in featurizer.feature_types)
feature_tensor, _ = featurizer.featurize_complexes([self.ligand_file],
[self.protein_file])
self.assertIsInstance(feature_tensor, np.ndarray)
total_len = voxel_total_len + flat_total_len - 3 - 2**ecfp_power
self.assertEqual(feature_tensor.shape, (1, total_len))
def test_default_featurizer(self):
# test if default parameters work
featurizer = rgf.RdkitGridFeaturizer()
self.assertIsInstance(featurizer, rgf.RdkitGridFeaturizer)
feature_tensor, _ = featurizer.featurize([self.ligand_file],
[self.protein_file])
self.assertIsInstance(feature_tensor, np.ndarray)
def test_rotations(self):
featurizer = rgf.RdkitGridFeaturizer(nb_rotations=3,
feature_types=['voxel_combined'],
flatten=False,
sanitize=True)
feature_tensors, _ = featurizer.featurize([self.ligand_file],
[self.protein_file])
self.assertEqual(feature_tensors.shape, (1, 4, 16, 16, 16, 40))
def test_rotations(self):
featurizer = rgf.RdkitGridFeaturizer(nb_rotations=3,
feature_types=['voxel_combined'],
flatten=False,
sanitize=True)
feature_tensors = featurizer.featurize_complexes([self.ligand_file],
[self.protein_file])
self.assertEqual(len(feature_tensors), 4)
def test_force_flatten(self):
# test if input is flattened when flat features are used
featurizer = rgf.RdkitGridFeaturizer(
feature_types=['ecfp_hashed'], flatten=False)
featurizer.flatten = True # False should be ignored with ecfp_hashed
feature_tensor, _ = featurizer.featurize_complexes([self.ligand_file],
[self.protein_file])
self.assertIsInstance(feature_tensor, np.ndarray)
self.assertEqual(feature_tensor.shape, (1, 2 * 2**featurizer.ecfp_power))
def test_example_featurizer(self):
# check if use-case from examples works
featurizer = rgf.RdkitGridFeaturizer(
voxel_width=16.0,
feature_types=['ecfp', 'splif', 'hbond', 'salt_bridge'],
ecfp_power=9,
splif_power=9,
flatten=True)
feature_tensor, _ = featurizer.featurize([self.ligand_file],
[self.protein_file])
self.assertIsInstance(feature_tensor, np.ndarray)
def test_custom_cutoffs(self):
custom_cutoffs = {
'hbond_dist_bins': [(2., 3.), (3., 3.5)],
'hbond_angle_cutoffs': [5, 90],
'splif_contact_bins': [(0, 3.5), (3.5, 6.0)],
'ecfp_cutoff': 5.0,
'sybyl_cutoff': 3.0,
'salt_bridges_cutoff': 4.0,
'pi_stack_dist_cutoff': 5.0,
'pi_stack_angle_cutoff': 15.0,
'cation_pi_dist_cutoff': 5.5,
'cation_pi_angle_cutoff': 20.0,
}
rgf_featurizer = rgf.RdkitGridFeaturizer(**custom_cutoffs)
self.assertEqual(rgf_featurizer.cutoffs, custom_cutoffs)
def test_voxelize(self):
prot_xyz, prot_rdk = rgf.load_molecule(self.protein_file)
lig_xyz, lig_rdk = rgf.load_molecule(self.ligand_file)
centroid = rgf.compute_centroid(lig_xyz)
prot_xyz = rgf.subtract_centroid(prot_xyz, centroid)
lig_xyz = rgf.subtract_centroid(lig_xyz, centroid)
prot_ecfp_dict, lig_ecfp_dict = rgf.featurize_binding_pocket_ecfp(
prot_xyz, prot_rdk, lig_xyz, lig_rdk)
box_w = 20
f_power = 5
rgf_featurizer = rgf.RdkitGridFeaturizer(
box_width=box_w,
ecfp_power=f_power,
feature_types=['all_combined'],
flatten=True,
sanitize=True)
prot_tensor = rgf_featurizer._voxelize(
rgf.convert_atom_to_voxel,
rgf.hash_ecfp,
prot_xyz,
feature_dict=prot_ecfp_dict,
channel_power=f_power)
self.assertEqual(prot_tensor.shape, tuple([box_w] * 3 + [2**f_power]))
all_features = prot_tensor.sum()
# protein is too big for the box, some features should be missing
self.assertGreater(all_features, 0)
self.assertLess(all_features, prot_rdk.GetNumAtoms())
lig_tensor = rgf_featurizer._voxelize(
rgf.convert_atom_to_voxel,
rgf.hash_ecfp,
lig_xyz,
feature_dict=lig_ecfp_dict,
channel_power=f_power)
self.assertEqual(lig_tensor.shape, tuple([box_w] * 3 + [2**f_power]))
all_features = lig_tensor.sum()
# whole ligand should fit in the box
self.assertEqual(all_features, lig_rdk.GetNumAtoms())
def load_pdbbind(featurizer="grid",
split="random",
subset="core",
reload=True):
"""Load and featurize raw PDBBind dataset."""
pdbbind_tasks = ["-logKd/Ki"]
data_dir = deepchem.utils.get_data_dir()
if reload:
save_dir = os.path.join(data_dir,
"pdbbind/" + featurizer + "/" + str(split))
loaded, all_dataset, transformers = deepchem.utils.save.load_dataset_from_disk(
save_dir)
if loaded:
return pdbbind_tasks, all_dataset, transformers
dataset_file = os.path.join(data_dir, "pdbbind_v2015.tar.gz")
data_folder = os.path.join(data_dir, "v2015")
if not os.path.exists(dataset_file):
logger.warning(
"About to download PDBBind full dataset. Large file, 2GB")
deepchem.utils.download_url(
'http://deepchem.io.s3-website-us-west-1.amazonaws.com/datasets/' +
"pdbbind_v2015.tar.gz")
if os.path.exists(data_folder):
logger.info("Data directory for %s already exists" % subset)
else:
print("Untarring full dataset")
deepchem.utils.untargz_file(dataset_file, dest_dir=data_dir)
if subset == "core":
index_file = os.path.join(data_folder, "INDEX_core_name.2013")
labels_file = os.path.join(data_folder, "INDEX_core_data.2013")
elif subset == "refined":
index_file = os.path.join(data_folder, "INDEX_refined_name.2013")
labels_file = os.path.join(data_folder, "INDEX_refined_data.2013")
else:
raise ValueError("Other subsets not supported")
# Extract locations of data
pdbs = []
with open(index_file, "r") as g:
lines = g.readlines()
for line in lines:
line = line.split(" ")
pdb = line[0]
if len(pdb) == 4:
pdbs.append(pdb)
protein_files = [
os.path.join(data_folder, pdb, "%s_protein.pdb" % pdb) for pdb in pdbs
]
ligand_files = [
os.path.join(data_folder, pdb, "%s_ligand.sdf" % pdb) for pdb in pdbs
]
# Extract labels
labels = []
with open(labels_file, "r") as f:
lines = f.readlines()
for line in lines:
# Skip comment lines
if line[0] == "#":
continue
# Lines have format
# PDB code, resolution, release year, -logKd/Ki, Kd/Ki, reference, ligand name
line = line.split()
# The base-10 logarithm, -log kd/pk
log_label = line[3]
labels.append(log_label)
# Featurize Data
if featurizer == "grid":
# TODO: This is not the correct setting. Set hyperparameters correctly
ecfp_power = 5
splif_power = 5
featurizer = rgf.RdkitGridFeaturizer(
voxel_width=16.0,
feature_types=['ecfp', 'splif', 'hbond', 'salt_bridge'],
ecfp_power=ecfp_power,
splif_power=splif_power,
flatten=True)
else:
raise ValueError("Featurizer not supported")
print("Featurizing Complexes")
features = featurizer.featurize_complexes(ligand_files,
protein_files,
log_every_n=1)
dataset = deepchem.data.DiskDataset.from_numpy(features, labels)
# No transformations of data
transformers = []
# TODO(rbharath): This should be modified to contain a cluster split so
# structures of the same protein aren't in both train/test
splitters = {
'index': deepchem.splits.IndexSplitter(),
'random': deepchem.splits.RandomSplitter(),
}
splitter = splitters[split]
train, valid, test = splitter.train_valid_test_split(dataset)
all_dataset = (train, valid, test)
if reload:
deepchem.utils.save.save_dataset_to_disk(save_dir, train, valid, test,
transformers)
return pdbbind_tasks, all_dataset, transformers
def load_pdbbind(featurizer="grid",
split="random",
subset="core",
reload=True):
"""Load and featurize raw PDBBind dataset.
Parameters
----------
data_dir: String, optional
Specifies the data directory to store the featurized dataset.
split: Str
Either "random" or "index"
feat: Str
Either "grid" or "atomic" for grid and atomic featurizations.
subset: Str
Only "core" or "refined" for now.
"""
pdbbind_tasks = ["-logKd/Ki"]
data_dir = deepchem.utils.get_data_dir()
if reload:
save_dir = os.path.join(data_dir,
"pdbbind/" + featurizer + "/" + str(split))
loaded, all_dataset, transformers = deepchem.utils.save.load_dataset_from_disk(
save_dir)
if loaded:
return pdbbind_tasks, all_dataset, transformers
dataset_file = os.path.join(data_dir, "pdbbind_v2015.tar.gz")
data_folder = os.path.join(data_dir, "v2015")
if not os.path.exists(dataset_file):
logger.warning(
"About to download PDBBind full dataset. Large file, 2GB")
deepchem.utils.download_url(
'http://deepchem.io.s3-website-us-west-1.amazonaws.com/datasets/' +
"pdbbind_v2015.tar.gz")
if os.path.exists(data_folder):
logger.info("Data directory for %s already exists" % subset)
else:
print("Untarring full dataset")
deepchem.utils.untargz_file(dataset_file, dest_dir=data_dir)
if subset == "core":
index_file = os.path.join(data_folder, "INDEX_core_name.2013")
labels_file = os.path.join(data_folder, "INDEX_core_data.2013")
elif subset == "refined":
index_file = os.path.join(data_folder, "INDEX_refined_name.2013")
labels_file = os.path.join(data_folder, "INDEX_refined_data.2013")
else:
raise ValueError("Other subsets not supported")
# Extract locations of data
pdbs = []
with open(index_file, "r") as g:
lines = g.readlines()
for line in lines:
line = line.split(" ")
pdb = line[0]
if len(pdb) == 4:
pdbs.append(pdb)
protein_files = [
os.path.join(data_folder, pdb, "%s_protein.pdb" % pdb) for pdb in pdbs
]
ligand_files = [
os.path.join(data_folder, pdb, "%s_ligand.sdf" % pdb) for pdb in pdbs
]
# Extract labels
labels = []
with open(labels_file, "r") as f:
lines = f.readlines()
for line in lines:
# Skip comment lines
if line[0] == "#":
continue
# Lines have format
# PDB code, resolution, release year, -logKd/Ki, Kd/Ki, reference, ligand name
line = line.split()
# The base-10 logarithm, -log kd/pk
log_label = line[3]
labels.append(log_label)
labels = np.array(labels)
# Featurize Data
if featurizer == "grid":
featurizer = rgf.RdkitGridFeaturizer(voxel_width=2.0,
feature_types=[
'ecfp', 'splif', 'hbond',
'salt_bridge', 'pi_stack',
'cation_pi', 'charge'
],
flatten=True)
elif featurizer == "atomic":
# Pulled from PDB files. For larger datasets with more PDBs, would use
# max num atoms instead of exact.
frag1_num_atoms = 70 # for ligand atoms
frag2_num_atoms = 24000 # for protein atoms
complex_num_atoms = 24070 # in total
max_num_neighbors = 4
# Cutoff in angstroms
neighbor_cutoff = 4
featurizer = ComplexNeighborListFragmentAtomicCoordinates(
frag1_num_atoms, frag2_num_atoms, complex_num_atoms,
max_num_neighbors, neighbor_cutoff)
elif featurizer == "atomic_conv":
frag1_num_atoms = 70 # for ligand atoms
frag2_num_atoms = 24000 # for protein atoms
complex_num_atoms = 24070 # in total
max_num_neighbors = 4
# Cutoff in angstroms
neighbor_cutoff = 4
featurizer = AtomicConvFeaturizer(labels=labels,
frag1_num_atoms=frag1_num_atoms,
frag2_num_atoms=frag2_num_atoms,
complex_num_atoms=complex_num_atoms,
neighbor_cutoff=neighbor_cutoff,
max_num_neighbors=max_num_neighbors,
batch_size=64)
else:
raise ValueError("Featurizer not supported")
print("Featurizing Complexes")
features, failures = featurizer.featurize_complexes(
ligand_files, protein_files)
# Delete labels for failing elements
labels = np.delete(labels, failures)
dataset = deepchem.data.DiskDataset.from_numpy(features, labels)
print('Featurization complete.')
# No transformations of data
transformers = []
if split == None:
return pdbbind_tasks, (dataset, None, None), transformers
# TODO(rbharath): This should be modified to contain a cluster split so
# structures of the same protein aren't in both train/test
splitters = {
'index': deepchem.splits.IndexSplitter(),
'random': deepchem.splits.RandomSplitter(),
}
splitter = splitters[split]
train, valid, test = splitter.train_valid_test_split(dataset)
all_dataset = (train, valid, test)
if reload:
deepchem.utils.save.save_dataset_to_disk(save_dir, train, valid, test,
transformers)
return pdbbind_tasks, all_dataset, transformers
def load_pdbbind(reload=True,
data_dir=None,
subset="core",
load_binding_pocket=False,
featurizer="grid",
split="random",
split_seed=None,
save_dir=None,
save_timestamp=False):
"""Load raw PDBBind dataset by featurization and split.
Parameters
----------
reload: Bool, optional
Reload saved featurized and splitted dataset or not.
data_dir: Str, optional
Specifies the directory storing the raw dataset.
load_binding_pocket: Bool, optional
Load binding pocket or full protein.
subset: Str
Specifies which subset of PDBBind, only "core" or "refined" for now.
featurizer: Str
Either "grid" or "atomic" for grid and atomic featurizations.
split: Str
Either "random" or "index".
split_seed: Int, optional
Specifies the random seed for splitter.
save_dir: Str, optional
Specifies the directory to store the featurized and splitted dataset when
reload is False. If reload is True, it will load saved dataset inside save_dir.
save_timestamp: Bool, optional
Save featurized and splitted dataset with timestamp or not. Set it as True
when running similar or same jobs simultaneously on multiple compute nodes.
"""
pdbbind_tasks = ["-logKd/Ki"]
deepchem_dir = deepchem.utils.get_data_dir()
if data_dir == None:
data_dir = DEFAULT_DATA_DIR
data_folder = os.path.join(data_dir, "pdbbind", "v2015")
if save_dir == None:
save_dir = os.path.join(DEFAULT_DATA_DIR, "from-pdbbind")
if load_binding_pocket:
save_folder = os.path.join(
save_dir, "protein_pocket-%s-%s-%s" % (subset, featurizer, split))
else:
save_folder = os.path.join(
save_dir, "full_protein-%s-%s-%s" % (subset, featurizer, split))
if save_timestamp:
save_folder = "%s-%s-%s" % (
save_folder, time.strftime("%Y%m%d", time.localtime()),
re.search("\.(.*)", str(time.time())).group(1))
if reload:
if not os.path.exists(save_folder):
print("Dataset does not exist at {}. Reconstructing...".format(
save_folder))
else:
print("\nLoading featurized and splitted dataset from:\n%s\n" %
save_folder)
loaded, all_dataset, transformers = deepchem.utils.save.load_dataset_from_disk(
save_folder)
if loaded:
return pdbbind_tasks, all_dataset, transformers
dataset_file = os.path.join(data_dir, "pdbbind_v2015.tar.gz")
if not os.path.exists(dataset_file):
logger.warning(
"About to download PDBBind full dataset. Large file, 2GB")
deepchem.utils.download_url(
'http://deepchem.io.s3-website-us-west-1.amazonaws.com/datasets/' +
"pdbbind_v2015.tar.gz",
dest_dir=data_dir)
if os.path.exists(data_folder):
logger.info("PDBBind full dataset already exists.")
else:
print("Untarring full dataset...")
deepchem.utils.untargz_file(dataset_file,
dest_dir=os.path.join(data_dir, "pdbbind"))
print("\nRaw dataset:\n%s" % data_folder)
print("\nFeaturized and splitted dataset:\n%s" % save_folder)
if subset == "core":
index_labels_file = os.path.join(data_folder, "INDEX_core_data.2013")
elif subset == "refined":
index_labels_file = os.path.join(data_folder,
"INDEX_refined_data.2015")
else:
raise ValueError("Other subsets not supported")
# Extract locations of data
with open(index_labels_file, "r") as g:
pdbs = [line[:4] for line in g.readlines() if line[0] != "#"]
if load_binding_pocket:
protein_files = [
os.path.join(data_folder, pdb, "%s_pocket.pdb" % pdb)
for pdb in pdbs
]
else:
protein_files = [
os.path.join(data_folder, pdb, "%s_protein.pdb" % pdb)
for pdb in pdbs
]
ligand_files = [
os.path.join(data_folder, pdb, "%s_ligand.sdf" % pdb) for pdb in pdbs
]
# Extract labels
with open(index_labels_file, "r") as g:
labels = np.array([
# Lines have format
# PDB code, resolution, release year, -logKd/Ki, Kd/Ki, reference, ligand name
# The base-10 logarithm, -log kd/pk
float(line.split()[3]) for line in g.readlines() if line[0] != "#"
])
# Featurize Data
if featurizer == "grid":
featurizer = rgf.RdkitGridFeaturizer(voxel_width=2.0,
feature_types=[
'ecfp', 'splif', 'hbond',
'salt_bridge', 'pi_stack',
'cation_pi', 'charge'
],
flatten=True)
elif featurizer == "atomic" or featurizer == "atomic_conv":
# Pulled from PDB files. For larger datasets with more PDBs, would use
# max num atoms instead of exact.
frag1_num_atoms = 70 # for ligand atoms
if load_binding_pocket:
frag2_num_atoms = 1000
complex_num_atoms = 1070
else:
frag2_num_atoms = 24000 # for protein atoms
complex_num_atoms = 24070 # in total
max_num_neighbors = 4
# Cutoff in angstroms
neighbor_cutoff = 4
if featurizer == "atomic":
featurizer = ComplexNeighborListFragmentAtomicCoordinates(
frag1_num_atoms=frag1_num_atoms,
frag2_num_atoms=frag2_num_atoms,
complex_num_atoms=complex_num_atoms,
max_num_neighbors=max_num_neighbors,
neighbor_cutoff=neighbor_cutoff)
if featurizer == "atomic_conv":
featurizer = AtomicConvFeaturizer(
labels=labels,
frag1_num_atoms=frag1_num_atoms,
frag2_num_atoms=frag2_num_atoms,
complex_num_atoms=complex_num_atoms,
neighbor_cutoff=neighbor_cutoff,
max_num_neighbors=max_num_neighbors,
batch_size=64)
else:
raise ValueError("Featurizer not supported")
print("\nFeaturizing Complexes for \"%s\" ...\n" % data_folder)
feat_t1 = time.time()
features, failures = featurizer.featurize_complexes(
ligand_files, protein_files)
feat_t2 = time.time()
print("\nFeaturization finished, took %0.3f s." % (feat_t2 - feat_t1))
# Delete labels and ids for failing elements
labels = np.delete(labels, failures)
labels = labels.reshape((len(labels), 1))
ids = np.delete(pdbs, failures)
print("\nConstruct dataset excluding failing featurization elements...")
dataset = deepchem.data.DiskDataset.from_numpy(features, y=labels, ids=ids)
# No transformations of data
transformers = []
# Split dataset
print("\nSplit dataset...\n")
if split == None:
return pdbbind_tasks, (dataset, None, None), transformers
# TODO(rbharath): This should be modified to contain a cluster split so
# structures of the same protein aren't in both train/test
splitters = {
'index': deepchem.splits.IndexSplitter(),
'random': deepchem.splits.RandomSplitter(),
}
splitter = splitters[split]
train, valid, test = splitter.train_valid_test_split(dataset,
seed=split_seed)
all_dataset = (train, valid, test)
print("\nSaving dataset to \"%s\" ..." % save_folder)
deepchem.utils.save.save_dataset_to_disk(save_folder, train, valid, test,
transformers)
return pdbbind_tasks, all_dataset, transformers
def load_pdbbind_from_dir(data_folder,
index_files,
featurizer="grid",
split="random",
ex_ids=[],
save_dir=None):
"""Load and featurize raw PDBBind dataset from a local directory with the option to avoid certain IDs.
Parameters
----------
data_dir: String,
Specifies the data directory to store the featurized dataset.
index_files: List
List of data and labels index file paths relative to the path in data_dir
split: Str
Either "random" or "index"
feat: Str
Either "grid" or "atomic" for grid and atomic featurizations.
subset: Str
Only "core" or "refined" for now.
ex_ids: List
List of PDB IDs to avoid loading if present
save_dir: String
Path to store featurized datasets
"""
pdbbind_tasks = ["-logKd/Ki"]
index_file = os.path.join(data_folder, index_files[0])
labels_file = os.path.join(data_folder, index_files[1])
# Extract locations of data
pdbs = []
with open(index_file, "r") as g:
lines = g.readlines()
for line in lines:
line = line.split(" ")
pdb = line[0]
if len(pdb) == 4:
pdbs.append(pdb)
protein_files = [
os.path.join(data_folder, pdb, "%s_protein.pdb" % pdb)
for pdb in pdbs
if pdb not in ex_ids
]
ligand_files = [
os.path.join(data_folder, pdb, "%s_ligand.sdf" % pdb)
for pdb in pdbs
if pdb not in ex_ids
]
# Extract labels
labels_tmp = {}
with open(labels_file, "r") as f:
lines = f.readlines()
for line in lines:
# Skip comment lines
if line[0] == "#":
continue
# Lines have format
# PDB code, resolution, release year, -logKd/Ki, Kd/Ki, reference, ligand name
line = line.split()
# The base-10 logarithm, -log kd/pk
log_label = line[3]
labels_tmp[line[0]] = log_label
labels = np.array([labels_tmp[pdb] for pdb in pdbs])
print(labels)
# Featurize Data
if featurizer == "grid":
featurizer = rgf.RdkitGridFeaturizer(
voxel_width=2.0,
feature_types=[
'ecfp', 'splif', 'hbond', 'salt_bridge', 'pi_stack', 'cation_pi',
'charge'
],
flatten=True)
elif featurizer == "atomic":
# Pulled from PDB files. For larger datasets with more PDBs, would use
# max num atoms instead of exact.
frag1_num_atoms = 70 # for ligand atoms
frag2_num_atoms = 24000 # for protein atoms
complex_num_atoms = 24070 # in total
max_num_neighbors = 4
# Cutoff in angstroms
neighbor_cutoff = 4
featurizer = ComplexNeighborListFragmentAtomicCoordinates(
frag1_num_atoms, frag2_num_atoms, complex_num_atoms, max_num_neighbors,
neighbor_cutoff)
else:
raise ValueError("Featurizer not supported")
print("Featurizing Complexes")
features, failures = featurizer.featurize(ligand_files, protein_files)
# Delete labels for failing elements
labels = np.delete(labels, failures)
dataset = deepchem.data.DiskDataset.from_numpy(features, labels)
# No transformations of data
transformers = []
if split == None:
return pdbbind_tasks, (dataset, None, None), transformers
# TODO(rbharath): This should be modified to contain a cluster split so
# structures of the same protein aren't in both train/test
splitters = {
'index': deepchem.splits.IndexSplitter(),
'random': deepchem.splits.RandomSplitter(),
}
splitter = splitters[split]
train, valid, test = splitter.train_valid_test_split(dataset)
all_dataset = (train, valid, test)
if save_dir:
deepchem.utils.data_utils.save_dataset_to_disk(save_dir, train, valid, test,
transformers)
return pdbbind_tasks, all_dataset, transformers
