def soft_sphere_cell_list(
displacement_or_metric,
box_size,
R_example,
species=None,
sigma=1.0,
epsilon=1.0,
alpha=2.0):
"""Convenience wrapper to compute soft spheres using a cell list."""
sigma = np.array(sigma, dtype=f32)
epsilon = np.array(epsilon, dtype=f32)
alpha = np.array(alpha, dtype=f32)
if species is None:
if sigma.shape or epsilon.shape or alpha.shape:
raise ValueError(
('At the moment per-particle (as opposed to per-species) parameters are'
' not supported using cell lists. Please open a feature request!'))
energy_fn = smap.cartesian_product(
soft_sphere,
space.canonicalize_displacement_or_metric(displacement_or_metric),
sigma=sigma,
epsilon=epsilon,
alpha=alpha)
return smap.cell_list(energy_fn, box_size, np.max(sigma), R_example, species)
def lennard_jones_cell_list(
displacement_or_metric,
box_size,
R_example,
species=None,
sigma=1.0,
epsilon=1.0,
alpha=2.0,
r_onset=2.0,
r_cutoff=2.5):
"""Convenience wrapper to compute soft spheres using a cell list."""
sigma = np.array(sigma, dtype=f32)
epsilon = np.array(epsilon, dtype=f32)
r_onset = f32(r_onset * np.max(sigma))
r_cutoff = f32(r_cutoff * np.max(sigma))
if species is None:
if sigma.shape or epsilon.shape:
raise ValueError(
('At the moment per-particle (as opposed to per-species) parameters are'
' not supported using cell lists. Please open a feature request!'))
energy_fn = smap.cartesian_product(
multiplicative_isotropic_cutoff(lennard_jones, r_onset, r_cutoff),
space.canonicalize_displacement_or_metric(displacement_or_metric),
sigma=sigma,
epsilon=epsilon)
return smap.cell_list(energy_fn, box_size, r_cutoff, R_example, species)
def test_cell_list_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)
R = box_size * random.uniform(
key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
cell_list_energy = smap.cartesian_product(
energy.soft_sphere, space.metric(displacement))
cell_list_energy = \
jit(smap.cell_list(cell_list_energy, box_size, cell_size, R))
self.assertAllClose(
np.array(energy_fn(R), dtype=dtype), cell_list_energy(R), True)
def test_pair_cell_list_energy(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(9.0)
cell_size = f32(1.0)
displacement, _ = space.periodic(box_size)
metric = space.metric(displacement)
exact_energy_fn = energy.soft_sphere_pair(displacement)
energy_fn = smap.cartesian_product(energy.soft_sphere, metric)
R = box_size * random.uniform(
key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
cell_energy_fn = smap.cell_list(energy_fn, box_size, cell_size, R)
self.assertAllClose(
np.array(exact_energy_fn(R), dtype=dtype),
cell_energy_fn(R), True)
def test_cell_list_direct_force_jit(self, spatial_dimension, dtype):
key = random.PRNGKey(1)
box_size = f32(9.0)
cell_size = f32(1.0)
displacement, _ = space.periodic(box_size)
energy_fn = energy.soft_sphere_pair(displacement)
force_fn = quantity.force(energy_fn)
R = box_size * random.uniform(
key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
grid_energy_fn = smap.cartesian_product(
energy.soft_sphere, space.metric(displacement))
grid_force_fn = quantity.force(grid_energy_fn)
grid_force_fn = jit(smap.cell_list(grid_force_fn, box_size, cell_size, R))
self.assertAllClose(
np.array(force_fn(R), dtype=dtype), grid_force_fn(R), True)
def test_cell_list_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)
force_fn = quantity.force(energy_fn)
R = box_size * random.uniform(
key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
cell_energy_fn = smap.cartesian_product(
energy.soft_sphere, space.metric(displacement))
cell_force_fn = quantity.force(cell_energy_fn)
cell_force_fn = smap.cell_list(cell_force_fn, box_size, cell_size, R)
df = np.sum((force_fn(R) - cell_force_fn(R)) ** 2, axis=1)
self.assertAllClose(
np.array(force_fn(R), dtype=dtype), cell_force_fn(R), True)