Commit c70e2963 authored by hgarrereyn's avatar hgarrereyn
Browse files

metric improvements, add config files

parent d26f81d2
......@@ -169,7 +169,7 @@ class FragmentDataset(Dataset):
def __init__(self, fragment_file, rec_typer, lig_typer, filter_rec=None,
fdist_min=None, fdist_max=None, fmass_min=None,
fmass_max=None, verbose=False):
fmass_max=None, verbose=False, skip_remap=False):
"""Initializes the fragment dataset.
Args:
......@@ -177,6 +177,7 @@ class FragmentDataset(Dataset):
rec_typer: function to map receptor rows to layer index
lig_typer: function to map ligand rows to layer index
filter_rec: list of receptor ids to use (or None to use all)
skip_remap: if True, don't prepare atom type information
(filtering options):
fdist_min: minimum fragment distance
......@@ -185,6 +186,7 @@ class FragmentDataset(Dataset):
fmass_max: maximum fragment mass (Da)
"""
self.verbose = verbose
self._skip_remap = skip_remap
self.rec = self._load_rec(fragment_file, rec_typer)
self.frag = self._load_fragments(fragment_file, lig_typer)
......@@ -205,8 +207,9 @@ class FragmentDataset(Dataset):
r = tqdm.tqdm(r, desc='Remap receptor atoms')
rec_remapped = np.zeros(len(rec_types)).astype(np.int32)
for i in r:
rec_remapped[i] = rec_typer(rec_types[i])
if not self._skip_remap:
for i in r:
rec_remapped[i] = rec_typer(rec_types[i])
# create rec mapping
rec_mapping = {}
......@@ -239,7 +242,11 @@ class FragmentDataset(Dataset):
frag_coords = frag_data[:,:3].astype(np.float32)
frag_types = frag_data[:,3].astype(np.int32)
frag_remapped = np.vectorize(lig_typer)(frag_types)
frag_remapped = None
if self._skip_remap:
frag_remapped = np.zeros(len(frag_types))
else:
frag_remapped = np.vectorize(lig_typer)(frag_types)
# find and save connection point
r = range(len(frag_lookup))
......
......@@ -7,6 +7,7 @@ import torch.nn as nn
import torch.nn.functional as F
import wandb
import tqdm
import numpy as np
from leadopt.models.voxel import VoxelFingerprintNet
from leadopt.data_util import FragmentDataset, FingerprintDataset, LIG_TYPER,\
......@@ -413,6 +414,66 @@ class VoxelNet(LeadoptModel):
val_metrics.clear()
def run_test(self, save_path):
# load test dataset
test_dat = FragmentDataset(
self._args['fragments'],
rec_typer=REC_TYPER[self._args['rec_typer']],
lig_typer=LIG_TYPER[self._args['lig_typer']],
filter_rec=partitions.TEST,
fdist_min=self._args['fdist_min'],
fdist_max=self._args['fdist_max'],
fmass_min=self._args['fmass_min'],
fmass_max=self._args['fmass_max'],
verbose=True
)
fingerprints = FingerprintDataset(self._args['fingerprints'])
self._models['voxel'].eval()
predicted_fp = np.zeros((
len(test_dat),
self._args['samples_per_example'],
self._args['output_size']))
smiles = [test_dat[i]['smiles'] for i in range(len(test_dat))]
correct_fp = fingerprints.for_smiles(smiles).numpy()
# (example_idx, sample_idx)
queries = []
for i in range(len(test_dat)):
queries += [(i,x) for x in range(self._args['samples_per_example'])]
# run inference
pbar = tqdm.tqdm(
range(0, len(queries), self._args['batch_size']), desc='Inference')
for i in pbar:
batch = queries[i:i+self._args['batch_size']]
torch_grid, examples = get_batch(
test_dat,
self._args['rec_channels'],
self._args['lig_channels'],
batch_size=self._args['batch_size'],
batch_set=[x[0] for x in batch],
width=self._args['grid_width'],
res=self._args['grid_res'],
ignore_receptor=self._args['ignore_receptor'],
ignore_parent=self._args['ignore_parent']
)
predicted = self._models['voxel'](torch_grid)
for j in range(len(batch)):
example_idx, sample_idx = batch[j]
predicted_fp[example_idx][sample_idx] = predicted[j].detach().cpu().numpy()
np.save(os.path.join(save_path, 'predicted_fp.npy'), predicted_fp)
np.save(os.path.join(save_path, 'correct_fp.npy'), correct_fp)
print('done.')
MODELS = {
'voxelnet': VoxelNet
......
{
"version": "voxelnet",
"no_partitions": false,
"fragments": "./data/fragments_desc.h5",
"fingerprints": "./data/fp_rdk_desc.h5",
"learning_rate": 0.001,
"num_epochs": 100,
"test_steps": 400,
"batch_size": 32,
"grid_width": 24,
"grid_res": 1.0,
"fdist_min": null,
"fdist_max": 3,
"fmass_min": null,
"fmass_max": 150,
"ignore_receptor": false,
"ignore_parent": false,
"rec_typer": "simple",
"lig_typer": "simple",
"rec_channels": 4,
"lig_channels": 3,
"in_channels": 7,
"output_size": 2048,
"pad": false,
"blocks": [32, 64],
"fc": [2048],
"use_all_labels": true,
"dist_fn": "cos",
"loss": "direct"
}
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.01, "num_epochs": 100, "test_steps": 400, "batch_size": 32, "grid_width": 16, "grid_res": 0.5, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [32, 64], "fc": [1024, 1024], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 16, "grid_width": 16, "grid_res": 2, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 64], "fc": [512, 512], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.0001, "num_epochs": 100, "test_steps": 400, "batch_size": 128, "grid_width": 32, "grid_res": 0.5, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [16, 32], "fc": [256], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.0001, "num_epochs": 100, "test_steps": 400, "batch_size": 128, "grid_width": 24, "grid_res": 1, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 64], "fc": [256], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.01, "num_epochs": 100, "test_steps": 400, "batch_size": 32, "grid_width": 24, "grid_res": 0.5, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [32, 64], "fc": [1024, 1024], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 16, "grid_width": 32, "grid_res": 1, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 64], "fc": [256, 256], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 128, "grid_width": 24, "grid_res": 0.5, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 128], "fc": [1024, 1024], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.01, "num_epochs": 100, "test_steps": 400, "batch_size": 32, "grid_width": 32, "grid_res": 1, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [32, 64], "fc": [2048], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 64, "grid_width": 32, "grid_res": 1, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [16, 32], "fc": [2048], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.0001, "num_epochs": 100, "test_steps": 400, "batch_size": 128, "grid_width": 32, "grid_res": 2, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [16, 32], "fc": [1024, 1024], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.0001, "num_epochs": 100, "test_steps": 400, "batch_size": 16, "grid_width": 24, "grid_res": 0.5, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 64], "fc": [512], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 128, "grid_width": 16, "grid_res": 2, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 64], "fc": [1024], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 16, "grid_width": 24, "grid_res": 0.5, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [32, 64], "fc": [256, 256], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 32, "grid_width": 16, "grid_res": 2, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [16, 32], "fc": [256], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 16, "grid_width": 24, "grid_res": 1, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 64], "fc": [1024], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 16, "grid_width": 32, "grid_res": 1, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 128], "fc": [256, 256], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
{"version": "voxelnet", "no_partitions": false, "fragments": "./data/fragments_desc.h5", "fingerprints": "./data/fp_rdk_desc.h5", "learning_rate": 0.001, "num_epochs": 100, "test_steps": 400, "batch_size": 32, "grid_width": 24, "grid_res": 0.5, "fdist_min": null, "fdist_max": 3, "fmass_min": null, "fmass_max": 150, "ignore_receptor": false, "ignore_parent": false, "rec_typer": "simple", "lig_typer": "simple", "rec_channels": 4, "lig_channels": 3, "in_channels": 7, "output_size": 2048, "pad": false, "blocks": [64, 128], "fc": [256], "use_all_labels": true, "dist_fn": "cos", "loss": "direct"}
\ No newline at end of file
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment