| # Copyright 2021 Google LLC. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| """Tests for lsh_tree.""" |
| |
| import typing |
| from unittest import mock |
| |
| from absl.testing import absltest |
| import numpy as np |
| from scipy import stats |
| |
| from clustering import clustering_params |
| from clustering import lsh |
| from clustering import lsh_tree |
| from clustering import test_utils |
| |
| |
| def get_test_origin_points(nonprivate_count=5, dim=10): |
| """Points with defaults for parameters not needed for the test.""" |
| return np.zeros((nonprivate_count, dim)) |
| |
| |
| def get_test_sim_hash(dim=10, max_hash_len=1): |
| """SimHash with defaults for parameters not needed for the test.""" |
| return lsh.SimHash(dim=dim, max_hash_len=max_hash_len) |
| |
| |
| class TestLshTreeNode(lsh_tree.LshTreeNode): |
| """Test node for testing without the real children hashing logic. |
| |
| This test implementation always returns a fraction of the points with the next |
| hash_prefix containing 0, and the rest with 1, setting the private count equal |
| to the real count + 1. Equality is checked by just checking the hash_prefix |
| and private_count. |
| """ |
| # Fraction of the nonprivate_points to add a 0 to the hash_prefix. |
| frac_zero: float |
| |
| def __init__(self, *args, frac_zero: float = 0.5, **kwargs): |
| """Initializes test node. |
| |
| Args: |
| *args: extra arguments for the parent class |
| frac_zero: fraction of the nonprivate_points to add a 0 to the hash_prefix |
| **kwargs: extra arguments for the parent class |
| """ |
| super(TestLshTreeNode, self).__init__(*args, **kwargs) |
| self.frac_zero = frac_zero |
| |
| def get_private_count(self) -> int: |
| """Returns a fake private count.""" |
| self.private_count = len(self.nonprivate_points) + 1 |
| return self.private_count |
| |
| def children(self) -> typing.List[lsh_tree.LshTreeNode]: |
| """Returns fake children for this node.""" |
| cutoff = int(len(self.nonprivate_points) * self.frac_zero) |
| return [ |
| TestLshTreeNode( |
| self.hash_prefix + '0', |
| self.nonprivate_points[:cutoff], |
| self.coreset_param, |
| self.sim_hash, |
| private_count=cutoff + 1, |
| frac_zero=self.frac_zero), |
| TestLshTreeNode( |
| self.hash_prefix + '1', |
| self.nonprivate_points[cutoff:], |
| self.coreset_param, |
| self.sim_hash, |
| private_count=len(self.nonprivate_points) - cutoff + 1, |
| frac_zero=self.frac_zero) |
| ] |
| |
| def __eq__(self, other): |
| """Returns whether hash_prefix and private_count are the same for tests.""" |
| if not isinstance(other, TestLshTreeNode): |
| return False |
| return self.hash_prefix == other.hash_prefix and (self.private_count |
| == other.private_count) |
| |
| |
| class LshTreeTest(absltest.TestCase): |
| |
| @mock.patch.object(stats.dlaplace, 'rvs', return_value=-5, autospec=True) |
| def test_get_private_count_basic(self, mock_dlaplace_fn): |
| nonprivate_count = 30 |
| nonprivate_points = get_test_origin_points( |
| nonprivate_count=nonprivate_count) |
| coreset_param = test_utils.get_test_coreset_param( |
| epsilon=5, frac_sum=0.2, frac_group_count=0.8, max_depth=9) |
| sim_hash = get_test_sim_hash() |
| lsh_tree_node = lsh_tree.LshTreeNode( |
| hash_prefix='', |
| nonprivate_points=nonprivate_points, |
| coreset_param=coreset_param, |
| sim_hash=sim_hash) |
| self.assertEqual(lsh_tree_node.get_private_count(), 25) |
| mock_dlaplace_fn.assert_called_once_with(0.4) |
| |
| def test_get_private_count_cache(self): |
| nonprivate_count = 30 |
| nonprivate_points = get_test_origin_points( |
| nonprivate_count=nonprivate_count) |
| coreset_param = test_utils.get_test_coreset_param(epsilon=0.01) |
| sim_hash = get_test_sim_hash() |
| lsh_tree_node = lsh_tree.LshTreeNode( |
| hash_prefix='', |
| nonprivate_points=nonprivate_points, |
| coreset_param=coreset_param, |
| sim_hash=sim_hash) |
| |
| first_private_count = lsh_tree_node.get_private_count() |
| self.assertEqual(first_private_count, lsh_tree_node.get_private_count()) |
| |
| def test_get_children(self): |
| hash_prefix, dim, max_hash_len = '0', 5, 2 |
| datapoints = np.array([[1.5, 0, 1, 0.5, 0], [1, 1, 0, 0.7, 0.1], |
| [0.8, 0.1, 1, 0.2, 0.4], [0.1, 0.5, 0.3, 0.7, 0.8], |
| [-0.5, 0.1, -0.3, -0.4, 0.2]]) |
| # Returns children regardless of whether the node should branch. The |
| # filtering in the algorithm is done after. |
| coreset_param = test_utils.get_test_coreset_param(max_depth=max_hash_len) |
| projection_vectors = np.array([[0, 1, 1, -1, 0], [1, 0, -1, 0, 0]]) |
| sh = lsh.SimHash(dim, max_hash_len, projection_vectors) |
| node = lsh_tree.LshTreeNode(hash_prefix, datapoints, coreset_param, sh) |
| children = node.children() |
| |
| self.assertSameElements([child.hash_prefix for child in children], |
| ['00', '01']) |
| for child in children: |
| self.assertEqual(child.coreset_param, coreset_param) |
| self.assertEqual(child.sim_hash, sh) |
| if child.hash_prefix == '00': |
| self.assertTrue((child.nonprivate_points == datapoints[[0, 1]]).all()) |
| if child.hash_prefix == '01': |
| self.assertTrue((child.nonprivate_points == datapoints[[2, 3, |
| 4]]).all()) |
| |
| def test_get_children_one_empty(self): |
| hash_prefix, dim, max_hash_len = '0', 5, 2 |
| datapoints = np.array([[1.5, 0, 1, 0.5, 0], [1, 1, 0, 0.7, 0.1]]) |
| # Returns children regardless of whether the node should branch. The |
| # filtering in the algorithm is done after. |
| coreset_param = test_utils.get_test_coreset_param(max_depth=max_hash_len) |
| projection_vectors = np.array([[0, 1, 1, -1, 0], [1, 0, -1, 0, 0]]) |
| sh = lsh.SimHash(dim, max_hash_len, projection_vectors) |
| node = lsh_tree.LshTreeNode(hash_prefix, datapoints, coreset_param, sh) |
| children = node.children() |
| |
| self.assertSameElements([child.hash_prefix for child in children], |
| ['00', '01']) |
| for child in children: |
| self.assertEqual(child.coreset_param, coreset_param) |
| self.assertEqual(child.sim_hash, sh) |
| if child.hash_prefix == '00': |
| self.assertTrue((child.nonprivate_points == datapoints[[0, 1]]).all()) |
| if child.hash_prefix == '01': |
| self.assertEmpty(child.nonprivate_points) |
| |
| def test_get_children_error(self): |
| hash_prefix, dim, max_hash_len = '00', 5, 2 |
| datapoints = np.array([[1.5, 0, 1, 0.5, 0], [1, 1, 0, 0.7, 0.1]]) |
| # Returns children regardless of whether the node should branch. The |
| # filtering in the algorithm is done after. |
| coreset_param = test_utils.get_test_coreset_param(max_depth=max_hash_len) |
| projection_vectors = np.array([[0, 1, 1, -1, 0], [1, 0, -1, 0, 0]]) |
| sh = lsh.SimHash(dim, max_hash_len, projection_vectors) |
| node = lsh_tree.LshTreeNode(hash_prefix, datapoints, coreset_param, sh) |
| |
| with self.assertRaises(ValueError): |
| node.children() |
| |
| def test_filter_branching_nodes_too_few_points(self): |
| sim_hash = get_test_sim_hash() |
| # private_count, not the nonprivate_count, should be used for the check. |
| level: lsh_tree.LshTreeLevel = [ |
| lsh_tree.LshTreeNode( |
| '0', |
| get_test_origin_points(nonprivate_count=15), |
| test_utils.get_test_coreset_param( |
| min_num_points_in_branching_node=10), |
| sim_hash, |
| private_count=1), |
| ] |
| self.assertEmpty(lsh_tree.LshTree.filter_branching_nodes(level)) |
| |
| def test_filter_branching_nodes_enough_points(self): |
| sim_hash = get_test_sim_hash() |
| level: lsh_tree.LshTreeLevel = [ |
| lsh_tree.LshTreeNode( |
| '0', |
| get_test_origin_points(nonprivate_count=15), |
| test_utils.get_test_coreset_param( |
| min_num_points_in_branching_node=10), |
| sim_hash, |
| private_count=20), |
| ] |
| self.assertSequenceEqual( |
| lsh_tree.LshTree.filter_branching_nodes(level), level) |
| |
| def test_get_next_level_empty_list(self): |
| self.assertEmpty(lsh_tree.LshTree.get_next_level([])) |
| |
| def test_get_next_level(self): |
| sim_hash = get_test_sim_hash() |
| coreset_param = test_utils.get_test_coreset_param( |
| min_num_points_in_branching_node=10, min_num_points_in_node=5) |
| level: lsh_tree.LshTreeLevel = [ |
| TestLshTreeNode( |
| '0', |
| get_test_origin_points(nonprivate_count=16), |
| coreset_param, |
| sim_hash, |
| private_count=20), |
| ] |
| expected_next_level = [ |
| TestLshTreeNode( |
| '00', |
| get_test_origin_points(nonprivate_count=8), |
| coreset_param, |
| sim_hash, |
| private_count=9), |
| TestLshTreeNode( |
| '01', |
| get_test_origin_points(nonprivate_count=8), |
| coreset_param, |
| sim_hash, |
| private_count=9), |
| ] |
| branching_nodes = lsh_tree.LshTree.filter_branching_nodes(level) |
| self.assertSequenceEqual( |
| lsh_tree.LshTree.get_next_level(branching_nodes), expected_next_level) |
| |
| def test_get_next_level_filters_children_node(self): |
| sim_hash = get_test_sim_hash() |
| coreset_param = test_utils.get_test_coreset_param( |
| min_num_points_in_branching_node=10, min_num_points_in_node=9) |
| level: lsh_tree.LshTreeLevel = [ |
| # The children test nodes have a private count of 6, which is less than |
| # min_num_points_in_node. |
| TestLshTreeNode( |
| '0', |
| get_test_origin_points(nonprivate_count=10), |
| coreset_param, |
| sim_hash, |
| private_count=11), |
| # The children test nodes have a private count of 3 and 9, only the node |
| # with 9 should be in the result. |
| TestLshTreeNode( |
| '1', |
| get_test_origin_points(nonprivate_count=10), |
| coreset_param, |
| sim_hash, |
| private_count=11, |
| frac_zero=0.2), |
| ] |
| expected_next_level = [ |
| TestLshTreeNode( |
| '11', |
| get_test_origin_points(nonprivate_count=8), |
| coreset_param, |
| sim_hash, |
| private_count=9), |
| ] |
| branching_nodes = lsh_tree.LshTree.filter_branching_nodes(level) |
| self.assertSequenceEqual( |
| lsh_tree.LshTree.get_next_level(branching_nodes), expected_next_level) |
| |
| def test_root_node(self): |
| nonprivate_points = [[1, 2, 1], [0.4, 0.2, 0.8], [3, 0, 3]] |
| data = clustering_params.Data(nonprivate_points, radius=4.3) |
| coreset_param = test_utils.get_test_coreset_param(radius=4.3, max_depth=20) |
| root = lsh_tree.root_node(data, coreset_param) |
| self.assertEqual(root.hash_prefix, '') |
| self.assertSequenceEqual(root.nonprivate_points, nonprivate_points) |
| self.assertEqual(root.coreset_param, coreset_param) |
| self.assertEqual(root.sim_hash.dim, 3) |
| self.assertEqual(root.sim_hash.max_hash_len, 20) |
| self.assertIsNotNone(root.private_count) |
| |
| def test_root_node_provide_private_count(self): |
| nonprivate_points = [[1, 2, 1], [0.4, 0.2, 0.8], [3, 0, 3]] |
| data = clustering_params.Data(nonprivate_points, radius=4.3) |
| coreset_param = test_utils.get_test_coreset_param(radius=4.3, max_depth=20) |
| root = lsh_tree.root_node(data, coreset_param, private_count=10) |
| self.assertEqual(root.hash_prefix, '') |
| self.assertSequenceEqual(root.nonprivate_points, nonprivate_points) |
| self.assertEqual(root.coreset_param, coreset_param) |
| self.assertEqual(root.sim_hash.dim, 3) |
| self.assertEqual(root.sim_hash.max_hash_len, 20) |
| self.assertEqual(root.private_count, 10) |
| |
| def test_lsh_tree_empty_root_errors(self): |
| test_root = lsh_tree.LshTreeNode( |
| '0', |
| get_test_origin_points(nonprivate_count=15), |
| test_utils.get_test_coreset_param(), |
| get_test_sim_hash(), |
| private_count=0) |
| with self.assertRaises(ValueError): |
| lsh_tree.LshTree(test_root) |
| |
| def test_lsh_tree_negative_count_root_errors(self): |
| test_root = lsh_tree.LshTreeNode( |
| '0', |
| get_test_origin_points(nonprivate_count=15), |
| test_utils.get_test_coreset_param(), |
| get_test_sim_hash(), |
| private_count=-10) |
| with self.assertRaises(ValueError): |
| lsh_tree.LshTree(test_root) |
| |
| def test_lsh_tree(self): |
| # Test tree: |
| # Nodes are nonprivate count + 1. |
| # Branches to the left are 0, to the right are 1. |
| # Nodes in parentheses are filtered out. |
| # 64+1 |
| # / \ |
| # 8+1 56+1 |
| # / \ / \ |
| # (1+1) 7+1 7+1 49+1 |
| # / \ |
| # (6+1) 43+1 |
| nonprivate_count = 64 |
| sh = get_test_sim_hash() |
| cp = test_utils.get_test_coreset_param( |
| min_num_points_in_node=8, |
| min_num_points_in_branching_node=9, |
| max_depth=3) |
| test_root = TestLshTreeNode( |
| '', get_test_origin_points(nonprivate_count), cp, sh, frac_zero=0.125) |
| expected_tree = { |
| 0: [TestLshTreeNode('', get_test_origin_points(64), cp, sh)], |
| 1: [ |
| TestLshTreeNode('0', get_test_origin_points(8), cp, sh), |
| TestLshTreeNode('1', get_test_origin_points(56), cp, sh) |
| ], |
| 2: [ |
| TestLshTreeNode('01', get_test_origin_points(7), cp, sh), |
| TestLshTreeNode('10', get_test_origin_points(7), cp, sh), |
| TestLshTreeNode('11', get_test_origin_points(49), cp, sh) |
| ], |
| 3: [TestLshTreeNode('111', get_test_origin_points(43), cp, sh)], |
| } |
| tree = lsh_tree.LshTree(test_root) |
| self.assertEqual(tree.tree, expected_tree) |
| |
| def test_lsh_tree_branching_node_becomes_leaf(self): |
| # Test tree: |
| # Nodes are nonprivate count + 1. |
| # Branches to the left are 0, to the right are 1. |
| # Nodes in parentheses are filtered out. |
| # 64+1 |
| # / \ |
| # 32+1 32+1 |
| # / \ / \ |
| # (16+1) (16+1) (16+1) (16+1) |
| nonprivate_count = 64 |
| sh = get_test_sim_hash() |
| cp = test_utils.get_test_coreset_param( |
| min_num_points_in_node=20, |
| min_num_points_in_branching_node=30, |
| max_depth=5) |
| test_root = TestLshTreeNode( |
| '', get_test_origin_points(nonprivate_count), cp, sh, frac_zero=0.5) |
| expected_tree = { |
| 0: [TestLshTreeNode('', get_test_origin_points(64), cp, sh)], |
| 1: [ |
| TestLshTreeNode('0', get_test_origin_points(32), cp, sh), |
| TestLshTreeNode('1', get_test_origin_points(32), cp, sh), |
| ] |
| } |
| tree = lsh_tree.LshTree(test_root) |
| self.assertEqual(tree.tree, expected_tree) |
| |
| def test_lsh_tree_leaves(self): |
| # Test tree: |
| # Nodes are nonprivate count + 1. |
| # Branches to the left are 0, to the right are 1. |
| # Nodes in parentheses are filtered out. |
| # 64+1 |
| # / \ |
| # 8+1 56+1 |
| # / \ / \ |
| # (1+1) 7+1 7+1 49+1 |
| # / \ |
| # (6+1) 43+1 |
| nonprivate_count = 64 |
| sh = get_test_sim_hash() |
| cp = test_utils.get_test_coreset_param( |
| min_num_points_in_node=8, |
| min_num_points_in_branching_node=9, |
| max_depth=3) |
| test_root = TestLshTreeNode( |
| '', get_test_origin_points(nonprivate_count), cp, sh, frac_zero=0.125) |
| expected_leaves = [ |
| TestLshTreeNode('01', get_test_origin_points(7), cp, sh), |
| TestLshTreeNode('10', get_test_origin_points(7), cp, sh), |
| TestLshTreeNode('111', get_test_origin_points(43), cp, sh) |
| ] |
| tree = lsh_tree.LshTree(test_root) |
| self.assertEqual(tree.leaves, expected_leaves) |
| |
| def test_lsh_tree_leaves_branching_node_becomes_leaf(self): |
| # Test tree: |
| # Nodes are nonprivate count + 1. |
| # Branches to the left are 0, to the right are 1. |
| # Nodes in parentheses are filtered out. |
| # 64+1 |
| # / \ |
| # 32+1 32+1 |
| # / \ / \ |
| # (16+1) (16+1) (16+1) (16+1) |
| nonprivate_count = 64 |
| sh = get_test_sim_hash() |
| cp = test_utils.get_test_coreset_param( |
| min_num_points_in_node=20, |
| min_num_points_in_branching_node=30, |
| max_depth=5) |
| test_root = TestLshTreeNode( |
| '', get_test_origin_points(nonprivate_count), cp, sh, frac_zero=0.5) |
| expected_leaves = [ |
| TestLshTreeNode('0', get_test_origin_points(32), cp, sh), |
| TestLshTreeNode('1', get_test_origin_points(32), cp, sh), |
| ] |
| tree = lsh_tree.LshTree(test_root) |
| self.assertEqual(tree.leaves, expected_leaves) |
| |
| |
| if __name__ == '__main__': |
| absltest.main() |
