def test_get_cells(self):
"""Test that tensorflow can compute grid cells."""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
with self.session() as sess:
ndim = 3
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim).eval()
assert len(cells.shape) == 2
assert cells.shape[0] == 4**ndim
ndim = 2
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim).eval()
assert len(cells.shape) == 2
assert cells.shape[0] == 4**ndim
def test_get_cells(self):
"""Test that tensorflow can compute grid cells."""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
with self.session() as sess:
ndim = 3
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim).eval()
assert len(cells.shape) == 2
assert cells.shape[0] == 4**ndim
ndim = 2
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim).eval()
assert len(cells.shape) == 2
assert cells.shape[0] == 4**ndim
def test_put_atoms_in_cells(self):
"""Test that atoms can be partitioned into spatial cells."""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
ndim = 3
k = 5
# The number of cells which we should theoretically have
n_cells = ((stop - start) / nbr_cutoff)**ndim
with self.session() as sess:
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
coords = np.random.rand(N, ndim)
_, atoms_in_cells = put_atoms_in_cells(coords, cells, N, n_cells, ndim, k)
atoms_in_cells = atoms_in_cells.eval()
assert len(atoms_in_cells) == n_cells
# Each atom neighbors tensor should be (k, ndim) shaped.
for atoms in atoms_in_cells:
assert atoms.shape == (k, ndim)
def test_get_cells_for_atoms(self):
"""Test that atoms are placed in the correct cells."""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
ndim = 3
k = 5
# The number of cells which we should theoretically have
n_cells = ((stop - start) / nbr_cutoff)**ndim
# TODO(rbharath): The test below only checks that shapes work out.
# Need to do a correctness implementation vs. a simple CPU impl.
with self.session() as sess:
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
coords = np.random.rand(N, ndim)
cells_for_atoms = get_cells_for_atoms(coords, cells, N, n_cells, ndim)
cells_for_atoms = cells_for_atoms.eval()
assert cells_for_atoms.shape == (N, 1)
def test_put_atoms_in_cells(self):
"""Test that atoms can be partitioned into spatial cells."""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
ndim = 3
k = 5
# The number of cells which we should theoretically have
n_cells = ((stop - start) / nbr_cutoff)**ndim
with self.session() as sess:
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
coords = np.random.rand(N, ndim)
_, atoms_in_cells = put_atoms_in_cells(coords, cells, N, n_cells,
ndim, k)
atoms_in_cells = atoms_in_cells.eval()
assert len(atoms_in_cells) == n_cells
# Each atom neighbors tensor should be (k, ndim) shaped.
for atoms in atoms_in_cells:
assert atoms.shape == (k, ndim)
def test_get_cells_for_atoms(self):
"""Test that atoms are placed in the correct cells."""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
ndim = 3
k = 5
# The number of cells which we should theoretically have
n_cells = ((stop - start) / nbr_cutoff)**ndim
# TODO(rbharath): The test below only checks that shapes work out.
# Need to do a correctness implementation vs. a simple CPU impl.
with self.session() as sess:
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
coords = np.random.rand(N, ndim)
cells_for_atoms = get_cells_for_atoms(coords, cells, N, n_cells,
ndim)
cells_for_atoms = cells_for_atoms.eval()
assert cells_for_atoms.shape == (N, 1)
def test_compute_neighbor_cells(self):
"""Test that indices of neighboring cells can be computed."""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
ndim = 3
# The number of cells which we should theoretically have
n_cells = ((stop - start) / nbr_cutoff)**ndim
# TODO(rbharath): The test below only checks that shapes work out.
# Need to do a correctness implementation vs. a simple CPU impl.
with self.session() as sess:
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
nbr_cells = compute_neighbor_cells(cells, ndim, n_cells)
nbr_cells = nbr_cells.eval()
assert len(nbr_cells) == n_cells
nbr_cells = [nbr_cell for nbr_cell in nbr_cells]
for nbr_cell in nbr_cells:
assert nbr_cell.shape == (26, )
def test_compute_closest_neighbors(self):
"""Test that closest neighbors can be computed properly"""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
ndim = 3
k = 5
# The number of cells which we should theoretically have
n_cells = ((stop - start) / nbr_cutoff)**ndim
# TODO(rbharath): The test below only checks that shapes work out.
# Need to do a correctness implementation vs. a simple CPU impl.
with self.session() as sess:
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
nbr_cells = compute_neighbor_cells(cells, ndim, n_cells)
coords = np.random.rand(N, ndim)
_, atoms_in_cells = put_atoms_in_cells(coords, cells, N, n_cells, ndim, k)
nbrs = compute_closest_neighbors(coords, cells, atoms_in_cells, nbr_cells,
N, n_cells)
def test_compute_neighbor_cells(self):
"""Test that indices of neighboring cells can be computed."""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
ndim = 3
# The number of cells which we should theoretically have
n_cells = ((stop - start) / nbr_cutoff)**ndim
# TODO(rbharath): The test below only checks that shapes work out.
# Need to do a correctness implementation vs. a simple CPU impl.
with self.session() as sess:
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
nbr_cells = compute_neighbor_cells(cells, ndim, n_cells)
nbr_cells = nbr_cells.eval()
assert len(nbr_cells) == n_cells
nbr_cells = [nbr_cell for nbr_cell in nbr_cells]
for nbr_cell in nbr_cells:
assert nbr_cell.shape == (26,)
def test_compute_closest_neighbors(self):
"""Test that closest neighbors can be computed properly"""
N = 10
start = 0
stop = 4
nbr_cutoff = 1
ndim = 3
k = 5
# The number of cells which we should theoretically have
n_cells = ((stop - start) / nbr_cutoff)**ndim
# TODO(rbharath): The test below only checks that shapes work out.
# Need to do a correctness implementation vs. a simple CPU impl.
with self.session() as sess:
cells = get_cells(start, stop, nbr_cutoff, ndim=ndim)
nbr_cells = compute_neighbor_cells(cells, ndim, n_cells)
coords = np.random.rand(N, ndim)
_, atoms_in_cells = put_atoms_in_cells(coords, cells, N, n_cells,
ndim, k)
nbrs = compute_closest_neighbors(coords, cells, atoms_in_cells,
nbr_cells, N, n_cells)