def test_count_nonempty(self): bt = SkLearnBallTreeHashIndex() # Make 1000 random bit vectors of length 256 m = np.random.randint(0, 2, 234 * 256).reshape(234, 256) bt.build_index(m) self.assertEqual(bt.count(), 234)
def test_save_model_with_cache(self, m_savez): cache_element = DataMemoryElement() bt = SkLearnBallTreeHashIndex(cache_element, random_seed=0) m = numpy.random.randint(0, 2, 1000 * 256).reshape(1000, 256) bt.build_index(m) nose.tools.assert_true(m_savez.called) nose.tools.assert_equal(m_savez.call_count, 1)
def test_save_model_no_cache(self, m_savez): bt = SkLearnBallTreeHashIndex() m = numpy.random.randint(0, 2, 1000 * 256).reshape(1000, 256) bt.build_index(m) # Underlying serialization function should not have been called # because no cache element set. nose.tools.assert_false(m_savez.called)
def test_build_index(self): bt = SkLearnBallTreeHashIndex(random_seed=0) # Make 1000 random bit vectors of length 256 m = np.random.randint(0, 2, 1000 * 256).reshape(1000, 256) bt.build_index(m) # deterministically sort index of built and source data to determine # that an index was built. self.assertIsNotNone(bt.bt) np.testing.assert_array_almost_equal( sorted(np.array(bt.bt.data).tolist()), sorted(m.tolist()))
def test_remove_from_index_invalid_key_single(self): bt = SkLearnBallTreeHashIndex(random_seed=0) index = np.ndarray((1000, 256), bool) for i in range(1000): index[i] = int_to_bit_vector_large(i, 256) bt.build_index(index) # Copy post-build index for checking no removal occurred bt_data = np.copy(bt.bt.data) self.assertRaises(KeyError, bt.remove_from_index, [ int_to_bit_vector_large(1001, 256), ]) np.testing.assert_array_equal(bt_data, np.asarray(bt.bt.data))
def test_remove_from_index_last_element_with_cache(self): """ Test removing final element also clears the cache element. """ c = DataMemoryElement() bt = SkLearnBallTreeHashIndex(cache_element=c, random_seed=0) index = np.ndarray((1, 256), bool) index[0] = int_to_bit_vector_large(1, 256) bt.build_index(index) self.assertEqual(bt.count(), 1) self.assertFalse(c.is_empty()) bt.remove_from_index(index) self.assertEqual(bt.count(), 0) self.assertTrue(c.is_empty())
def test_remove_from_index_invalid_key_multiple(self): # Test that mixed valid and invalid keys raises KeyError and does not # modify the index. bt = SkLearnBallTreeHashIndex(random_seed=0) index = np.ndarray((1000, 256), bool) for i in range(1000): index[i] = int_to_bit_vector_large(i, 256) bt.build_index(index) # Copy post-build index for checking no removal occurred bt_data = np.copy(bt.bt.data) self.assertRaises(KeyError, bt.remove_from_index, [ int_to_bit_vector_large(42, 256), int_to_bit_vector_large(1008, 256), ]) np.testing.assert_array_equal(bt_data, np.asarray(bt.bt.data))
def test_model_reload(self): fd, fp = tempfile.mkstemp('.npz') os.close(fd) os.remove(fp) # shouldn't exist before construction try: bt = SkLearnBallTreeHashIndex(fp) m = numpy.random.randint(0, 2, 1000 * 256).reshape(1000, 256) bt.build_index(m) q = numpy.random.randint(0, 2, 256).astype(bool) bt_neighbors, bt_dists = bt.nn(q, 10) bt2 = SkLearnBallTreeHashIndex(fp) bt2_neighbors, bt2_dists = bt2.nn(q, 10) nose.tools.assert_is_not(bt, bt2) nose.tools.assert_is_not(bt.bt, bt2.bt) numpy.testing.assert_equal(bt2_neighbors, bt_neighbors) numpy.testing.assert_equal(bt2_dists, bt_dists) finally: os.remove(fp)
def test_load_model(self): # Create two index instances, building model with one, and loading # the other with the cache of the first instance. Each should have # distinct model instances, but should otherwise have equal model # values and parameters. cache_element = DataMemoryElement() bt1 = SkLearnBallTreeHashIndex(cache_element, random_seed=0) m = np.random.randint(0, 2, 1000 * 256).reshape(1000, 256) bt1.build_index(m) bt2 = SkLearnBallTreeHashIndex(cache_element) self.assertIsNotNone(bt2.bt) q = np.random.randint(0, 2, 256).astype(bool) bt_neighbors, bt_dists = bt1.nn(q, 10) bt2_neighbors, bt2_dists = bt2.nn(q, 10) self.assertIsNot(bt1, bt2) self.assertIsNot(bt1.bt, bt2.bt) np.testing.assert_equal(bt2_neighbors, bt_neighbors) np.testing.assert_equal(bt2_dists, bt_dists)
def test_update_index_additive(self): # Test updating an existing index, i.e. rebuilding using the union of # previous and new data. bt = SkLearnBallTreeHashIndex(random_seed=0) # Make 1000 random bit vectors of length 256 m1 = np.random.randint(0, 2, 1000 * 256).reshape(1000, 256)\ .astype(bool) m2 = np.random.randint(0, 2, 100 * 256).reshape(100, 256).astype(bool) # Build initial index bt.build_index(m1) # Current model should only contain m1's data. np.testing.assert_array_almost_equal( sorted(np.array(bt.bt.data).tolist()), sorted(m1.tolist())) # "Update" index with new hashes bt.update_index(m2) # New model should contain the union of the data. np.testing.assert_array_almost_equal( sorted(np.array(bt.bt.data).tolist()), sorted(np.concatenate([m1, m2], 0).tolist()))
def test_remove_from_index(self): # Test that we actually remove from the index. bt = SkLearnBallTreeHashIndex(random_seed=0) index = np.ndarray((1000, 256), bool) for i in range(1000): index[i] = int_to_bit_vector_large(i, 256) bt.build_index(index) # Copy post-build index for checking no removal occurred bt_data = np.copy(bt.bt.data) bt.remove_from_index([ int_to_bit_vector_large(42, 256), int_to_bit_vector_large(998, 256), ]) # Make sure expected arrays are missing from data block. new_data = np.asarray(bt.bt.data) self.assertEqual(new_data.shape, (998, 256)) new_data_set = set(tuple(r) for r in new_data.tolist()) self.assertNotIn(tuple(int_to_bit_vector_large(42, 256)), new_data_set) self.assertNotIn(tuple(int_to_bit_vector_large(998, 256)), new_data_set)
def main(): args = cli_parser().parse_args() initialize_logging(logging.getLogger('smqtk'), logging.DEBUG) initialize_logging(logging.getLogger('__main__'), logging.DEBUG) log = logging.getLogger(__name__) hash2uuids_fp = os.path.abspath(args.hash2uuids_fp) bit_len = args.bit_len leaf_size = args.leaf_size rand_seed = args.rand_seed balltree_model_fp = os.path.abspath(args.balltree_model_fp) assert os.path.isfile(hash2uuids_fp), "Bad path: '%s'" % hash2uuids_fp assert os.path.isdir(os.path.dirname(balltree_model_fp)), \ "Bad path: %s" % balltree_model_fp log.debug("hash2uuids_fp : %s", hash2uuids_fp) log.debug("bit_len : %d", bit_len) log.debug("leaf_size : %d", leaf_size) log.debug("rand_seed : %d", rand_seed) log.debug("balltree_model_fp: %s", balltree_model_fp) log.info("Loading hash2uuids table") with open(hash2uuids_fp) as f: hash2uuids = cPickle.load(f) log.info("Computing hash-code vectors") hash_vectors = [] #[int_to_bit_vector_large(h, bit_len) for h in hash2uuids] rs = [0] * 7 for h in hash2uuids: hash_vectors.append( int_to_bit_vector_large(h, bit_len) ) report_progress(log.debug, rs, 1.) log.info("Initializing ball tree") btree = SkLearnBallTreeHashIndex(balltree_model_fp, leaf_size, rand_seed) log.info("Building ball tree") btree.build_index(hash_vectors)
def test_remove_from_index(self): # Test that we actually remove from the index. bt = SkLearnBallTreeHashIndex(random_seed=0) index = np.ndarray((1000, 256), bool) for i in range(1000): index[i] = int_to_bit_vector_large(i, 256) bt.build_index(index) # BallTree data should now contain 1000 entries self.assertEqual(bt.bt.data.shape, (1000, 256)) bt.remove_from_index([ int_to_bit_vector_large(42, 256), int_to_bit_vector_large(998, 256), ]) # Make sure data block is of the expected shape (two rows shorter) new_data = np.asarray(bt.bt.data) self.assertEqual(new_data.shape, (998, 256)) # Make sure expected arrays are missing from data block. new_data_set = set(tuple(r) for r in new_data.tolist()) self.assertNotIn(tuple(int_to_bit_vector_large(42, 256)), new_data_set) self.assertNotIn(tuple(int_to_bit_vector_large(998, 256)), new_data_set)
def test_remove_from_index_last_element(self): """ Test removing the final the only element / final elements from the index. """ # Add one hash, remove one hash. bt = SkLearnBallTreeHashIndex(random_seed=0) index = np.ndarray((1, 256), bool) index[0] = int_to_bit_vector_large(1, 256) bt.build_index(index) self.assertEqual(bt.count(), 1) bt.remove_from_index(index) self.assertEqual(bt.count(), 0) self.assertIsNone(bt.bt) # Add many hashes, remove many hashes in batches until zero bt = SkLearnBallTreeHashIndex(random_seed=0) index = np.ndarray((1000, 256), bool) for i in range(1000): index[i] = int_to_bit_vector_large(i, 256) bt.build_index(index) # Remove first 250 bt.remove_from_index(index[:250]) self.assertEqual(bt.count(), 750) self.assertIsNotNone(bt.bt) # Remove second 250 bt.remove_from_index(index[250:500]) self.assertEqual(bt.count(), 500) self.assertIsNotNone(bt.bt) # Remove third 250 bt.remove_from_index(index[500:750]) self.assertEqual(bt.count(), 250) self.assertIsNotNone(bt.bt) # Remove final 250 bt.remove_from_index(index[750:]) self.assertEqual(bt.count(), 0) self.assertIsNone(bt.bt)
def test_save_model_with_cache(self, m_savez): bt = SkLearnBallTreeHashIndex('some_file.npz') m = numpy.random.randint(0, 2, 1000 * 256).reshape(1000, 256) bt.build_index(m) nose.tools.assert_true(m_savez.called) nose.tools.assert_equal(m_savez.call_count, 1)
def test_save_model_no_cache(self, m_savez): bt = SkLearnBallTreeHashIndex() m = numpy.random.randint(0, 2, 1000 * 256).reshape(1000, 256) bt.build_index(m) nose.tools.assert_false(m_savez.called)