<commit_before>#!/bin/env python

print("------------------------ running MatmulIndex tests")

import time
import numpy as np
import dig


def all_eq(x, y):
    if (len(x) == 0) and (len(y) == 0):
        return True
    return np.max(np.abs(x - y)) < .001


def test_index(idx_func=dig.MatmulIndex):
    N = 100 * 1000
    D = 100
    r2 = 1. * D

    # ------------------------ data and index construction

    X = np.random.randn(N, D)
    X = np.cumsum(X, axis=1)

    q = np.cumsum(np.random.randn(D))

    index = idx_func(X)
    t = index.getIndexConstructionTimeMs()
    print "index time: {}".format(t)

    # ------------------------ dists for ground truth

    t0 = time.clock()
    diffs = X - q
    trueDists = np.sum(diffs * diffs, axis=1)
    t_python = (time.clock() - t0) * 1000
    print "-> brute force time\t= {}".format(t_python)

    # ------------------------------------------------ single query

    # ------------------------ range query

    neighborIdxs = index.radius(q, np.sqrt(r2))
    t = index.getQueryTimeMs()
    neighborIdxs = np.sort(neighborIdxs)

    trueNeighborIdxs = np.where(trueDists <= r2)[0]

    if len(trueNeighborIdxs < 100):
        print "neighborIdxs: ", neighborIdxs
        print "trueNeighborIdxs: ", trueNeighborIdxs
    print "-> range query time\t= {}".format(t)

    assert(len(neighborIdxs) == len(trueNeighborIdxs))
    assert(all_eq(neighborIdxs, trueNeighborIdxs))

    # # ------------------------ knn query

    nn10 = index.knn(q, 10)
    t = index.getQueryTimeMs()
    trueNN10 = np.argsort(trueDists)[:10]
    # print "nn10, true nn10 = {}, {}".format(nn10, trueNN10)
    print "-> knn time\t\t= {}".format(t)
    assert(all_eq(nn10, trueNN10))

    # ------------------------------------------------ batch of queries

    # ------------------------ dists for ground truth

    Q = 4
    queries = np.random.randn(Q, D)
    queries = np.cumsum(queries, axis=1)

    # compute idxs in sorted order
    rowNorms = np.sum(X * X, axis=1, keepdims=True)
    t0 = time.clock()
    queryNorms = np.sum(queries * queries, axis=1)
    dotProds = np.dot(X, queries.T)
    dists = (-2 * dotProds) + rowNorms + queryNorms
    t_python = (time.clock() - t0) * 1000
    print "-> python batch brute force time ms\t= {:03} ({:.3}/query)".format(
        t_python, t_python / Q)

    # ------------------------ knn query

    k = 5
    idxs_sorted = np.argsort(dists, axis=0)
    true_knn = idxs_sorted[0:k, :].T

    nested_neighbors = index.knn_batch(queries, k)
    assert(all_eq(nested_neighbors, true_knn))
    t_cpp = index.getQueryTimeMs()
    print "-> cpp batch brute force time ms\t= {:03} ({:.3}/query)".format(
        t_cpp, t_cpp / Q)

    # ------------------------ radius

    # true_knn = np.where(dists < dists.mean())
    # print true_knn[:2]

    nested_neighbors = index.radius_batch(queries, r2)
    print nested_neighbors


if __name__ == '__main__':
    test_index()<commit_msg>add python batch range query test + refactor test code<commit_after>#!/bin/env python

print("------------------------ running MatmulIndex tests")

import time
import numpy as np
import dig


def all_eq(x, y):
    if (len(x) == 0) and (len(y) == 0):
        return True
    return np.max(np.abs(x - y)) < .001


def sq_dists_to_vectors(X, queries, rowNorms=None, queryNorms=None):
    Q = queries.shape[0]

    if rowNorms is None:
        rowNorms = np.sum(X * X, axis=1, keepdims=True)

    t0 = time.clock()

    if queryNorms is None:
        queryNorms = np.sum(queries * queries, axis=1)

    dotProds = np.dot(X, queries.T)
    dists = (-2 * dotProds) + rowNorms + queryNorms

    t_python = (time.clock() - t0) * 1000
    print "-> python batch{} knn time\t= {:03} ({:.3}/query)".format(
        Q, t_python, t_python / Q)

    idxs_sorted = np.argsort(dists, axis=0)
    return dists, idxs_sorted, t_python


def test_knn_batch_query(index, queries, idxs_sorted, k, name):
    Q = queries.shape[0]

    nested_neighbors = index.knn_batch(queries, k)
    true_knn = idxs_sorted[0:k, :].T
    assert(all_eq(nested_neighbors, true_knn))

    t_cpp = index.getQueryTimeMs()
    print "-> {} batch{} {}nn time \t= {:03} ({:.3}/query)".format(
        name, Q, k, t_cpp, t_cpp / Q)

    return t_cpp


def test_radius_batch_query(index, queries, dists, name, r2):
    nested_neighbors = index.radius_batch(queries, r2)

    for q in range(queries.shape[0]):
        # pull out idxs > 0 for this query in full dist mat
        query_dists = dists[:, q]
        true_idxs = np.where(query_dists < r2)[0]

        # find idxs > -1 in row for this query
        returned_idxs = nested_neighbors[q]
        where_valid = np.where(returned_idxs >= 0)[0]
        if len(where_valid):
            returned_idxs = returned_idxs[where_valid]
        else:
            returned_idxs = np.array([])

        # print "true_idxs:", true_idxs
        # print "returned_idxs:", returned_idxs
        assert(all_eq(returned_idxs, true_idxs))


def test_index(idx_func=dig.MatmulIndex, name="cpp", N=100, D=100, r2=-1):
    N = 100 * 1000
    D = 100
    if r2 <= 0:
        r2 = 1. * D

    # ------------------------ data and index construction

    X = np.random.randn(N, D)
    X = np.cumsum(X, axis=1)

    q = np.cumsum(np.random.randn(D))

    index = idx_func(X)
    t = index.getIndexConstructionTimeMs()
    print "{} index time: {}".format(name, t)

    # ------------------------ dists for ground truth

    t0 = time.clock()
    diffs = X - q
    trueDists = np.sum(diffs * diffs, axis=1)
    t_python = (time.clock() - t0) * 1000
    print "-> python full dists time\t= {}".format(t_python)

    # ------------------------------------------------ single query

    # ------------------------ range query

    neighborIdxs = index.radius(q, np.sqrt(r2))
    t = index.getQueryTimeMs()
    neighborIdxs = np.sort(neighborIdxs)

    trueNeighborIdxs = np.where(trueDists <= r2)[0]

    if len(trueNeighborIdxs < 100):
        print "neighborIdxs: ", neighborIdxs
        print "trueNeighborIdxs: ", trueNeighborIdxs
    print "-> {} range query time\t= {}".format(name, t)

    assert(len(neighborIdxs) == len(trueNeighborIdxs))
    assert(all_eq(neighborIdxs, trueNeighborIdxs))

    # # ------------------------ knn query

    nn10 = index.knn(q, 10)
    t = index.getQueryTimeMs()
    trueNN10 = np.argsort(trueDists)[:10]
    # print "nn10, true nn10 = {}, {}".format(nn10, trueNN10)
    print "-> {} knn time\t\t= {}".format(name, t)
    assert(all_eq(nn10, trueNN10))

    # ------------------------------------------------ batch of queries

    # ------------------------ dists for ground truth

    Q = 128
    queries = np.random.randn(Q, D)
    queries = np.cumsum(queries, axis=1)

    dists, idxs_sorted, t_python = sq_dists_to_vectors(X, queries)

    # ------------------------ knn query

    test_knn_batch_query(index, queries, idxs_sorted, name=name, k=3)
    test_knn_batch_query(index, queries, idxs_sorted, name=name, k=10)

    # ------------------------ radius

    test_radius_batch_query(index, queries, dists, name=name, r2=r2)


if __name__ == '__main__':
    test_index(dig.MatmulIndex, "matmul")
{'fraction_non_alphanumeric': 0.20786281890422417, 'fraction_numerical': 0.02383939774153074, 'mean_word_length': 2.11236442516269, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4777051', 'n_tokens_mistral': 2677, 'n_tokens_neox': 2502, 'n_words': 575}