Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
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])
types.data[...] = 0
types.data[:,10] = 1
batch_radii = torch.tensor(np.tile(radii, (batch_size, 1)), dtype=torch.float32, device='cuda')
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)