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