risc.py

from __future__ import absolute_import
import sys
sys.path.append('/risc/Code')
print(sys.path)
import pyrisc
from ann_benchmarks.algorithms.base import BaseANN
from scipy.sparse import csr_matrix
import numpy
import os

class Risc(BaseANN):

    def __init__(self, metric, method):
        if metric != "jaccard":
            raise NotImplementedError("BruteForce doesn't support metric %s, only jaccard metric is supported." % metric)
        methods = {'Risc': 1, 'Linearscan': 2, 'AOR': 3, 'DivideSkip': 4}
        self._metric = metric
        self._method = methods[method]
        self.name = method + "()"

    def pre_fit(self, X):
        X = numpy.concatenate((X, [numpy.ones(X.shape[1], dtype=numpy.bool)]), axis=0)
        print(X.shape)
        def matrToStrArray(sparseMatr):
            res = ""
            indptr = sparseMatr.indptr
            indices = sparseMatr.indices
            for row in range(sparseMatr.shape[0]):
                arr = [k for k in indices[indptr[row]: indptr[row + 1]]]
                arr.sort()
                res1 = "{" + ':1 , '.join([str(k) for k in arr]) + ':1}'
                res += res1 + "\n"
            return res

        # transform data and store in file
        data_trans = matrToStrArray(csr_matrix(X))
        # print(data_trans)
        text_file = open("train.txt", "w")
        text_file.write(data_trans)
        text_file.close()

        # call function with file
        self._featureId = pyrisc.getFeatureId("train.txt", "features.txt")
        self._data = pyrisc.readDatabase("train.txt", self._featureId)
        # self._data = pyrisc.readDatabase("train.txt", "features.txt")


    def fit(self, X):
        self._index = pyrisc.getIndex(self._data, self._method)

    def pre_query(self, v, n):
        # transform data and store in file
        nz = numpy.nonzero(v)[0]
        v = '{' + ':1 , '.join([str(k) for k in nz]) + ':1}\n'
        if os.path.isfile("query.txt"):
            os.remove("query.txt")
        text_file = open("query.txt", "w")
        text_file.write(v)
        text_file.close()

        # queries = pyrisc.readQueries("train.txt", "query.txt", "features.txt")
        queries = pyrisc.readQueries("query.txt", self._featureId)
        self._queryFP = pyrisc.dataBinary_getFingerPrint(queries, 0)

    def query(self, v, n, rq=False):
        if rq:
            self._results = pyrisc._experiments_runRange_InMemory(self._index, self._data, self._queryFP, 1.0-n, self._method)
        else:
            self._n = n
            self._results = pyrisc._experiments_runTopK_inMemory(self._index, self._data, self._queryFP, n, self._method)

    def post_query(self, rq=False):
        if os.path.isfile("result.txt"):
            os.remove("result.txt")
        if rq:
            pyrisc.writeResults_Range("result.txt", self._data, self._results)
        else:
            pyrisc.writeResults("result.txt", self._data, self._results, self._n)

        # read results from output file
        result = []
        with open("result.txt", "r") as fp:
            line = fp.readline()
            while line:
                if line.startswith("#"):
                    line = fp.readline()
                    continue
                # make 1 based index 0 based
                result.append(int(line[:-1])-1)
                line = fp.readline()
        return result