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