How to use the molgrid.MGrid4f function in molgrid

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

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

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)

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

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)

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

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

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