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_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_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_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_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_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_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
