def test_pair_grid_force_incommensurate(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(12.1)
cell_size = f32(3.0)
displacement, _ = space.periodic(box_size)
energy_fn = energy.soft_sphere_pair(displacement, quantity.Dynamic)
force_fn = quantity.force(energy_fn)
R = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
grid_force_fn = jit(smap.grid(force_fn, box_size, cell_size, R))
species = np.zeros((PARTICLE_COUNT, ), dtype=np.int64)
self.assertAllClose(np.array(force_fn(R, species, 1), dtype=dtype),
grid_force_fn(R), True)
def test_pair_grid_force_nonuniform(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
if spatial_dimension == 2:
box_size = f32(np.array([[8.0, 10.0]]))
else:
box_size = f32(np.array([[8.0, 10.0, 12.0]]))
cell_size = f32(2.0)
displacement, _ = space.periodic(box_size[0])
energy_fn = energy.soft_sphere_pair(displacement, quantity.Dynamic)
force_fn = quantity.force(energy_fn)
R = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
grid_force_fn = smap.grid(force_fn, box_size, cell_size, R)
species = np.zeros((PARTICLE_COUNT, ), dtype=np.int64)
self.assertAllClose(np.array(force_fn(R, species, 1), dtype=dtype),
grid_force_fn(R), True)
def test_pair_grid_energy(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f16(9.0)
cell_size = f16(2.0)
displacement, _ = space.periodic(box_size)
metric = space.metric(displacement)
energy_fn = smap.pair(energy.soft_sphere,
metric,
quantity.Dynamic,
reduce_axis=(1, ),
keepdims=True)
R = box_size * random.uniform(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
grid_energy_fn = smap.grid(energy_fn, box_size, cell_size, R)
species = np.zeros((PARTICLE_COUNT, ), dtype=np.int64)
self.assertAllClose(np.array(energy_fn(R, species, 1), dtype=dtype),
grid_energy_fn(R), True)
