def test_backward_vec():
g1 = molgrid.GridMaker(resolution=.1, dimension=6.0)
c = np.array([[1.0, 0, 0], [-1, -1, 0]], np.float32)
t = np.array([[0, 1.0, 0], [1.0, 0, 0]], np.float32)
r = np.array([2.0, 2.0], np.float32)
coords = molgrid.CoordinateSet(c, t, r)
shape = g1.grid_dimensions(3)
#make diff with gradient in center
diff = molgrid.MGrid4f(*shape)
diff[0, 30, 30, 30] = 1.0
diff[1, 30, 30, 30] = -1.0
cpuatoms = molgrid.MGrid2f(2, 3)
cputypes = molgrid.MGrid2f(2, 3)
gpuatoms = molgrid.MGrid2f(2, 3)
gputypes = molgrid.MGrid2f(2, 3)
g1.backward((0, 0, 0), coords, diff.cpu(), cpuatoms.cpu(), cputypes.cpu())
assert cputypes[0][0] > 0
assert cputypes[0][1] < 0
assert cputypes[0][2] == 0
g1.backward((0, 0, 0), coords, diff.gpu(), gpuatoms.gpu(), gputypes.gpu())
np.testing.assert_allclose(gpuatoms.tonumpy(),
cpuatoms.tonumpy(),
atol=1e-5)
np.testing.assert_allclose(gputypes.tonumpy(),
cputypes.tonumpy(),
atol=1e-5)
def test_backwards():
g1 = molgrid.GridMaker(resolution=.1, dimension=6.0)
c = np.array([[1.0, 0, 0]], np.float32)
t = np.array([0], np.float32)
r = np.array([2.0], np.float32)
coords = molgrid.CoordinateSet(molgrid.Grid2f(c), molgrid.Grid1f(t),
molgrid.Grid1f(r), 1)
shape = g1.grid_dimensions(1)
#make diff with gradient in center
diff = molgrid.MGrid4f(*shape)
diff[0, 30, 30, 30] = 1.0
cpuatoms = molgrid.MGrid2f(1, 3)
gpuatoms = molgrid.MGrid2f(1, 3)
#apply random rotation
T = molgrid.Transform((0, 0, 0), 0, True)
T.forward(coords, coords)
g1.backward((0, 0, 0), coords, diff.cpu(), cpuatoms.cpu())
g1.backward((0, 0, 0), coords, diff.gpu(), gpuatoms.gpu())
T.backward(cpuatoms.cpu(), cpuatoms.cpu(), False)
T.backward(gpuatoms.gpu(), gpuatoms.gpu(), False)
print(cpuatoms.tonumpy(), gpuatoms.tonumpy())
# results should be ~ -.6, 0, 0
np.testing.assert_allclose(cpuatoms.tonumpy(),
gpuatoms.tonumpy(),
atol=1e-5)
np.testing.assert_allclose(cpuatoms.tonumpy().flatten(),
[-0.60653067, 0, 0],
atol=1e-5)
def test_vector_sum_types():
fname = datadir + "/ligonly.types"
e = molgrid.ExampleProvider(data_root=datadir + "/structs",
make_vector_types=True)
e.populate(fname)
batch_size = 10
b = e.next_batch(batch_size)
sum = molgrid.MGrid2f(batch_size, e.num_types())
b.sum_types(sum)
sum2 = np.zeros(sum.shape, np.float32)
b.sum_types(sum2)
sum3 = torch.empty(sum.shape, dtype=torch.float32, device='cuda')
b.sum_types(sum3)
np.testing.assert_allclose(sum.tonumpy(),
sum3.detach().cpu().numpy(),
atol=1e-5)
np.testing.assert_allclose(sum.tonumpy(), sum2, atol=1e-5)
np.testing.assert_allclose(sum[0].tonumpy(), [
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 2., 3., 0., 0.,
0., 0., 0., 0., 2., 2., 1., 0., 0., 0.
],
atol=1e-5)
e = molgrid.ExampleProvider(molgrid.NullIndexTyper(),
molgrid.defaultGninaLigandTyper,
data_root=datadir + "/structs",
make_vector_types=True)
e.populate(fname)
b = e.next_batch(batch_size)
sum = molgrid.MGrid2f(batch_size, e.num_types())
b.sum_types(sum)
np.testing.assert_allclose(
sum[0].tonumpy(),
[2., 3., 0., 0., 0., 0., 0., 0., 2., 2., 1., 0., 0., 0.],
atol=1e-5)
def test_type_radii():
g1 = molgrid.GridMaker(resolution=.25,
dimension=6.0,
radius_type_indexed=True)
c = np.array([[0, 0, 0]], np.float32)
t = np.array([0], np.float32)
r = np.array([1.0], np.float32)
coords = molgrid.CoordinateSet(molgrid.Grid2f(c), molgrid.Grid1f(t),
molgrid.Grid1f(r), 2)
coords.make_vector_types(True, [3.0, 1.0])
shape = g1.grid_dimensions(3) #includes dummy type
reference = molgrid.MGrid4f(*shape)
gpudata = molgrid.MGrid4f(*shape)
assert g1.get_radii_type_indexed()
g1.forward((0, 0, 0), coords, reference.cpu())
g1.forward((0, 0, 0), coords, gpudata.gpu())
np.testing.assert_allclose(reference.tonumpy(),
gpudata.tonumpy(),
atol=1e-5)
assert reference.tonumpy().sum() > 2980 #radius of 1 would be 116
reference.fill_zero()
reference[0][20][12][12] = -1
reference[1][20][12][12] = 1
reference[2][20][12][12] = 2
cpuatoms = molgrid.MGrid2f(1, 3)
cputypes = molgrid.MGrid2f(1, 3)
gpuatoms = molgrid.MGrid2f(1, 3)
gputypes = molgrid.MGrid2f(1, 3)
g1.backward((0, 0, 0), coords, reference.cpu(), cpuatoms.cpu(),
cputypes.cpu())
assert cpuatoms[0][0] < 0
assert cpuatoms[0][1] == 0
assert cpuatoms[0][2] == 0
assert cputypes[0][0] < 0
assert cputypes[0][1] == 0
assert cputypes[0][2] == 0
g1.backward((0, 0, 0), coords, reference.gpu(), gpuatoms.gpu(),
gputypes.gpu())
np.testing.assert_allclose(gpuatoms.tonumpy(),
cpuatoms.tonumpy(),
atol=1e-5)
np.testing.assert_allclose(gputypes.tonumpy(),
cputypes.tonumpy(),
atol=1e-5)
def test_gnina_example_provider():
fname = datadir + "/small.types"
e = molgrid.ExampleProvider(data_root=datadir + "/structs")
e.populate(fname)
batch_size = 100
batch = e.next_batch(batch_size)
#extract labels
nlabels = e.num_labels()
assert nlabels == 3
labels = molgrid.MGrid2f(batch_size, nlabels)
gpulabels = molgrid.MGrid2f(batch_size, nlabels)
batch.extract_labels(labels.cpu())
batch.extract_labels(gpulabels.gpu())
assert np.array_equal(labels.tonumpy(), gpulabels.tonumpy())
label0 = molgrid.MGrid1f(batch_size)
label1 = molgrid.MGrid1f(batch_size)
label2 = molgrid.MGrid1f(batch_size)
batch.extract_label(0, label0.cpu())
batch.extract_label(1, label1.cpu())
batch.extract_label(2, label2.gpu())
assert label0[0] == 1
assert label1[0] == approx(6.05)
assert label2[0] == approx(0.162643)
assert labels[0, 0] == 1
assert labels[0][1] == approx(6.05)
assert labels[0][2] == approx(0.162643)
for i in range(nlabels):
assert label0[i] == labels[i][0]
assert label1[i] == labels[i][1]
assert label2[i] == labels[i][2]
ex = batch[0]
crec = ex.coord_sets[0]
assert crec.size() == 1781
assert list(crec.coords[0]) == approx([45.042, 12.872, 13.001])
assert crec.radii[0] == approx(1.8)
assert list(crec.type_index)[:10] == [
6.0, 1.0, 1.0, 7.0, 0.0, 6.0, 1.0, 1.0, 7.0, 1.0
]
clig = ex.coord_sets[1]
assert clig.size() == 10
assert list(clig.coords[9]) == approx([27.0536, 3.2453, 32.4511])
assert list(clig.type_index) == [
8.0, 1.0, 1.0, 9.0, 10.0, 0.0, 0.0, 1.0, 9.0, 8.0
]
batch = e.next_batch(1)
a = np.array([0], dtype=np.float32)
batch.extract_label(1, a)
def test_numpy():
mg3d = molgrid.MGrid3d(3,3,3)
mg2f = molgrid.MGrid2f(2,4)
a2f = np.arange(27).astype(np.float32).reshape(3,-1)
a3d = np.arange(27).astype(np.float).reshape(3,3,-1)
g2f = molgrid.Grid2f(a2f)
g3d = molgrid.Grid3d(a3d)
g2f[0,0] = 100
g3d[0,0,0] = 101
assert a2f[0,0] == 100
assert a3d[0,0,0] == 101
mg2f.copyFrom(g2f)
mg3d.copyFrom(g3d)
assert mg2f[1,3] == 7
mg2f[1,3] = 200
assert g2f[1,3] == 12
assert a2f[1,3] == 12
mg2f.copyTo(g2f)
assert g2f[0,7] == 200
assert a2f[0,7] == 200
def test_coords2grid():
gmaker = molgrid.GridMaker(resolution=0.5,
dimension=23.5,
radius_scale=1,
radius_type_indexed=True)
n_types = molgrid.defaultGninaLigandTyper.num_types()
radii = np.array(list(molgrid.defaultGninaLigandTyper.get_type_radii()),
np.float32)
dims = gmaker.grid_dimensions(n_types)
grid_size = dims[0] * dims[1] * dims[2] * dims[3]
c2grid = molgrid.Coords2Grid(gmaker, center=(0, 0, 0))
n_atoms = 2
batch_size = 1
coords = nn.Parameter(torch.randn(n_atoms, 3, device='cuda'))
types = nn.Parameter(torch.randn(n_atoms, n_types + 1, device='cuda'))
coords.data[0, :] = torch.tensor([1, 0, 0])
coords.data[1, :] = torch.tensor([-1, 0, 0])
types.data[...] = 0
types.data[:, 10] = 1
batch_radii = torch.tensor(np.tile(radii, (batch_size, 1)),
dtype=torch.float32,
device='cuda')
grid_gen = c2grid(coords.unsqueeze(0),
types.unsqueeze(0)[:, :, :-1], batch_radii)
assert float(grid_gen[0][10].sum()) == approx(float(grid_gen.sum()))
assert grid_gen.sum() > 0
target = torch.zeros_like(grid_gen)
target[0, :, 24, 24, 24] = 1000.0
grad_coords = molgrid.MGrid2f(n_atoms, 3)
grad_types = molgrid.MGrid2f(n_atoms, n_types)
r = molgrid.MGrid1f(len(radii))
r.copyFrom(radii)
grid_loss = F.mse_loss(target, grid_gen)
grid_loss.backward()
print(grid_loss)
print(coords.grad.detach().cpu().numpy())
def test_example_provider_iterator_interface():
fname = datadir+"/small.types"
BSIZE=25
e = molgrid.ExampleProvider(data_root=datadir+"/structs",default_batch_size=BSIZE)
e.populate(fname)
e2 = molgrid.ExampleProvider(data_root=datadir+"/structs",default_batch_size=BSIZE)
e2.populate(fname)
nlabels = e.num_labels()
labels = molgrid.MGrid2f(BSIZE,nlabels)
labels2 = molgrid.MGrid2f(BSIZE,nlabels)
for (i, b) in enumerate(e):
b2 = e2.next_batch()
b.extract_labels(labels.cpu())
b2.extract_labels(labels2.cpu())
np.testing.assert_allclose(labels,labels2)
if i > 10:
break
def test_vector_types_mol():
'''Test vector types with a real molecule'''
fname = datadir+"/small.types"
e = molgrid.ExampleProvider(data_root=datadir+"/structs")
e.populate(fname)
ex = e.next()
ev = molgrid.ExampleProvider(data_root=datadir+"/structs",make_vector_types=True)
ev.populate(fname)
exv = ev.next()
assert exv.has_vector_types()
assert not ex.has_vector_types()
gmaker = molgrid.GridMaker()
dims = gmaker.grid_dimensions(ex.num_types()) # this should be grid_dims or get_grid_dims
mgridout = molgrid.MGrid4f(*dims)
mgridgpu = molgrid.MGrid4f(*dims)
mgridoutv = molgrid.MGrid4f(*dims)
mgridgpuv = molgrid.MGrid4f(*dims)
d = np.ones(dims,np.float32)
diff = molgrid.MGrid4f(*dims)
diff.copyFrom(d)
gmaker.forward(ex, mgridout.cpu())
gmaker.forward(ex, mgridgpu.gpu())
center = ex.coord_sets[-1].center()
c = ex.merge_coordinates()
backcoordscpu = molgrid.MGrid2f(c.size(),3)
backcoordsgpu = molgrid.MGrid2f(c.size(),3)
gmaker.backward(center, c, diff.cpu(), backcoordscpu.cpu())
gmaker.backward(center, c, diff.gpu(), backcoordsgpu.gpu())
#vector types
gmaker.set_radii_type_indexed(True)
gmaker.forward(exv, mgridoutv.cpu())
gmaker.forward(exv, mgridgpuv.gpu())
cv = exv.merge_coordinates()
vbackcoordscpu = molgrid.MGrid2f(cv.size(),3)
vbackcoordsgpu = molgrid.MGrid2f(cv.size(),3)
vbacktypescpu = molgrid.MGrid2f(cv.size(),cv.num_types())
vbacktypesgpu = molgrid.MGrid2f(cv.size(),cv.num_types())
gmaker.backward(center, cv, diff.cpu(), vbackcoordscpu.cpu(),vbacktypescpu.cpu())
gmaker.backward(center, cv, diff.gpu(), vbackcoordsgpu.gpu(),vbacktypesgpu.gpu())
np.testing.assert_allclose(mgridout.tonumpy(),mgridoutv.tonumpy(),atol=1e-5)
np.testing.assert_allclose(mgridgpu.tonumpy(),mgridgpuv.tonumpy(),atol=1e-5)
np.testing.assert_allclose(mgridoutv.tonumpy(),mgridgpuv.tonumpy(),atol=1e-5)
np.testing.assert_allclose(vbackcoordscpu.tonumpy(),backcoordscpu.tonumpy(),atol=1e-5)
np.testing.assert_allclose(vbackcoordsgpu.tonumpy(),backcoordsgpu.tonumpy(),atol=1e-5)
np.testing.assert_allclose(vbackcoordscpu.tonumpy(),vbackcoordsgpu.tonumpy(),atol=1e-4)
np.testing.assert_allclose(vbacktypescpu.tonumpy(),vbacktypesgpu.tonumpy(),atol=1e-4)
def test_grid():
g = molgrid.MGrid2f(10, 2)
assert g.size() == 20
g2 = molgrid.Grid2f(g.cpu())
g2[5][1] = 3.0
assert g[5, 1] == g2[5, 1]
g1 = g2[5]
g1[1] = 4.0
assert g[5, 1] == 4.0
gclone = g.clone()
gclone[5, 1] = 5.5
assert g[5, 1] == 4.0
assert gclone[5][1] == 5.5
def test_mgrid_copyto_tensor():
mg2 = molgrid.MGrid2f(3,4)
mg3 = molgrid.MGrid3d(3,4,2)
for i in range(3):
for j in range(4):
mg2[i,j] = i*j+1
for k in range(2):
mg3[i,j,k] = i*j+1+k
t2 = torch.FloatTensor(3,4)
t3 = torch.DoubleTensor(3,4,2)
mg2.copyTo(t2)
mg3.copyTo(t3)
for i in range(3):
for j in range(4):
assert t2[i,j] == i*j+1
for k in range(2):
assert t3[i,j,k] == i*j+1+k
def test_mgrid_copyfrom_tensor_cuda():
mg2 = molgrid.MGrid2f(3,4)
mg3 = molgrid.MGrid3d(3,4,2)
t2 = torch.cuda.FloatTensor(3,4)
t3 = torch.cuda.DoubleTensor(3,4,2)
for i in range(3):
for j in range(4):
t2[i,j] = i*j+1
for k in range(2):
t3[i,j,k] = i*j+1+k
mg2.copyFrom(t2)
mg3.copyFrom(t3)
for i in range(3):
for j in range(4):
assert mg2[i,j] == i*j+1
for k in range(2):
assert mg3[i,j,k] == i*j+1+k
def test_vector_types():
g1 = molgrid.GridMaker(resolution=.25,dimension=6.0)
c = np.array([[0,0,0],[2,0,0]],np.float32)
t = np.array([0,1],np.float32)
vt = np.array([[1.0,0],[0,1.0]],np.float32)
vt2 = np.array([[0.5,0.0],[0.0,0.5]],np.float32)
r = np.array([1.0,1.0],np.float32)
coords = molgrid.CoordinateSet(molgrid.Grid2f(c),molgrid.Grid1f(t),molgrid.Grid1f(r),2)
vcoords = molgrid.CoordinateSet(molgrid.Grid2f(c),molgrid.Grid2f(vt),molgrid.Grid1f(r))
v2coords = molgrid.CoordinateSet(molgrid.Grid2f(c),molgrid.Grid2f(vt2),molgrid.Grid1f(r))
shape = g1.grid_dimensions(2)
reference = molgrid.MGrid4f(*shape)
vgrid = molgrid.MGrid4f(*shape)
v2grid = molgrid.MGrid4f(*shape)
v3grid = molgrid.MGrid4f(*shape)
g1.forward((0,0,0),coords, reference.cpu())
g1.forward((0,0,0),vcoords, vgrid.cpu())
g1.forward((0,0,0),v2coords, v2grid.cpu())
g1.forward((0,0,0),c,vt,r, v3grid.cpu())
np.testing.assert_allclose(reference.tonumpy(),vgrid.tonumpy(),atol=1e-5)
np.testing.assert_allclose(vgrid.tonumpy(),v3grid.tonumpy(),atol=1e-6)
v2g = v2grid.tonumpy()
g = reference.tonumpy()
np.testing.assert_allclose(g[0,:],v2g[0,:]*2.0,atol=1e-5)
np.testing.assert_allclose(g[1,:],v2g[1,:]*2.0,atol=1e-5)
vgridgpu = molgrid.MGrid4f(*shape)
v2gridgpu = molgrid.MGrid4f(*shape)
g1.forward((0,0,0),vcoords, vgridgpu.gpu())
g1.forward((0,0,0),v2coords, v2gridgpu.gpu())
np.testing.assert_allclose(reference.tonumpy(),vgridgpu.tonumpy(),atol=1e-5)
v2gpu = v2gridgpu.tonumpy()
np.testing.assert_allclose(g[0,:],v2gpu[0,:]*2.0,atol=1e-5)
np.testing.assert_allclose(g[1,:],v2gpu[1,:]*2.0,atol=1e-5)
#create target grid with equal type density at 1,0,0
tc = molgrid.Grid2f(np.array([[1,0,0]],np.float32))
tv = molgrid.Grid2f(np.array([[0.5,0.5]],np.float32))
tr = molgrid.Grid1f(np.array([1.0],np.float32))
targetc = molgrid.CoordinateSet(tc,tv,tr)
tgrid = molgrid.MGrid4f(*shape)
g1.forward((0,0,0),targetc,tgrid.cpu())
gradc = molgrid.MGrid2f(2,3)
gradt = molgrid.MGrid2f(2,2)
g1.backward((0,0,0),vcoords,tgrid.cpu(),gradc.cpu(),gradt.cpu())
assert gradc[0,0] == approx(-gradc[1,0],abs=1e-4)
assert gradc[0,0] > 0
gradc.fill_zero()
gradt.fill_zero()
g1.backward((0,0,0),vcoords,tgrid.gpu(),gradc.gpu(),gradt.gpu())
assert gradc[0,0] == approx(-gradc[1,0],abs=1e-4)
assert gradc[0,0] > 0
