def test_pair_correlation_neighbor_list_species(self, dim, dtype, format):
if format is partition.OrderedSparse:
self.skipTest('OrderedSparse not supported for pair correlation '
'function.')
N = 100
L = 10.
displacement, _ = space.periodic(L)
R = random.uniform(random.PRNGKey(0), (N, dim), dtype=dtype)
species = np.where(np.arange(N) < N // 2, 0, 1)
rs = np.linspace(0, 2, 60, dtype=dtype)
g = quantity.pair_correlation(displacement, rs, f32(0.1), species)
nbr_fn, g_neigh = quantity.pair_correlation_neighbor_list(
displacement, L, rs, f32(0.1), species, format=format)
nbrs = nbr_fn.allocate(R)
g_0, g_1 = g(R)
g_0 = np.mean(g_0, axis=0)
g_1 = np.mean(g_1, axis=0)
g_0_neigh, g_1_neigh = g_neigh(R, neighbor=nbrs)
g_0_neigh = np.mean(g_0_neigh, axis=0)
g_1_neigh = np.mean(g_1_neigh, axis=0)
self.assertAllClose(g_0, g_0_neigh)
self.assertAllClose(g_1, g_1_neigh)
def test_pair_correlation_neighbor_list(self, dim, dtype):
N = 100
L = 10.
displacement, _ = space.periodic(L)
R = random.uniform(random.PRNGKey(0), (N, dim), dtype=dtype)
rs = np.linspace(0, 2, 60, dtype=dtype)
g = quantity.pair_correlation(displacement, rs, f32(0.1))
nbr_fn, g_neigh = quantity.pair_correlation_neighbor_list(
displacement, L, rs, f32(0.1))
nbrs = nbr_fn(R)
g_0 = g(R)
g_0 = np.mean(g_0, axis=0)
g_0_neigh = g_neigh(R, neighbor=nbrs)
g_0_neigh = np.mean(g_0_neigh, axis=0)
self.assertAllClose(g_0, g_0_neigh)