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_radius_multiples(): g1 = molgrid.GridMaker(resolution=.1, dimension=6.0) 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), 1) shape = g1.grid_dimensions(1) cpugrid = molgrid.MGrid4f(*shape) cpugrid2 = molgrid.MGrid4f(*shape) gpugrid = molgrid.MGrid4f(*shape) g1.forward((0, 0, 0), coords, cpugrid.cpu()) g1.forward((0, 0, 0), coords, gpugrid.gpu()) g1.forward((0, 0, 0), c, t, r, cpugrid2.cpu()) np.testing.assert_allclose(cpugrid.tonumpy(), gpugrid.tonumpy(), atol=1e-5) np.testing.assert_allclose(cpugrid.tonumpy(), cpugrid2.tonumpy(), atol=1e-6) g = cpugrid.tonumpy() assert g[0, 30, 30, 30] == approx(1) #cut a line across line = g[0, 30, 30, :] xvals = np.abs(np.arange(-3, 3.1, .1)) gauss = np.exp(-2 * xvals**2) for i in range(20, 41): assert line[i] == approx(gauss[i]) for i in list(range(0, 15)) + list(range(45, 61)): assert line[i] == approx(0) quad = 4 * np.exp(-2) * xvals**2 - 12 * np.exp(-2) * xvals + 9 * np.exp(-2) for i in list(range(15, 20)) + list(range(41, 45)): assert line[i] == approx(quad[i], abs=1e-5) #funkier grid g2 = molgrid.GridMaker(resolution=.1, dimension=6.0, radius_scale=0.5, gassian_radius_multiple=3.0) cpugrid = molgrid.MGrid4f(*shape) gpugrid = molgrid.MGrid4f(*shape) g2.forward((0, 0, 0), coords, cpugrid.cpu()) g2.forward((0, 0, 0), coords, gpugrid.gpu()) np.testing.assert_allclose(cpugrid.tonumpy(), gpugrid.tonumpy(), atol=1e-5) g = cpugrid.tonumpy() assert g[0, 30, 30, 30] == approx(1) #cut a line across line = g[0, 30, :, 30] xvals = np.abs(np.arange(-3, 3.1, .1)) * 2.0 gauss = np.exp(-2 * xvals**2) #should be guassian the whole way, although quickly hits numerical zero for i in range(0, 61): assert line[i] == approx(gauss[i], abs=1e-5)
def test_a_grid(): fname = datadir+"/small.types" e = molgrid.ExampleProvider(data_root=datadir+"/structs") e.populate(fname) ex = e.next() c = ex.coord_sets[1] assert np.min(c.type_index.tonumpy()) >= 0 gmaker = molgrid.GridMaker() dims = gmaker.grid_dimensions(c.max_type) # this should be grid_dims or get_grid_dims center = c.center() center = tuple(center) mgridout = molgrid.MGrid4f(*dims) mgridgpu = molgrid.MGrid4f(*dims) npout = np.zeros(dims, dtype=np.float32) torchout = torch.zeros(dims, dtype=torch.float32) cudaout = torch.zeros(dims, dtype=torch.float32, device='cuda') gmaker.forward(center, c, mgridout.cpu()) gmaker.forward(center, c, mgridgpu.gpu()) gmaker.forward(center, c, npout) gmaker.forward(center, c, torchout) gmaker.forward(center, c, cudaout) newt = gmaker.make_tensor(center, c) newa = gmaker.make_ndarray(center, c) assert 1.438691 == approx(mgridout.tonumpy().max()) assert 1.438691 == approx(mgridgpu.tonumpy().max()) assert 1.438691 == approx(npout.max()) assert 1.438691 == approx(torchout.numpy().max()) assert 1.438691 == approx(cudaout.cpu().numpy().max()) assert 1.438691 == approx(newt.cpu().numpy().max()) assert 1.438691 == approx(newa.max()) #should overwrite by default, yes? gmaker.forward(center, c, mgridout.cpu()) gmaker.forward(center, c, mgridgpu.gpu()) assert 1.438691 == approx(mgridout.tonumpy().max()) assert 1.438691 == approx(mgridgpu.tonumpy().max()) dims = gmaker.grid_dimensions(e.num_types()) mgridout = molgrid.MGrid4f(*dims) mgridgpu = molgrid.MGrid4f(*dims) gmaker.forward(ex, mgridout.cpu()) gmaker.forward(ex, mgridgpu.gpu()) gmaker.forward(ex, mgridout.cpu()) gmaker.forward(ex, mgridgpu.gpu()) assert 2.094017 == approx(mgridout.tonumpy().max()) assert 2.094017 == approx(mgridgpu.tonumpy().max())
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_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_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_dx(): fname = datadir + "/small.types" e = molgrid.ExampleProvider(data_root=datadir + "/structs") e.populate(fname) ex = e.next() c = ex.coord_sets[1] assert np.min(c.type_index.tonumpy()) >= 0 gmaker = molgrid.GridMaker() dims = gmaker.grid_dimensions( c.max_type) # this should be grid_dims or get_grid_dims center = c.coord.tonumpy().mean(axis=0) center = tuple(center.astype(float)) mgridout = molgrid.MGrid4f(*dims) gmaker.forward(center, c, mgridout.cpu()) molgrid.write_dx("tmp.dx", mgridout[0].cpu(), center, 0.5) mgridin = molgrid.read_dx("tmp.dx") os.remove("tmp.dx") g = mgridin.grid().tonumpy() go = mgridout[0].tonumpy() np.testing.assert_array_almost_equal(g, go, decimal=5) assert center == approx(list(mgridin.center())) assert mgridin.resolution() == 0.5
def test_dx(): fname = datadir + "/small.types" e = molgrid.ExampleProvider(data_root=datadir + "/structs") e.populate(fname) ex = e.next() c = ex.coord_sets[1] assert np.min(c.type_index.tonumpy()) >= 0 gmaker = molgrid.GridMaker() dims = gmaker.grid_dimensions( e.type_size()) # this should be grid_dims or get_grid_dims center = tuple(c.center()) mgridout = molgrid.MGrid4f(*dims) gmaker.forward(center, c, mgridout.cpu()) molgrid.write_dx("tmp.dx", mgridout[0].cpu(), center, 0.5) mgridin = molgrid.read_dx("tmp.dx") os.remove("tmp.dx") g = mgridin.grid().tonumpy() go = mgridout[0].tonumpy() np.testing.assert_array_almost_equal(g, go, decimal=5) assert center == approx(list(mgridin.center())) assert mgridin.resolution() == 0.5 #dump everything molgrid.write_dx_grids("/tmp/tmp", e.get_type_names(), mgridout.cpu(), center, gmaker.get_resolution(), 0.5) checkgrid = molgrid.MGrid4f(*dims) molgrid.read_dx_grids("/tmp/tmp", e.get_type_names(), checkgrid.cpu()) np.testing.assert_array_almost_equal(mgridout.tonumpy(), 2.0 * checkgrid.tonumpy(), decimal=5)
def test_vector_types(): g1 = molgrid.GridMaker(resolution=.25, dimension=6.0) c = np.array([[0, 0, 0]], np.float32) t = np.array([0], np.float32) vt = np.array([[1.0, 0]], np.float32) vt2 = np.array([[0.5, 0.5]], np.float32) r = np.array([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[0, :], 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[0, :], v2gpu[1, :] * 2.0, atol=1e-5)
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