def test_cosine_angles(self, dtype):
displacement, _ = space.free()
displacement = space.map_product(displacement)
R = np.array([[0, 0], [0, 1], [1, 1]], dtype=dtype)
dR = displacement(R, R)
cangles = quantity.cosine_angles(dR)
c45 = 1 / np.sqrt(2)
true_cangles = np.array([[[0, 0, 0], [0, 1, c45], [0, c45, 1]],
[[1, 0, 0], [0, 0, 0], [0, 0, 1]],
[[1, c45, 0], [c45, 1, 0], [0, 0, 0]]],
dtype=dtype)
self.assertAllClose(cangles, true_cangles)
def test_cosine_angles_neighbors(self, dtype):
displacement, _ = space.free()
displacement = vmap(vmap(displacement, (None, 0)), 0)
R = np.array([[0, 0], [0, 1], [1, 1]], dtype=dtype)
R_neigh = np.array(
[[[0, 1], [1, 1]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]],
dtype=dtype)
dR = displacement(R, R_neigh)
cangles = quantity.cosine_angles(dR)
c45 = 1 / np.sqrt(2)
true_cangles = np.array(
[[[1, c45], [c45, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]],
dtype=dtype)
self.assertAllClose(cangles, true_cangles)
