def test_pair_neighbor_list_scalar_params_matrix(self, spatial_dimension,
dtype):
key = random.PRNGKey(0)
def truncated_square(dr, sigma):
return np.where(dr < sigma, dr**2, f32(0.))
N = NEIGHBOR_LIST_PARTICLE_COUNT
box_size = 2. * N**(1. / spatial_dimension)
key, split = random.split(key)
disp, _ = space.periodic(box_size)
d = space.metric(disp)
neighbor_square = jit(smap.pair_neighbor_list(truncated_square, d))
mapped_square = jit(smap.pair(truncated_square, d))
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = box_size * random.uniform(split, (N, spatial_dimension),
dtype=dtype)
sigma = random.uniform(key, (N, N), minval=0.5, maxval=1.5)
sigma = 0.5 * (sigma + sigma.T)
neighbor_fn = partition.neighbor_list(disp, box_size,
np.max(sigma), 0.)
nbrs = neighbor_fn(R)
self.assertAllClose(mapped_square(R, sigma=sigma),
neighbor_square(R, nbrs, sigma=sigma))
def test_pair_dynamic_species_vector(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
square = lambda dr, param=1.0: param * np.sum(dr**2, axis=2)
params = f32(np.array([[1.0, 2.0], [2.0, 3.0]]))
key, split = random.split(key)
species = random.randint(split, (PARTICLE_COUNT, ), 0, 2)
disp, _ = space.free()
mapped_square = smap.pair(square,
disp,
species=quantity.Dynamic,
param=params)
disp = vmap(vmap(disp, (0, None), 0), (None, 0), 0)
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = random.uniform(split, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
total = 0.0
for i in range(2):
for j in range(2):
param = params[i, j]
R_1 = R[species == i]
R_2 = R[species == j]
total = total + 0.5 * np.sum(square(disp(R_1, R_2), param))
self.assertAllClose(mapped_square(R, species, 2),
np.array(total, dtype=dtype), True)
def test_pair_neighbor_list_force_scalar_diverging_potential(
self, spatial_dimension, dtype):
key = random.PRNGKey(0)
def potential(dr, sigma):
return np.where(dr < sigma, dr**-6, f32(0.))
N = NEIGHBOR_LIST_PARTICLE_COUNT
box_size = 4. * N**(1. / spatial_dimension)
key, split = random.split(key)
disp, _ = space.periodic(box_size)
d = space.metric(disp)
neighbor_square = jit(
quantity.force(smap.pair_neighbor_list(potential, d)))
mapped_square = jit(quantity.force(smap.pair(potential, d)))
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = box_size * random.uniform(split, (N, spatial_dimension),
dtype=dtype)
sigma = random.uniform(key, (), minval=0.5, maxval=4.5)
neighbor_fn = partition.neighbor_list(disp, box_size, sigma, 0.0)
nbrs = neighbor_fn(R)
self.assertAllClose(mapped_square(R, sigma=sigma),
neighbor_square(R, nbrs, sigma=sigma))
def test_pair_neighbor_list_scalar_params_species(self, spatial_dimension,
dtype):
key = random.PRNGKey(0)
def truncated_square(dr, sigma):
return np.where(dr < sigma, dr**2, f32(0.))
tol = 2e-6 if dtype == np.float32 else None
N = NEIGHBOR_LIST_PARTICLE_COUNT
box_size = 2. * N**(1. / spatial_dimension)
species = np.zeros((N, ), np.int32)
species = np.where(np.arange(N) > N / 3, 1, species)
species = np.where(np.arange(N) > 2 * N / 3, 2, species)
key, split = random.split(key)
disp, _ = space.periodic(box_size)
d = space.metric(disp)
neighbor_square = jit(
smap.pair_neighbor_list(truncated_square, d, species=species))
mapped_square = jit(smap.pair(truncated_square, d, species=species))
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = box_size * random.uniform(split, (N, spatial_dimension),
dtype=dtype)
sigma = random.uniform(key, (3, 3), minval=0.5, maxval=1.5)
sigma = 0.5 * (sigma + sigma.T)
neighbor_fn = jit(
partition.neighbor_list(disp, box_size, np.max(sigma), R))
idx = neighbor_fn(R)
self.assertAllClose(mapped_square(R, sigma=sigma),
neighbor_square(R, idx, sigma=sigma), True,
tol, tol)
def test_pair_neighbor_list_scalar(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
def truncated_square(dr, sigma):
return np.where(dr < sigma, dr**2, f32(0.))
tol = 2e-10 if dtype == np.float32 else None
N = NEIGHBOR_LIST_PARTICLE_COUNT
box_size = 4. * N**(1. / spatial_dimension)
key, split = random.split(key)
disp, _ = space.periodic(box_size)
d = space.metric(disp)
neighbor_square = jit(smap.pair_neighbor_list(truncated_square, d))
mapped_square = jit(smap.pair(truncated_square, d))
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = box_size * random.uniform(split, (N, spatial_dimension),
dtype=dtype)
sigma = random.uniform(key, (), minval=0.5, maxval=2.5)
neighbor_fn = jit(partition.neighbor_list(disp, box_size, sigma,
R))
idx = neighbor_fn(R)
self.assertAllClose(mapped_square(R, sigma=sigma),
neighbor_square(R, idx, sigma=sigma), True,
tol, tol)
def test_pair_neighbor_list_scalar_nonadditive(self, spatial_dimension,
dtype, format):
key = random.PRNGKey(0)
def truncated_square(dR, sigma):
dr = space.distance(dR)
return np.where(dr < sigma, dr**2, f32(0.))
N = PARTICLE_COUNT
box_size = 2. * N**(1. / spatial_dimension)
key, split = random.split(key)
disp, _ = space.periodic(box_size)
neighbor_square = jit(
smap.pair_neighbor_list(truncated_square,
disp,
sigma=lambda x, y: x * y))
mapped_square = jit(smap.pair(truncated_square, disp, sigma=1.0))
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = box_size * random.uniform(split, (N, spatial_dimension),
dtype=dtype)
sigma = random.uniform(key, (N, ), minval=0.5, maxval=1.5)
sigma_pair = sigma[:, None] * sigma[None, :]
neighbor_fn = partition.neighbor_list(disp,
box_size,
np.max(sigma)**2,
0.,
format=format)
nbrs = neighbor_fn.allocate(R)
self.assertAllClose(mapped_square(R, sigma=sigma_pair),
neighbor_square(R, nbrs, sigma=sigma))
def test_pair_dynamic_species_scalar(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
square = lambda dr, param=1.0: param * dr**2
params = f32(np.array([[1.0, 2.0], [2.0, 3.0]]))
key, split = random.split(key)
species = random.randint(split, (PARTICLE_COUNT, ), 0, 2)
displacement, _ = space.free()
metric = lambda Ra, Rb, **kwargs: \
np.sum(displacement(Ra, Rb, **kwargs) ** 2, axis=-1)
mapped_square = smap.pair(square,
metric,
species=quantity.Dynamic,
param=params)
metric = space.map_product(metric)
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = random.uniform(split, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
total = 0.0
for i in range(2):
for j in range(2):
param = params[i, j]
R_1 = R[species == i]
R_2 = R[species == j]
total = total + 0.5 * np.sum(
square(metric(R_1, R_2), param))
self.assertAllClose(mapped_square(R, species, 2),
np.array(total, dtype=dtype), True)
def test_pair_neighbor_list_vector(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
def truncated_square(dR, sigma):
dr = np.reshape(space.distance(dR), dR.shape[:-1] + (1, ))
return np.where(dr < sigma, dR**2, f32(0.))
tol = 5e-6 if dtype == np.float32 else 1e-14
N = PARTICLE_COUNT
box_size = 2. * N**(1. / spatial_dimension)
key, split = random.split(key)
disp, _ = space.periodic(box_size)
neighbor_square = jit(
smap.pair_neighbor_list(truncated_square, disp, reduce_axis=(1, )))
mapped_square = jit(
smap.pair(truncated_square, disp, reduce_axis=(1, )))
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = box_size * random.uniform(split, (N, spatial_dimension),
dtype=dtype)
sigma = random.uniform(key, (), minval=0.5, maxval=1.5)
neighbor_fn = jit(partition.neighbor_list(disp, box_size, sigma,
R))
idx = neighbor_fn(R)
self.assertAllClose(mapped_square(R, sigma=sigma),
neighbor_square(R, idx, sigma=sigma), True,
tol, tol)
def test_pair_scalar_dummy_arg(self, spatial_dimension, dtype):
key = random.PRNGKey(0)
square = lambda dr, param=f32(1.0), **unused_kwargs: param * dr**2
key, split = random.split(key)
R = random.normal(key, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
displacement, shift = space.free()
mapped = smap.pair(square, space.metric(displacement))
mapped(R, t=f32(0))
def soft_sphere_pair(
displacement_or_metric, species=None, sigma=1.0, epsilon=1.0, alpha=2.0):
"""Convenience wrapper to compute soft sphere energy over a system."""
sigma = np.array(sigma, dtype=f32)
epsilon = np.array(epsilon, dtype=f32)
alpha = np.array(alpha, dtype=f32)
return smap.pair(
soft_sphere,
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
sigma=sigma,
epsilon=epsilon,
alpha=alpha)
def morse_pair(
displacement_or_metric,
species=None, sigma=1.0, epsilon=5.0, alpha=5.0, r_onset=2.0, r_cutoff=2.5):
"""Convenience wrapper to compute Morse energy over a system."""
sigma = np.array(sigma, dtype=f32)
epsilon = np.array(epsilon, dtype=f32)
alpha = np.array(alpha, dtype=f32)
return smap.pair(
multiplicative_isotropic_cutoff(morse, r_onset, r_cutoff),
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
sigma=sigma,
epsilon=epsilon,
alpha=alpha)
def lennard_jones_pair(
displacement_or_metric,
species=None, sigma=1.0, epsilon=1.0, r_onset=2.0, r_cutoff=2.5):
"""Convenience wrapper to compute Lennard-Jones energy over a system."""
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))
return smap.pair(
multiplicative_isotropic_cutoff(lennard_jones, r_onset, r_cutoff),
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
sigma=sigma,
epsilon=epsilon)
def test_pair_no_species_vector(self, spatial_dimension, dtype):
square = lambda dr: np.sum(dr ** 2, axis=2)
disp, _ = space.free()
mapped_square = smap.pair(square, disp)
disp = space.map_product(disp)
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = random.uniform(
split, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
mapped_ref = np.array(0.5 * np.sum(square(disp(R, R))), dtype=dtype)
self.assertAllClose(mapped_square(R), mapped_ref)
def soft_sphere_pair(displacement_or_metric: DisplacementOrMetricFn,
species: Array = None,
sigma: Array = 1.0,
epsilon: Array = 1.0,
alpha: Array = 2.0,
per_particle: bool = False):
"""Convenience wrapper to compute soft sphere energy over a system."""
sigma = np.array(sigma, dtype=f32)
epsilon = np.array(epsilon, dtype=f32)
alpha = np.array(alpha, dtype=f32)
return smap.pair(
soft_sphere,
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
sigma=sigma,
epsilon=epsilon,
alpha=alpha,
reduce_axis=(1, ) if per_particle else None)
def test_pair_no_species_scalar(self, spatial_dimension, dtype):
square = lambda dr: dr**2
displacement, _ = space.free()
metric = lambda Ra, Rb, **kwargs: \
np.sum(displacement(Ra, Rb, **kwargs) ** 2, axis=-1)
mapped_square = smap.pair(square, metric)
metric = space.map_product(metric)
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = random.uniform(split, (PARTICLE_COUNT, spatial_dimension),
dtype=dtype)
self.assertAllClose(
mapped_square(R),
np.array(0.5 * np.sum(square(metric(R, R))), dtype=dtype))
def bks_pair(displacement_or_metric, species, Q_sq, exp_coeff, exp_decay,
attractive_coeff, repulsive_coeff, coulomb_alpha, cutoff):
Q_sq = np.array(Q_sq, f32)
exp_coeff = np.array(exp_coeff, f32)
exp_decay = np.array(exp_decay, f32)
attractive_coeff = np.array(attractive_coeff, f32)
repulsive_coeff = np.array(repulsive_coeff, f32)
return smap.pair(bks,
displacement_or_metric,
species=species,
Q_sq=Q_sq,
exp_coeff=exp_coeff,
exp_decay=exp_decay,
attractive_coeff=attractive_coeff,
repulsive_coeff=repulsive_coeff,
coulomb_alpha=coulomb_alpha,
cutoff=cutoff)
def test_stress_lammps_periodic_general(self, dim, dtype):
key = random.PRNGKey(0)
N = 64
(box, R, V), (E, C) = test_util.load_lammps_stress_data(dtype)
displacement_fn, _ = space.periodic_general(box)
energy_fn = smap.pair(
lambda dr, **kwargs: jnp.where(dr < f32(2.5),
energy.lennard_jones(dr), f32(0.0)),
space.canonicalize_displacement_or_metric(displacement_fn))
ad_stress = quantity.stress(energy_fn, R, box, velocity=V)
tol = 5e-5
self.assertAllClose(energy_fn(R) / len(R), E, atol=tol, rtol=tol)
self.assertAllClose(C, ad_stress, atol=tol, rtol=tol)
def test_pair_no_species_vector_nonadditive(self, spatial_dimension, dtype):
square = lambda dr, params: params * np.sum(dr ** 2, axis=2)
disp, _ = space.free()
mapped_square = smap.pair(square, disp, params=lambda x, y: x * y)
disp = space.map_product(disp)
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, R_key, params_key = random.split(key, 3)
R = random.uniform(
R_key, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
params = random.uniform(
params_key, (PARTICLE_COUNT,), dtype=dtype, minval=0.1, maxval=1.5)
pp_params = params[None, :] * params[:, None]
mapped_ref = np.array(0.5 * np.sum(square(disp(R, R), pp_params)),
dtype=dtype)
self.assertAllClose(mapped_square(R, params=params), mapped_ref)
def lennard_jones_pair(displacement_or_metric: DisplacementOrMetricFn,
species: Array = None,
sigma: Array = 1.0,
epsilon: Array = 1.0,
r_onset: Array = 2.0,
r_cutoff: Array = 2.5,
per_particle: bool = False) -> Callable[[Array], Array]:
"""Convenience wrapper to compute Lennard-Jones energy over a system."""
sigma = np.array(sigma, dtype=f32)
epsilon = np.array(epsilon, dtype=f32)
r_onset = r_onset * np.max(sigma)
r_cutoff = r_cutoff * np.max(sigma)
return smap.pair(
multiplicative_isotropic_cutoff(lennard_jones, r_onset, r_cutoff),
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
sigma=sigma,
epsilon=epsilon,
reduce_axis=(1, ) if per_particle else None)
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)
def soft_sphere_pair(displacement_or_metric: DisplacementOrMetricFn,
species: Array = None,
sigma: Array = 1.0,
epsilon: Array = 1.0,
alpha: Array = 2.0,
per_particle: bool = False):
"""Convenience wrapper to compute soft sphere energy over a system."""
sigma = maybe_downcast(sigma)
epsilon = maybe_downcast(epsilon)
alpha = maybe_downcast(alpha)
return smap.pair(
soft_sphere,
space.canonicalize_displacement_or_metric(displacement_or_metric),
ignore_unused_parameters=True,
species=species,
sigma=sigma,
epsilon=epsilon,
alpha=alpha,
reduce_axis=(1, ) if per_particle else None)
def test_pair_no_species_scalar_dynamic(self, spatial_dimension, dtype):
square = lambda dr, epsilon: epsilon * dr ** 2
displacement, _ = space.free()
metric = lambda Ra, Rb, **kwargs: \
np.sum(displacement(Ra, Rb, **kwargs) ** 2, axis=-1)
mapped_square = smap.pair(square, metric, epsilon=1.0)
metric = space.map_product(metric)
key = random.PRNGKey(0)
for _ in range(STOCHASTIC_SAMPLES):
key, split1, split2 = random.split(key, 3)
R = random.uniform(
split1, (PARTICLE_COUNT, spatial_dimension), dtype=dtype)
epsilon = random.uniform(split2, (PARTICLE_COUNT,), dtype=dtype)
mat_epsilon = 0.5 * (epsilon[:, np.newaxis] + epsilon[np.newaxis, :])
self.assertAllClose(
mapped_square(R, epsilon=epsilon),
np.array(0.5 * np.sum(
square(metric(R, R), mat_epsilon)), dtype=dtype))
def morse_pair(displacement_or_metric: DisplacementOrMetricFn,
species: Array = None,
sigma: Array = 1.0,
epsilon: Array = 5.0,
alpha: Array = 5.0,
r_onset: float = 2.0,
r_cutoff: float = 2.5,
per_particle: bool = False) -> Callable[[Array], Array]:
"""Convenience wrapper to compute Morse energy over a system."""
sigma = np.array(sigma, dtype=f32)
epsilon = np.array(epsilon, dtype=f32)
alpha = np.array(alpha, dtype=f32)
return smap.pair(
multiplicative_isotropic_cutoff(morse, r_onset, r_cutoff),
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
sigma=sigma,
epsilon=epsilon,
alpha=alpha,
reduce_axis=(1, ) if per_particle else None)
def bks_pair(displacement_or_metric: DisplacementOrMetricFn, species: Array,
Q_sq: Array, exp_coeff: Array, exp_decay: Array,
attractive_coeff: Array, repulsive_coeff: Array,
coulomb_alpha: Array, cutoff: float) -> Callable[[Array], Array]:
"""Convenience wrapper to compute BKS energy over a system."""
Q_sq = np.array(Q_sq, f32)
exp_coeff = np.array(exp_coeff, f32)
exp_decay = np.array(exp_decay, f32)
attractive_coeff = np.array(attractive_coeff, f32)
repulsive_coeff = np.array(repulsive_coeff, f32)
return smap.pair(bks,
displacement_or_metric,
species=species,
Q_sq=Q_sq,
exp_coeff=exp_coeff,
exp_decay=exp_decay,
attractive_coeff=attractive_coeff,
repulsive_coeff=repulsive_coeff,
coulomb_alpha=coulomb_alpha,
cutoff=cutoff)
def bks_pair(displacement_or_metric: DisplacementOrMetricFn, species: Array,
Q_sq: Array, exp_coeff: Array, exp_decay: Array,
attractive_coeff: Array, repulsive_coeff: Array,
coulomb_alpha: Array, cutoff: float) -> Callable[[Array], Array]:
"""Convenience wrapper to compute BKS energy over a system."""
Q_sq = maybe_downcast(Q_sq)
exp_coeff = maybe_downcast(exp_coeff)
exp_decay = maybe_downcast(exp_decay)
attractive_coeff = maybe_downcast(attractive_coeff)
repulsive_coeff = maybe_downcast(repulsive_coeff)
return smap.pair(bks,
displacement_or_metric,
species=species,
ignore_unused_parameters=True,
Q_sq=Q_sq,
exp_coeff=exp_coeff,
exp_decay=exp_decay,
attractive_coeff=attractive_coeff,
repulsive_coeff=repulsive_coeff,
coulomb_alpha=coulomb_alpha,
cutoff=cutoff)
def harmonic_morse_pair(displacement_or_metric,
species=None,
D0=5.0,
alpha=10.0,
r0=1.0,
k=50.0):
"""The harmonic morse function over all pairs of particles in a system."""
# Initialize various parameters of the harmonic morse function
D0 = jnp.array(D0, dtype=jnp.float32)
alpha = jnp.array(alpha, dtype=jnp.float32)
r0 = jnp.array(r0, dtype=jnp.float32)
k = jnp.array(k, dtype=jnp.float32)
# Pass the harmonic morse function defined above along with its parameters and a
# displacement/metric function.
return smap.pair(
harmonic_morse,
space.canonicalize_displacement_or_metric(displacement_or_metric),
species=species,
D0=D0,
alpha=alpha,
r0=r0,
k=k)
def bks_pair(displacement_or_metric: DisplacementOrMetricFn,
species: Array,
Q_sq: Array,
exp_coeff: Array,
exp_decay: Array,
attractive_coeff: Array,
repulsive_coeff: Array,
coulomb_alpha: Array,
cutoff: float) -> Callable[[Array], Array]:
Q_sq = np.array(Q_sq, f32)
exp_coeff = np.array(exp_coeff, f32)
exp_decay = np.array(exp_decay, f32)
attractive_coeff = np.array(attractive_coeff, f32)
repulsive_coeff = np.array(repulsive_coeff, f32)
return smap.pair(bks, displacement_or_metric,
species=species,
Q_sq=Q_sq,
exp_coeff=exp_coeff,
exp_decay=exp_decay,
attractive_coeff=attractive_coeff,
repulsive_coeff=repulsive_coeff,
coulomb_alpha=coulomb_alpha,
cutoff=cutoff)
def test_pair_neighbor_list_vector(self, spatial_dimension, dtype, format):
if format is partition.OrderedSparse:
self.skipTest('Vector valued pair_neighbor_list not supported.')
key = random.PRNGKey(0)
def truncated_square(dR, sigma):
dr = np.reshape(space.distance(dR), dR.shape[:-1] + (1, ))
return np.where(dr < sigma, dR**2, f32(0.))
N = PARTICLE_COUNT
box_size = 2. * N**(1. / spatial_dimension)
key, split = random.split(key)
disp, _ = space.periodic(box_size)
neighbor_square = jit(
smap.pair_neighbor_list(truncated_square,
disp,
sigma=1.0,
reduce_axis=(1, )))
mapped_square = jit(
smap.pair(truncated_square, disp, sigma=1.0, reduce_axis=(1, )))
for _ in range(STOCHASTIC_SAMPLES):
key, split = random.split(key)
R = box_size * random.uniform(split, (N, spatial_dimension),
dtype=dtype)
sigma = random.uniform(key, (), minval=0.5, maxval=1.5)
neighbor_fn = partition.neighbor_list(disp,
box_size,
sigma,
0.,
format=format)
nbrs = neighbor_fn.allocate(R)
self.assertAllClose(mapped_square(R, sigma=sigma),
neighbor_square(R, nbrs, sigma=sigma))
