def test_l2_metric_dist_vectorization(self): """ Test the vectorization of dist. """ n_samples = self.n_discretized_curves curves_ab = self.l2_metric_s2.geodesic(self.curve_a, self.curve_b) curves_bc = self.l2_metric_s2.geodesic(self.curve_b, self.curve_c) curves_ab = curves_ab(self.times) curves_bc = curves_bc(self.times) result = self.l2_metric_s2.dist(curves_ab, curves_bc) self.assertAllClose(gs.shape(result), (n_samples, 1))
def test_log_vectorization(self): n_samples = self.n_samples n = self.n p = self.p one_point = self.space.random_uniform() one_base_point = self.space.random_uniform() n_points = self.space.random_uniform(n_samples=n_samples) n_base_points = self.space.random_uniform(n_samples=n_samples) result = self.metric.log(one_point, one_base_point) self.assertAllClose(gs.shape(result), (1, n, p)) result = self.metric.log(n_points, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, n, p)) result = self.metric.log(one_point, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, n, p)) result = self.metric.log(n_points, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, n, p))
def test_log_vectorization(self): n_samples = self.n_samples dim = self.dimension one_point = gs.array([0.0, 1.0]) one_base_point = gs.array([2.0, 10.0]) n_points = gs.array([[2.0, 1.0], [-2.0, -4.0], [-5.0, 1.0]]) n_base_points = gs.array([[2.0, 10.0], [8.0, -1.0], [-3.0, 6.0]]) result = self.metric.log(one_point, one_base_point) expected = one_point - one_base_point self.assertAllClose(result, expected) result = self.metric.log(n_points, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, dim)) result = self.metric.log(one_point, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, dim)) result = self.metric.log(n_points, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, dim))
def test_exp_vectorization(self): n_samples = self.n_samples one_base_point = self.space.random_uniform(n_samples=1) n_base_point = self.space.random_uniform(n_samples=n_samples) n_tangent_vec_same_base = self.space.random_tangent_vec_uniform( n_samples=n_samples, base_point=one_base_point) n_tangent_vec = self.space.random_tangent_vec_uniform( n_samples=n_samples, base_point=n_base_point) # Test with the 1 base_point, and several different tangent_vecs result = self.metric.exp(n_tangent_vec_same_base, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.space.n, self.space.n)) # Test with the same number of base_points and tangent_vecs result = self.metric.exp(n_tangent_vec, n_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.space.n, self.space.n))
def test_mean_matrices_shape(self): m, n = (2, 2) point = gs.array([[1.0, 4.0], [2.0, 3.0]]) metric = MatricesMetric(m, n) mean = FrechetMean(metric=metric, point_type="matrix") points = [point, point, point] mean.fit(points) result = mean.estimate_ self.assertAllClose(gs.shape(result), (m, n))
def test_compose_vectorization(self): n_samples = self.n_samples n_points_a = self.group.random_point(n_samples=n_samples) n_points_b = self.group.random_point(n_samples=n_samples) one_point = self.group.random_point(n_samples=1) result = self.group.compose(one_point, n_points_a) self.assertAllClose(gs.shape(result), (n_samples, *self.group.get_point_type_shape())) result = self.group.compose(n_points_a, one_point) if not geomstats.tests.tf_backend(): self.assertAllClose( gs.shape(result), (n_samples, *self.group.get_point_type_shape())) result = self.group.compose(n_points_a, n_points_b) self.assertAllClose( gs.shape(result), (n_samples, *self.group.get_point_type_shape()))
def test_srv_metric_pointwise_inner_product(self): curves_ab = self.l2_metric_s2.geodesic(self.curve_a, self.curve_b) curves_bc = self.l2_metric_s2.geodesic(self.curve_b, self.curve_c) curves_ab = curves_ab(self.times) curves_bc = curves_bc(self.times) tangent_vecs = self.l2_metric_s2.log(point=curves_bc, base_point=curves_ab) result = self.srv_metric_r3.pointwise_inner_product( tangent_vec_a=tangent_vecs, tangent_vec_b=tangent_vecs, base_curve=curves_ab) expected_shape = (self.n_discretized_curves, self.n_sampling_points) self.assertAllClose(gs.shape(result), expected_shape) result = self.srv_metric_r3.pointwise_inner_product( tangent_vec_a=tangent_vecs[0], tangent_vec_b=tangent_vecs[0], base_curve=curves_ab[0]) expected_shape = (self.n_sampling_points, ) self.assertAllClose(gs.shape(result), expected_shape)
def test_l2_metric_dist_vectorization(self): """Test the vectorization of dist.""" n_samples = self.n_landmark_sets landmarks_ab = self.l2_metric_s2.geodesic(self.landmarks_a, self.landmarks_b) landmarks_bc = self.l2_metric_s2.geodesic(self.landmarks_b, self.landmarks_c) landmarks_ab = landmarks_ab(self.times) landmarks_bc = landmarks_bc(self.times) result = self.l2_metric_s2.dist(landmarks_ab, landmarks_bc) self.assertAllClose(gs.shape(result), (n_samples, ))
def test_estimate_shape_default_gradient_descent_sphere(self): dim = 5 point_a = gs.array([1., 0., 0., 0., 0.]) point_b = gs.array([0., 1., 0., 0., 0.]) points = gs.array([point_a, point_b]) mean = FrechetMean( metric=self.sphere.metric, method='default', verbose=True) mean.fit(points) result = mean.estimate_ self.assertAllClose(gs.shape(result), (dim,))
def test_squared_norm_vectorization(self): n_samples = self.n_samples n_points = gs.array([ [2., 1.], [-2., -4.], [-5., 1.]]) result = self.metric.squared_norm(n_points) expected = gs.array([5., 20., 26.]) self.assertAllClose(gs.shape(result), (n_samples,)) self.assertAllClose(result, expected)
def test_load_cities(self): """Test that the cities coordinates belong to the sphere.""" sphere = Hypersphere(dim=2) data, _ = data_utils.load_cities() self.assertAllClose(gs.shape(data), (50, 3)) tokyo = data[0] self.assertAllClose(tokyo, gs.array([0.61993792, -0.52479018, 0.58332859])) result = sphere.belongs(data) self.assertTrue(gs.all(result))
def log(self, point, base_point, n_steps=N_STEPS, step='euler'): """Compute logarithm map associated to the affine connection. Solve the boundary value problem associated to the geodesic equation using the Christoffel symbols and conjugate gradient descent. Parameters ---------- point : array-like, shape=[n_samples, dim] Point on the manifold. base_point : array-like, shape=[n_samples, dim] Point on the manifold. n_steps : int The number of discrete time steps to take in the integration. step : str, {'euler', 'rk4'} The numerical scheme to use for integration. Returns ------- tangent_vec : array-like, shape=[n_samples, dim] Tangent vector at the base point. """ point = gs.to_ndarray(point, to_ndim=2) base_point = gs.to_ndarray(base_point, to_ndim=2) n_samples = len(base_point) def objective(velocity): """Define the objective function.""" velocity = velocity.reshape(base_point.shape) delta = self.exp(velocity, base_point, n_steps, step) - point loss = 1. / self.dim * gs.sum(delta**2, axis=1) return 1. / n_samples * gs.sum(loss) objective_grad = autograd.elementwise_grad(objective) def objective_with_grad(velocity): """Create helpful objective func wrapper for autograd comp.""" return objective(velocity), objective_grad(velocity) tangent_vec = gs.random.rand(gs.sum(gs.shape(base_point))) res = minimize(objective_with_grad, tangent_vec, method='L-BFGS-B', jac=True, options={ 'disp': False, 'maxiter': 25 }) tangent_vec = res.x tangent_vec = gs.reshape(tangent_vec, base_point.shape) return tangent_vec
def test_group_log_vectorization(self): n_samples = self.n_samples one_point = self.group.random_uniform(n_samples=1) one_base_point = self.group.random_uniform(n_samples=1) n_point = self.group.random_uniform(n_samples=n_samples) n_base_point = self.group.random_uniform(n_samples=n_samples) # Test with the 1 base point, and several different points result = self.group.log(n_point, one_base_point) self.assertAllClose( gs.shape(result), (n_samples, self.group.dim)) # Test with the several base point, and 1 point result = self.group.log(one_point, n_base_point) self.assertAllClose( gs.shape(result), (n_samples, self.group.dim)) # Test with the same number n of base point and point result = self.group.log(n_point, n_base_point) self.assertAllClose( gs.shape(result), (n_samples, self.group.dim))
def test_group_exp_vectorization(self): n_samples = self.n_samples one_tangent_vec = self.group.random_uniform(n_samples=1) one_base_point = self.group.random_uniform(n_samples=1) n_tangent_vec = self.group.random_uniform(n_samples=n_samples) n_base_point = self.group.random_uniform(n_samples=n_samples) # Test with the 1 base point, and n tangent vecs result = self.group.exp(n_tangent_vec, one_base_point) self.assertAllClose( gs.shape(result), (n_samples, self.group.dim)) # Test with the several base point, and one tangent vec result = self.group.exp(one_tangent_vec, n_base_point) self.assertAllClose( gs.shape(result), (n_samples, self.group.dim)) # Test with the same number n of base point and n tangent vec result = self.group.exp(n_tangent_vec, n_base_point) self.assertAllClose( gs.shape(result), (n_samples, self.group.dim))
def test_l2_metric_inner_product_vectorization(self): """Test the vectorization inner_product.""" n_samples = self.n_discretized_curves curves_ab = self.l2_metric_s2.geodesic(self.curve_a, self.curve_b) curves_bc = self.l2_metric_s2.geodesic(self.curve_b, self.curve_c) curves_ab = curves_ab(self.times) curves_bc = curves_bc(self.times) tangent_vecs = self.l2_metric_s2.log(point=curves_bc, base_point=curves_ab) result = self.l2_metric_s2.inner_product(tangent_vecs, tangent_vecs, curves_ab) self.assertAllClose(gs.shape(result), (n_samples,))
def test_exp_vectorization(self): n_samples = self.n_samples dim = self.dimension one_tangent_vec = gs.array([0., 1.]) one_base_point = gs.array([2., 10.]) n_tangent_vecs = gs.array([[2., 1.], [-2., -4.], [-5., 1.]]) n_base_points = gs.array([[2., 10.], [8., -1.], [-3., 6.]]) result = self.metric.exp(one_tangent_vec, one_base_point) expected = one_tangent_vec + one_base_point self.assertAllClose(result, expected) result = self.metric.exp(n_tangent_vecs, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, dim)) result = self.metric.exp(one_tangent_vec, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, dim)) result = self.metric.exp(n_tangent_vecs, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, dim))
def test_exp_vectorization_single_samples(self): dim = self.dimension + 1 one_vec = self.space.random_uniform() one_base_point = self.space.random_uniform() one_tangent_vec = self.space.to_tangent(one_vec, base_point=one_base_point) result = self.metric.exp(one_tangent_vec, one_base_point) self.assertAllClose(gs.shape(result), (dim, )) one_base_point = gs.to_ndarray(one_base_point, to_ndim=2) result = self.metric.exp(one_tangent_vec, one_base_point) self.assertAllClose(gs.shape(result), (1, dim)) one_tangent_vec = gs.to_ndarray(one_tangent_vec, to_ndim=2) result = self.metric.exp(one_tangent_vec, one_base_point) self.assertAllClose(gs.shape(result), (1, dim)) one_base_point = self.space.random_uniform() result = self.metric.exp(one_tangent_vec, one_base_point) self.assertAllClose(gs.shape(result), (1, dim))
def test_norm_vectorization_n_samples(self): n_samples = self.n_samples n_points = gs.array([ [2., 1.], [-2., -4.], [-5., 1.]]) result = self.metric.norm(n_points) expected = gs.array([2.2360679775, 4.472135955, 5.09901951359]) self.assertAllClose(gs.shape(result), (n_samples,)) self.assertAllClose(result, expected)
def test_squared_dist_vectorization(self): n_samples = self.n_samples one_point_a = gs.array([0., 1.]) one_point_b = gs.array([2., 10.]) n_points_a = gs.array([[2., 1.], [-2., -4.], [-5., 1.]]) n_points_b = gs.array([[2., 10.], [8., -1.], [-3., 6.]]) result = self.metric.squared_dist(one_point_a, one_point_b) vec = one_point_a - one_point_b expected = gs.dot(vec, gs.transpose(vec)) self.assertAllClose(result, expected) result = self.metric.squared_dist(n_points_a, one_point_b) self.assertAllClose(gs.shape(result), (n_samples, )) result = self.metric.squared_dist(one_point_a, n_points_b) self.assertAllClose(gs.shape(result), (n_samples, )) result = self.metric.squared_dist(n_points_a, n_points_b) expected = gs.array([81., 109., 29.]) self.assertAllClose(gs.shape(result), (n_samples, )) self.assertAllClose(result, expected)
def test_exp_vectorization(self): """Test of SPDMetricAffine.exp with power=1 and vectorization.""" n_samples = self.n_samples one_base_point = self.space.random_point(n_samples=1) n_base_point = self.space.random_point(n_samples=n_samples) n_tangent_vec_same_base = self.space.random_tangent_vec( n_samples=n_samples, base_point=one_base_point) n_tangent_vec = self.space.random_tangent_vec(n_samples=n_samples, base_point=n_base_point) metric = self.metric_affine # Test with the 1 base_point, and several different tangent_vecs result = metric.exp(n_tangent_vec_same_base, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.space.n, self.space.n)) # Test with the same number of base_points and tangent_vecs result = metric.exp(n_tangent_vec, n_base_point) self.assertAllClose(gs.shape(result), (n_samples, self.space.n, self.space.n))
def test_retraction_vectorization_shape(self): n_samples = self.n_samples n = self.n p = self.p one_point = self.point_a n_points = gs.tile(gs.to_ndarray(one_point, to_ndim=3), (n_samples, 1, 1)) one_tangent_vec = self.tangent_vector_1 n_tangent_vecs = gs.tile( gs.to_ndarray(self.tangent_vector_2, to_ndim=3), (n_samples, 1, 1)) result = self.metric.retraction(one_tangent_vec, one_point) self.assertAllClose(gs.shape(result), (n, p)) result = self.metric.retraction(n_tangent_vecs, one_point) self.assertAllClose(gs.shape(result), (n_samples, n, p)) result = self.metric.retraction(one_tangent_vec, n_points) self.assertAllClose(gs.shape(result), (n_samples, n, p)) result = self.metric.retraction(n_tangent_vecs, n_points) self.assertAllClose(gs.shape(result), (n_samples, n, p))
def test_lifting_vectorization_shape(self): n_samples = self.n_samples n = self.n p = self.p one_point = self.point_a one_base_point = self.point_b n_points = gs.tile(gs.to_ndarray(self.point_a, to_ndim=3), (n_samples, 1, 1)) n_base_points = gs.tile(gs.to_ndarray(self.point_b, to_ndim=3), (n_samples, 1, 1)) result = self.metric.lifting(one_point, one_base_point) self.assertAllClose(gs.shape(result), (n, p)) result = self.metric.lifting(n_points, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, n, p)) result = self.metric.lifting(one_point, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, n, p)) result = self.metric.lifting(n_points, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, n, p))
def test_exp_vectorization_shape(self): n_samples = self.n_samples n = self.n p = self.p one_base_point = self.point_a one_tangent_vec = self.tangent_vector_1 n_base_points = gs.tile(gs.to_ndarray(self.point_a, to_ndim=3), (n_samples, 1, 1)) n_tangent_vecs = gs.tile( gs.to_ndarray(self.tangent_vector_2, to_ndim=3), (n_samples, 1, 1)) # With single tangent vec and base point result = self.metric.exp(one_tangent_vec, one_base_point) self.assertAllClose(gs.shape(result), (n, p)) # With n_samples tangent vecs and base points result = self.metric.exp(n_tangent_vecs, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, n, p)) result = self.metric.exp(n_tangent_vecs, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, n, p))
def test_log_vectorization(self): dim = self.dimension n_samples = 3 one_point = gs.array([-1., 0.]) one_base_point = gs.array([1.0, 0.]) n_points = gs.array([[-1., 0.], [1., 0.], [2., math.sqrt(3)]]) n_base_points = gs.array([[2., -math.sqrt(3)], [4.0, math.sqrt(15)], [-4.0, math.sqrt(15)]]) result = self.metric.log(one_point, one_base_point) expected = one_point - one_base_point self.assertAllClose(result, expected) result = self.metric.log(n_points, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, dim)) result = self.metric.log(one_point, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, dim)) result = self.metric.log(n_points, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, dim))
def test_exp_vectorization(self): n_samples = self.n_samples n = self.n p = self.p one_base_point = self.point_a n_base_points = gs.tile( gs.to_ndarray(self.point_a, to_ndim=3), (n_samples, 1, 1)) one_tangent_vec = self.tangent_vector_1 result = self.metric.exp(one_tangent_vec, one_base_point) self.assertAllClose(gs.shape(result), (1, n, p)) n_tangent_vecs = gs.tile( gs.to_ndarray(self.tangent_vector_2, to_ndim=3), (n_samples, 1, 1)) result = self.metric.exp(n_tangent_vecs, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, n, p)) result = self.metric.exp(one_tangent_vec, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, n, p))
def test_exp_vectorization(self): dim = self.dimension n_samples = 3 one_tangent_vec = gs.array([-1.0, 0.0]) one_base_point = gs.array([1.0, 0.0]) n_tangent_vecs = gs.array([[-1.0, 0.0], [1.0, 0.0], [2.0, math.sqrt(3)]]) n_base_points = gs.array( [[2.0, -math.sqrt(3)], [4.0, math.sqrt(15)], [-4.0, math.sqrt(15)]] ) result = self.metric.exp(one_tangent_vec, one_base_point) expected = one_tangent_vec + one_base_point self.assertAllClose(result, expected) result = self.metric.exp(n_tangent_vecs, one_base_point) self.assertAllClose(gs.shape(result), (n_samples, dim)) result = self.metric.exp(one_tangent_vec, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, dim)) result = self.metric.exp(n_tangent_vecs, n_base_points) self.assertAllClose(gs.shape(result), (n_samples, dim))
def test_inner_product_vectorization(self): n_samples = 3 one_point_a = gs.array([-1., 0.]) one_point_b = gs.array([1.0, 0.]) n_points_a = gs.array([ [-1., 0.], [1., 0.], [2., math.sqrt(3)]]) n_points_b = gs.array([ [2., -math.sqrt(3)], [4.0, math.sqrt(15)], [-4.0, math.sqrt(15)]]) result = self.metric.inner_product(one_point_a, one_point_b) expected = gs.dot(one_point_a, gs.transpose(one_point_b)) expected -= (2 * one_point_a[self.time_like_dim] * one_point_b[self.time_like_dim]) result_no = self.metric.inner_product(n_points_a, one_point_b) result_on = self.metric.inner_product(one_point_a, n_points_b) result_nn = self.metric.inner_product(n_points_a, n_points_b) self.assertAllClose(result, expected) self.assertAllClose(gs.shape(result_no), (n_samples,)) self.assertAllClose(gs.shape(result_on), (n_samples,)) self.assertAllClose(gs.shape(result_nn), (n_samples,)) expected = np.zeros(n_samples) for i in range(n_samples): expected[i] = gs.dot(n_points_a[i], n_points_b[i]) expected[i] -= (2 * n_points_a[i, self.time_like_dim] * n_points_b[i, self.time_like_dim]) self.assertAllClose(result_nn, expected)
def cov_hessian_covmean_sphere(mean, dataset): """compute dataset covariance, Hessian and covariance of the Fréchet mean" Returns ------- cov: Cov = 1/n \\sum_i=1^n log_mean(x_i) log_mean(x_i)^t hess: H = 1/n \\sum_i=1^n \\partial^2 \\dist^2(mean, x_i) / \\partial mean^2 hinv: inverse Hessian restricted to the tangent space cov_mean_clt: covariance predicted on mean by CLT cov_mean_hc: covariance predicted on mean by high concentration expansion """ n_samples, dim = gs.shape(dataset) assert n_samples > 0, "dataset needs to have at least one data" assert dim > 1, "Dimension of embedding space should be at least 2" assert len(mean) == dim, "Mean should have dimension dim" sphere = Hypersphere(dim - 1) cov = np.zeros((dim, dim), 'float') hess = np.zeros((dim, dim), 'float') idmat = np.identity(dim, 'float') for item in dataset: xy = sphere.metric.log(item, mean) theta = sphere.metric.norm(xy, mean) cov += np.outer(xy, xy) # Numerical issues if theta < 1e-3: a = 1. / 3. + theta ** 2 / 45. + (2. / 945.) * theta ** 4 \ + theta ** 6 / 4725 + (2. / 93555.) * theta ** 8 b = 1. - theta ** 2 / 3. - theta ** 4 / 45. \ - (2. / 945.) * theta ** 6 - theta ** 8 / 4725 else: a = 1. / theta**2 - 1.0 / np.tan(theta) / theta b = theta / np.tan(theta) hess += a * np.outer(xy, xy) + b * (idmat - np.outer(mean, mean)) hess = 2 * hess / n_samples cov = cov / n_samples # ATTN inverse should be restricted to the tangent space because mean is # an eigenvector of Hessian with eigenvalue 0 # Thus, we add a component along mean.mean^t to make eigenvalue 1 before # inverting hinv = np.linalg.inv(hess + np.outer(mean, mean)) # and we remove the component along mean.mean^t afterward hinv -= np.dot(mean, np.matmul(hinv, mean)) * np.outer(mean, mean) # Compute covariance of the mean predicted cov_mean_clt = 4 * np.matmul(hinv, np.matmul(cov, hinv)) / n_samples cov_mean_hc = np.matmul( cov, 2.0 * (idmat - (1.0 - 1.0 / n_samples) / 3.0 * (cov - np.trace(cov) * idmat))) / n_samples return cov, hess, hinv, cov_mean_clt, cov_mean_hc
def test_srv_metric_pointwise_inner_products( self, times, curve_a, curve_b, curve_c, n_discretized_curves, n_sampling_points ): l2_metric_s2 = L2CurvesMetric(ambient_manifold=s2) srv_metric_r3 = SRVMetric(ambient_manifold=r3) curves_ab = l2_metric_s2.geodesic(curve_a, curve_b) curves_bc = l2_metric_s2.geodesic(curve_b, curve_c) curves_ab = curves_ab(times) curves_bc = curves_bc(times) tangent_vecs = l2_metric_s2.log(point=curves_bc, base_point=curves_ab) result = srv_metric_r3.l2_curves_metric.pointwise_inner_products( tangent_vec_a=tangent_vecs, tangent_vec_b=tangent_vecs, base_curve=curves_ab ) expected_shape = (n_discretized_curves, n_sampling_points) self.assertAllClose(gs.shape(result), expected_shape) result = srv_metric_r3.l2_curves_metric.pointwise_inner_products( tangent_vec_a=tangent_vecs[0], tangent_vec_b=tangent_vecs[0], base_curve=curves_ab[0], ) expected_shape = (n_sampling_points,) self.assertAllClose(gs.shape(result), expected_shape)
def test_dist_vectorization(self): n_samples = self.n_samples one_point_a = gs.array([0.0, 1.0]) one_point_b = gs.array([2.0, 10.0]) n_points_a = gs.array([[2.0, 1.0], [-2.0, -4.0], [-5.0, 1.0]]) n_points_b = gs.array([[2.0, 10.0], [8.0, -1.0], [-3.0, 6.0]]) result = self.metric.dist(one_point_a, one_point_b) vec = one_point_a - one_point_b expected = gs.sqrt(gs.dot(vec, gs.transpose(vec))) self.assertAllClose(result, expected) result = self.metric.dist(n_points_a, one_point_b) self.assertAllClose(gs.shape(result), (n_samples,)) result = self.metric.dist(one_point_a, n_points_b) self.assertAllClose(gs.shape(result), (n_samples,)) result = self.metric.dist(n_points_a, n_points_b) expected = gs.array([9.0, gs.sqrt(109.0), gs.sqrt(29.0)]) self.assertAllClose(gs.shape(result), (n_samples,)) self.assertAllClose(result, expected)