def test_morse_small_neighbor_list_energy(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(5.0)
displacement, _ = space.periodic(box_size)
metric = space.metric(displacement)
exact_energy_fn = energy.morse_pair(displacement)
R = box_size * random.uniform(key, (10, spatial_dimension),
dtype=dtype)
neighbor_fn, energy_fn = energy.morse_neighbor_list(
displacement, box_size)
nbrs = neighbor_fn(R)
self.assertAllClose(np.array(exact_energy_fn(R), dtype=dtype),
energy_fn(R, nbrs))
def test_morse_neighbor_list_force(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(15.0)
displacement, _ = space.periodic(box_size)
metric = space.metric(displacement)
exact_force_fn = quantity.force(energy.morse_pair(displacement))
r = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
neighbor_fn, energy_fn = energy.morse_neighbor_list(
displacement, box_size)
force_fn = quantity.force(energy_fn)
nbrs = neighbor_fn(r)
self.assertAllClose(np.array(exact_force_fn(r), dtype=dtype),
force_fn(r, nbrs))
def test_morse_neighbor_list_force(self, spatial_dimension, dtype, format):
key = random.PRNGKey(1)
box_size = f32(15.0)
displacement, _ = space.periodic(box_size)
metric = space.metric(displacement)
exact_force_fn = quantity.force(energy.morse_pair(displacement))
r = box_size * random.uniform(
key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
neighbor_fn, energy_fn = energy.morse_neighbor_list(
displacement, box_size, format=format)
force_fn = quantity.force(energy_fn)
nbrs = neighbor_fn.allocate(r)
if dtype == f32 and format is partition.OrderedSparse:
self.assertAllClose(
np.array(exact_force_fn(r), dtype=dtype),
force_fn(r, nbrs), atol=5e-5, rtol=5e-5)
else:
self.assertAllClose(
np.array(exact_force_fn(r), dtype=dtype),
force_fn(r, nbrs))
