How to use molgrid - 10 common examples

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

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)

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)

nr = Transform(q, float3(0,1,1), float3(2,0,-3))

    #random
    r = Transform(float3(0,1,1), 10.0, True)

    coord_data = [ [0,0,0],
                   [1,0,0],
                   [0,1,0],
                   [0,0,1],
                   [-1,.5,3],
                   [1,1,1],
                   [0,1,1],
                   [.333,.75,-9] ]
  
    coords = MGrid2f(8,3)
    coords2 = MGrid2f(8,3)

    for i in range(8):
        for j in range(3):
            coords[i][j] = coord_data[i][j]

    #does nr perform as expected?
    nr.forward(coords,coords2)

    assert tup(coords2[6]) == (2,1,-2) #at center
    assert tup(coords2[2]) == (2,1,-3) #on z-axis
    assert tup(coords2[5]) == (2,2,-2)

    #make sure input unchanged
    assert tup(coords[7]) == approx((0.333,.75,-9),abs=1e-5)

    # does random work both ways

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

def test_cached_example_provider():
    fname = datadir+"/small.types"
    e = molgrid.ExampleProvider(ligmolcache=datadir+'/lig.molcache2',recmolcache=datadir+'/rec.molcache2')
    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)

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

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