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)