How to use the molgrid.GridMaker function in molgrid

def test_batched_function():
    for dev in ('cuda','cpu'):
        gmaker = molgrid.GridMaker(resolution=.1,dimension=6.0)
        c = torch.tensor([[[1.0,0,0],[1,0,0]],[[0,1,0],[0,1,0]]],device=dev,dtype=torch.float32,requires_grad=True)
        vt = torch.tensor([[[0,1.0,0],[1.0,0,0]],[[0,1.0,0],[1.0,0,0]]],device=dev,dtype=torch.float32,requires_grad=True)
        r = torch.tensor([[2.0,2.0],[2.0,2.0]],device=dev,dtype=torch.float32)
        
        grid = BatchedCoords2GridFunction.apply(gmaker, (0,0,0), c, vt, r)
    
        shape = gmaker.grid_dimensions(3)    
        #make diff with gradient in center
        diff = torch.zeros(2,*shape,dtype=torch.float32,device=dev)
        diff[0,0,30,30,30] = 1.0  
        diff[0,1,30,30,30] = -1.0  
        diff[1,0,30,30,30] = 1.0  
        diff[1,1,30,30,30] = -1.0              
        grid.backward(diff)
        assert c.grad[0][0].cpu().numpy() == approx([0.60653,0,0],abs=1e-4)
        assert c.grad[0][1].cpu().numpy() == approx([-0.60653,0,0],abs=1e-4)

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.num_types()) # 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()))

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

#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,gaussian_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):

def test_make_vector_types_ex_provider(capsys):
    fname = datadir+"/ligonly.types"
    e = molgrid.ExampleProvider(molgrid.NullIndexTyper(),molgrid.defaultGninaLigandTyper, data_root=datadir+"/structs",make_vector_types=True)
    e.populate(fname)
    batch_size = 10
    b = e.next_batch(batch_size)

    gmaker = molgrid.GridMaker(dimension=23.5,radius_type_indexed=True)
    shape = gmaker.grid_dimensions(molgrid.defaultGninaLigandTyper.num_types())
    mgrid = molgrid.MGrid5f(batch_size,*shape)

    c = b[0].merge_coordinates()
    tv = c.type_vector.tonumpy()
    assert tv.shape == (10,14) 
    assert tv[0].sum() == 1.0
    assert tv[0][8] == 1.0
    
    
    e2 = molgrid.ExampleProvider(data_root=datadir+"/structs",make_vector_types=True)
    e2.populate(fname)
    b2 = e2.next_batch(batch_size)
    c2 = b2[0].merge_coordinates(unique_index_types=True)
    tv2 = c2.type_vector.tonumpy()
    assert tv2.shape == (10,28)

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

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

raise ValueError('Tensor base type %s not supported as grid type.'%str(t.dtype))
    
    return t

#extend grid maker to create pytorch Tensor
def make_grid_tensor(gridmaker, center, c):
    '''Create appropriately sized pytorch tensor of grid densities.  set_gpu_enabled can be used to control if result is located on the cpu or gpu'''
    dims = gridmaker.grid_dimensions(c.max_type) # this should be grid_dims or get_grid_dims
    if mg.get_gpu_enabled():
        t = torch.zeros(dims, dtype=torch.float32, device='cuda:0')
    else:
        t = torch.zeros(dims, dtype=torch.float32)
    gridmaker.forward(center, c, t)
    return t 

mg.GridMaker.make_tensor = make_grid_tensor
    
class Coords2GridFunction(torch.autograd.Function):
    '''Layer for converting from coordinate and type tensors to a molecular grid'''
    
    @staticmethod
    def forward(ctx, gmaker, center, coords, types, radii):
        '''coords are Nx3, types are NxT, radii are N'''
        ctx.save_for_backward(coords, types, radii)
        ctx.gmaker = gmaker
        ctx.center = center
        shape = gmaker.grid_dimensions(types.shape[1]) #ntypes == nchannels
        output = torch.empty(*shape,dtype=coords.dtype,device=coords.device)
        gmaker.forward(center, coords, types, radii, output)
        return output
        
    @staticmethod