Add tests

Author: Szubie

requirements.txt

@@ -4,3 +4,4 @@ numpy
 scikit-learn>0.20.0
 tqdm
 git+https://github.com/beringresearch/annoy.git#egg=annoy
+pytest

setup.py

@@ -17,4 +17,7 @@
           'annoy',
           'tqdm'
       ],
+      extras_require={
+          'tests': ['pytest']
+      },
       zip_safe=False)

tests/conftest.py

@@ -0,0 +1,13 @@
+import pytest
+from keras import backend as K
+
+
+@pytest.fixture(autouse=True)
+def clear_session_after_test():
+    """Test wrapper to clean up after TensorFlow and CNTK tests.
+
+    This wrapper runs for all the tests in the keras test suite.
+    """
+    yield
+    if K.backend() == 'tensorflow' or K.backend() == 'cntk':
+        K.clear_session()

tests/data/.test-annoy-index.index

@@ -0,0 +1,56 @@
+from ivis.data.knn import build_annoy_index, extract_knn
+
+from annoy import AnnoyIndex
+import numpy as np
+from scipy.sparse import csr_matrix
+from sklearn import datasets
+import tempfile
+import pytest
+import os
+
+
+@pytest.fixture(scope='function')
+def annoy_index_file():
+    _, filepath = tempfile.mkstemp('.index')
+    yield filepath
+    os.remove(filepath)
+
+def test_build_sparse_annoy_index(annoy_index_file):
+    data = np.random.choice([0, 1], size=(10, 5))
+    sparse_data = csr_matrix(data)
+
+    index = build_annoy_index(sparse_data, annoy_index_file)
+    assert os.path.exists(annoy_index_file)
+
+    loaded_index = AnnoyIndex(5)
+    loaded_index.load(annoy_index_file) 
+
+    assert index.f == loaded_index.f == 5
+    assert index.get_n_items() == loaded_index.get_n_items() == 10
+    assert index.get_nns_by_item(0, 5) == loaded_index.get_nns_by_item(0, 5)
+
+
+def test_dense_annoy_index(annoy_index_file):
+    data = np.random.choice([0, 1], size=(10, 5))
+    index = build_annoy_index(data, annoy_index_file)
+    assert os.path.exists(annoy_index_file)
+
+    loaded_index = AnnoyIndex(5)
+    loaded_index.load(annoy_index_file) 
+
+    assert index.f == loaded_index.f == 5
+    assert index.get_n_items() == loaded_index.get_n_items() == 10
+    assert index.get_nns_by_item(0, 5) == loaded_index.get_nns_by_item(0, 5)
+
+def test_knn_retrieval():
+    annoy_index_filepath = 'tests/data/.test-annoy-index.index'
+    expected_neighbour_list = np.load('tests/data/test_knn_k4.npy')
+    
+    iris = datasets.load_iris()
+    X = iris.data
+    
+    k = 4
+    search_k = -1
+    neighbour_list = extract_knn(X, annoy_index_filepath, k=k, search_k=search_k)
+
+    assert np.all(expected_neighbour_list == neighbour_list)

tests/data/test_knn.py

@@ -0,0 +1,27 @@
+from ivis.data.triplet_generators import generate_knn_triplets_from_neighbour_list
+from sklearn import datasets
+
+import numpy as np
+
+def test_generate_knn_triplets_from_neighbour_list():
+    neighbour_list = np.load('tests/data/test_knn_k4.npy')
+
+    iris = datasets.load_iris()
+    X = iris.data
+    batch_size = 32
+
+    data_generator = generate_knn_triplets_from_neighbour_list(X, neighbour_list, batch_size=batch_size)
+    
+    # Run generator thorugh one iteration of dataset and into the next
+    for i in range((X.shape[0] // batch_size) + 2):
+        batch = next(data_generator)
+
+        # Check that everything is the expected shape
+        assert isinstance(batch, tuple)
+        assert len(batch) == 2
+
+        assert len(batch[0]) == 3
+        assert len(batch[1]) == batch_size
+        assert batch[0][0].shape == (batch_size, X.shape[-1])
+
+

tests/data/test_knn_k4.npy

@@ -0,0 +1,18 @@
+from ivis import Ivis
+from sklearn import datasets
+import numpy as np
+
+
+def test_iris_embedding():
+    iris = datasets.load_iris()
+    x = iris.data
+    y = iris.target
+
+    val_x = x
+    val_y = y
+
+    ivis_iris = Ivis(n_epochs_without_progress=5)
+    ivis_iris.k = 15
+    ivis_iris.batch_size = 16
+
+    y_pred_iris = ivis_iris.fit_transform(x)

tests/data/test_triplet_generator.py

@@ -0,0 +1,17 @@
+from ivis.nn import losses
+from ivis.nn.losses import triplet_loss, get_loss_functions
+
+from keras import backend as K
+import numpy as np
+import tensorflow as tf
+
+
+def test_loss_function_call():
+    margin = 2
+
+    loss_dict = get_loss_functions(margin=margin)
+    
+    for loss_name in loss_dict.keys():
+        loss_function = triplet_loss(distance=loss_name, margin=margin)
+        assert loss_function.__name__ == loss_name
+        

tests/integration/test_iris.py

@@ -0,0 +1,32 @@
+from ivis.nn.network import get_base_networks, base_network, build_network
+
+import keras
+from keras.models import Model, Sequential
+from keras.layers import Dense
+import numpy as np
+
+def test_base_networks():
+    network_names = get_base_networks()
+    input_shape = (4,)
+
+    for name in network_names:
+        model = base_network(name, input_shape)
+        assert isinstance(model, Model)
+
+def test_build_network():
+
+    X = np.zeros(shape=(10, 5))
+    embedding_dims = 3
+
+    base_model = Sequential()
+    base_model.add(Dense(8, input_shape=(X.shape[-1],)))
+
+    model = build_network(base_model, embedding_dims=embedding_dims, embedding_l2=0.1)
+    encoder = model.layers[3]
+
+    assert model.layers[3].output_shape == (None, 3)
+    assert np.all(base_model.get_weights()[0] == encoder.get_weights()[0])
+    assert np.all([isinstance(layer, keras.layers.InputLayer) for layer in model.layers[:3]])
+
+    assert encoder.output_shape == (None, embedding_dims)
+