ConvGeNCode/library/NNSearch_experimental.py

import math
from ctypes import cdll, c_uint, c_double, c_void_p

import tensorflow as tf
import tensorflow.keras.layers as L
import numpy as np
import numpy.ctypeslib as npct

array_2d_uint = npct.ndpointer(dtype=np.uint, ndim=2, flags='CONTIGUOUS')
array_1d_double = npct.ndpointer(dtype=np.double, ndim=1, flags='CONTIGUOUS')
array_2d_double = npct.ndpointer(dtype=np.double, ndim=2, flags='CONTIGUOUS')


from sklearn.neighbors import NearestNeighbors
from library.timing import timing
from library.MaxHeap import MaxHeap


nbhLib = cdll.LoadLibrary("./library/c/libNeighborhood.so")
nbhLib.Neighborhood.rettype = None
nbhLib.Neighborhood.argtypes = [c_uint, c_uint, c_uint, array_2d_double, array_2d_uint]

nbhLib.NeighborhoodHeuristic.rettype = None
nbhLib.NeighborhoodHeuristic.argtypes = [c_uint, c_uint, c_uint, array_2d_double, array_2d_uint]



def scalarP(a, b):
    return sum(map(lambda c: c[0] * c[1], zip(a, b)))

def norm2(v):
    return sum(map(lambda z: z*z, v))

def dist(x,y):
    return norm2(x - y)

def maxby(data, fn, startValue=0.0):
    m = startValue
    for v in data:
        m = max(m, fn(v))
    return m

def distancesToPoint(p, points):
    w = np.array(np.repeat([p], len(points), axis=0))
    d = L.Subtract()([w, np.array(points)])
    t = L.Dot(axes=(1,1))([d,d])
    # As the concrete distance is not needed and sqrt(x) is strict monotone
    # we avoid here unneccessary calculating of expensive roots.
    return t.numpy()


def calculateCenter(points):
    if points.shape[0] == 1:
        return points[0]
    return tf.keras.layers.Average()(np.array(points)).numpy()


def centerPoints(points):
    points = np.array(points)
    center = L.Average()(list(points)).numpy()
    ctr = np.array(np.repeat([center], points.shape[0], axis=0))
    return L.Subtract()([ctr, points]).numpy()
    

def maxNormPoints(points):
    points = np.array(points)
    a = L.Lambda(lambda x: np.abs(x))(points)
    a = L.Reshape((points.shape[1], 1))(a)
    m = L.GlobalMaxPooling1D()(a)
    m = L.Reshape((1,))
    return m.numpy()


def twoNormSquaredPoints(points):
    points = np.array(points)
    return L.Dot(axes=(1,1))([points,points]).numpy()
    

def twoNormPoints(points):
    points = np.array(points)
    nsq = L.Dot(axes=(1,1))([points,points])
    return L.Lambda(lambda x: np.sqrt(x))(points).numpy()


def norms(points):
    points = np.array(points)
    a = L.Lambda(lambda x: np.abs(x))(points)
    a = L.Reshape((points.shape[1], 1))(a)
    m = L.GlobalMaxPooling1D()(a)
    m = L.Reshape((1,))(m)
    nsq = L.Dot(axes=(1,1))([points,points])
    return L.Concatenate()([m, nsq]).numpy()



class Ball:
    def __init__(self, points=None, indices=None, parent=None, center=None):
        if center is not None:
            self.center = center
        elif points is not None:
            self.center = calculateCenter(np.array(points))
        else:
            raise ParameterError("Missing points or center")
        self.radius = 0
        self.points = []
        self.indices = set()
        self.childs = []
        self.parent = parent

        if points is not None and indices is not None:
            for (i, r) in zip(indices, distancesToPoint(self.center, points)):
                self.add(i, r)
        elif points is not None:
            raise ParameterError("Missing indices")

    def findIt(self, r):
        if r == self.points[0][1]:
            return 0

        upper = len(self.points) - 1
        lower = 0
        while upper > lower + 1:
            h = (upper + lower) // 2
            if self.points[h][1] >= r:
                upper = h
            else:
                lower = h
        return upper


    def add(self, xi, r):
        if xi in self.indices:
            return False

        # Here we know, that x is not in points.
        newEntry = (xi, r)
        self.indices.add(xi)

        # Special case: empty list or new element will extend the radius:
        # Here we can avoid the search
        if self.points == [] or r >= self.radius:
            self.points.append(newEntry)
            self.radius = r
            return True

        # Special case: r <= min radius
        # Here we can avoid the search
        if self.points[0][1] >= r:
            self.points = [newEntry] + self.points
            return True

        # Here we know that min radius < r < max radius.
        # So len(points) >= 2.

        pos = self.findIt(r)

        # here shoul be: r(pos) >= r > r(pos - 1)
        self.points = self.points[:pos] + [newEntry] + self.points[pos: ]

        return True

    def remove(self, xi, r):
        if xi not in self.indices:
            return False


        # special case: remove the element with the heighest radius:
        if self.points[-1][0] == xi:
            self.indices.remove(xi)
            self.points = self.points[: -1]
            if self.points == []:
                self.radius = 0.0
            else:
                self.radius = self.points[-1][1]
            return True

        pos = self.findIt(r)
        nPoints = len(self.points)
        # pos is the smalest position with:
        # r(pos - 1) < r <= r(pos)

        # Just in case: check that we remove the correct index.
        while pos < nPoints and r == self.points[pos]:
            if self.points[pos][0] == xi:
                self.indices.remove(xi)
                self.points = self.points[:pos] + self.points[(pos + 1): ]
                return True
            pos += 1

        return False

    def divideBall(self, X):
        indices = [i for (i, _r) in self.points]
        points = np.array([X[i] for i in indices])
        
        distances = tf.keras.layers.Dot(axes=(1,1))([points, points]).numpy()
        ball = Ball(points, indices, center=np.zeros(len(points[0])))

        h = len(ball) // 2
        indicesA = [i for (i, _) in ball.points[:h]] 
        indicesB = [i for (i, _) in ball.points[h:]] 
        pointsA = np.array([X[i] for i in indicesA])
        pointsB = np.array([X[i] for i in indicesB])

        self.childs = []

        print(f"{len(points)} -> <{len(pointsA)}|{len(pointsB)}>")
        self.childs.append(Ball(pointsA, indicesA, self))
        self.childs.append(Ball(pointsB, indicesB, self))

        return self.childs

    def __len__(self):
        return len(self.points)

    def smalestBallFor(self, i):
        if i not in self.indices:
            return None

        for c in self.childs:
            b = c.smalestBallFor(i)
            if b is not None:
                return b

        return self



class NNSearch:
    def __init__(self, nebSize=5, timingDict=None):
        self.nebSize = nebSize
        self.neighbourhoods = []
        self.timingDict = timingDict


    def timerStart(self, name):
        if self.timingDict is not None:
            if name not in self.timingDict:
                self.timingDict[name] = timing(name)

            self.timingDict[name].start()

    def timerStop(self, name):
        if self.timingDict is not None:
            if name in self.timingDict:
                self.timingDict[name].stop()

    def neighbourhoodOfItem(self, i):
        return self.neighbourhoods[i]


    def fit(self, X, nebSize=None):
        self.timerStart("NN_fit_all")
        self.fit_chained(X, nebSize)
        #self.fit_bruteForce_np(X, nebSize)
        self.timerStop("NN_fit_all")
       
    def fit_optimized(self, X, nebSize=None):
        self.timerStart("NN_fit_all")
        if nebSize == None:
            nebSize = self.nebSize

        nPoints = len(X)
        nFeatures = len(X[0])

        if nFeatures > 15 or nebSize >= (nPoints // 2):
            print("Using brute force")
            self.fit_bruteForce_np(X, nebSize)
        else:
            print("Using chained")
            self.fit_chained(X, nebSize)
        self.timerStop("NN_fit_all")


    def fit_bruteForce(self, X, nebSize=None):
        if nebSize == None:
            nebSize = self.nebSize

        self.timerStart("NN_fit_bf_init")
        isGreaterThan = lambda x, y: x[1] > y[1]
        self.neighbourhoods = [MaxHeap(nebSize, isGreaterThan, (i, 0.0)) for i in range(len(X))]
        self.timerStop("NN_fit_bf_init")

        self.timerStart("NN_fit_bf_loop")
        for (i, x) in enumerate(X):
            nbh = self.neighbourhoods[i]

            for (j, y) in enumerate(X[i+1:]):
                j += i + 1
                self.timerStart("NN_fit_bf_dist")
                d = dist(x,y)
                self.timerStop("NN_fit_bf_dist")
                self.timerStart("NN_fit_bf_insert")
                nbh.insert((j,d))
                self.neighbourhoods[j].insert((i,d))
                self.timerStop("NN_fit_bf_insert")

            self.timerStart("NN_fit_bf_toList")
            self.neighbourhoods[i] = nbh.toArray(lambda v: v[0])
            self.timerStop("NN_fit_bf_toList")
        self.timerStop("NN_fit_bf_loop")


    def fit_bruteForce_np(self, X, nebSize=None):
        self.timerStart("NN_fit_bfnp_init")
        numOfPoints = len(X)
        nFeatures = len(X[0])

        tX = tf.convert_to_tensor(X)

        def distancesTo(x):
            w = np.repeat([x], numOfPoints, axis=0)
            d = tf.keras.layers.Subtract()([w,tX])
            t = tf.keras.layers.Dot(axes=(1,1))([d,d])
            return t.numpy()

        if nebSize == None:
            nebSize = self.nebSize

        isGreaterThan = lambda x, y: x[1] > y[1]
        self.neighbourhoods = [MaxHeap(nebSize, isGreaterThan, (i, 0.0)) for i in range(len(X))]
        self.timerStop("NN_fit_bfnp_init")

        self.timerStart("NN_fit_bfnp_loop")
        for (i, x) in enumerate(X):
            self.timerStart("NN_fit_bfnp_dist")
            distances = distancesTo(x)
            self.timerStop("NN_fit_bfnp_dist")
            nbh = self.neighbourhoods[i]

            for (j, y) in enumerate(X[i+1:]):
                j += i + 1
                d = distances[j]
                self.timerStart("NN_fit_bfnp_insert")
                nbh.insert((j,d))
                self.neighbourhoods[j].insert((i,d))
                self.timerStop("NN_fit_bfnp_insert")

            self.timerStart("NN_fit_bfnp_toList")
            self.neighbourhoods[i] = nbh.toArray(lambda v: v[0])
            self.timerStop("NN_fit_bfnp_toList")
        self.timerStop("NN_fit_bfnp_loop")


    def fit_chained(self, X, nebSize=None):
        self.timerStart("NN_fit_chained_init")
        if nebSize == None:
            nebSize = self.nebSize

        nPoints = len(X)
        nFeatures = len(X[0])

        neigh = NearestNeighbors(n_neighbors=nebSize)
        neigh.fit(X)
        self.timerStop("NN_fit_chained_init")
        self.timerStart("NN_fit_chained_toList")
        self.neighbourhoods = [
                (neigh.kneighbors([x], nebSize, return_distance=False))[0]
                for (i, x) in enumerate(X)
                ]
        self.timerStop("NN_fit_chained_toList")


    def fit_cLib(self, X, nebSize=None):
        self.timerStart("NN_fit_cLib_init")
        if nebSize == None:
            nebSize = self.nebSize
        
        nbh = np.array([np.zeros(nebSize, dtype=np.uint) for i in range(X.shape[0])])
        self.timerStop("NN_fit_cLib_init")
        self.timerStart("NN_fit_cLib_call")
        nbhLib.Neighborhood(nebSize, X.shape[0], X.shape[1], X, nbh)
        self.timerStop("NN_fit_cLib_call")
        self.timerStart("NN_fit_cLib_list")
        self.neighbourhoods = list(nbh)
        self.timerStop("NN_fit_cLib_list")

    def fit_cLibHeuristic(self, X, nebSize=None):
        self.timerStart("NN_fit_cLib_init")
        if nebSize == None:
            nebSize = self.nebSize
        
        nbh = np.array([np.zeros(nebSize, dtype=np.uint) for i in range(X.shape[0])])
        self.timerStop("NN_fit_cLib_init")
        self.timerStart("NN_fit_cLib_call")
        nbhLib.NeighborhoodHeuristic(nebSize, X.shape[0], X.shape[1], X, nbh)
        self.timerStop("NN_fit_cLib_call")
        self.timerStart("NN_fit_cLib_list")
        self.neighbourhoods = list(nbh)
        self.timerStop("NN_fit_cLib_list")
        

    # ===============================================================
    # Heuristic search
    # ===============================================================
    def fit_heuristic(self, X, nebSize=None, debugLayer=0, withDouble=True):
        if nebSize == None:
            nebSize = self.nebSize

        self.timerStart("NN_fit_heuristic_init")
        nPoints = len(X)
        nFeatures = len(X[0])
        nHeuristic = max(1, int(math.log(nFeatures)))
        isGreaterThan = lambda x, y: x[1] > y[1]
        self.neighbourhoods = [MaxHeap(maxSize=nebSize, isGreaterThan=isGreaterThan, smalestValue=(i, 0.0)) for i in range(len(X))]


        self.timerStart("NN_fit_heuristic_lineStart")
        z = X[0]
        farest = 0
        bestDist = 0.0
        for (i, x) in enumerate(X):
            d = dist(x, z)
            if d > bestDist:
                farest = i
                bestDist = d

        lineStart = farest
        z = X[lineStart]
        self.timerStop("NN_fit_heuristic_lineStart")

        # print(f"lineStart: {lineStart}@{z} ... {bestDist}")

        self.timerStart("NN_fit_heuristic_lineEnd")
        bestDist = 0.0
        for (i, x) in enumerate(X):
            d = dist(x, z)
            if d > bestDist:
                farest = i
                bestDist = d

        lineEnd = farest
        self.timerStop("NN_fit_heuristic_lineEnd")

        self.timerStart("NN_fit_heuristic_line")
        # print(f"lineEnd: {lineEnd}@{X[lineEnd]} ... {bestDist}")
        u = (X[lineEnd] - z)
        uFactor = (1 / math.sqrt(norm2(u)))
        u = uFactor * u
        # print(f"u: {u} ... {norm2(u)}")

        def heuristic(i,x):
            p = z + (scalarP(u, x - z) * u)
            dz = math.sqrt(dist(z, p))
            dx = math.sqrt(dist(x, p)) 
            return (i, dz, dx)

        line = [heuristic(i, x) for (i,x) in enumerate(X) ]
        line.sort(key= lambda a: a[1])
        self.timerStop("NN_fit_heuristic_line")
        self.timerStop("NN_fit_heuristic_init")

        self.timerStart("NN_fit_heuristic_loop")
        s = 0
        ff = False
        ptsDone = set()
        for (i,(xi, di, dix)) in enumerate(line):
            self.timerStart("NN_fit_heuristic_loop_init")
            h = self.neighbourhoods[xi]
            z = X[xi]
            self.timerStop("NN_fit_heuristic_loop_init")
            ptsDone.add(xi)

            self.timerStart("NN_fit_heuristic_loop_distance")
            ll = [(xj, norm2([dj - di, djx - dix])) for (xj, dj, djx) in line[i:]]
            # ll = [(xj, dist(
            #     np.array([di, dix] + list(X[xi][0:nHeuristic])),
            #     np.array([dj, djx] + list(X[xj][0:nHeuristic]))))
            #     for (xj, dj, djx) in line[1:]
            #     ]
            ll.sort(key = lambda a: a[1])
            kk = distancesToPoint(z, [X[j] for (j, _) in ll])
            self.timerStop("NN_fit_heuristic_loop_distance")

            for (d, (xj, djx)) in zip(kk, ll):
                ign = h.size >= nebSize and djx > h.getMax()[1]
                if ign:
                    break
                else:
                    #d = dist(X[xj], z)
                    self.timerStart("NN_fit_heuristic_insert")
                    s += 1
                    h.insert((xj, d))
                    k = self.neighbourhoods[xj]
                    if not isinstance(k, list):
                        k.insert((xi, d))
                    self.timerStop("NN_fit_heuristic_insert")
    
                # if xi == debugLayer:
                #     d = dist(X[xj], z)
                #     hint = ""
                #     if djx > d:
                #         hint += "!!"
                #     if ign:
                #         hint += "*"
                #     print(f"xj:{xj}   dx:{h.getMax()[1]:0.1f}   djx:{djx:0.1f}  d:{d:0.1f}" + hint)
            self.timerStart("NN_fit_heuristic_toArray")
            self.neighbourhoods[xi] = h.toArray()
            self.timerStop("NN_fit_heuristic_toArray")
        self.timerStop("NN_fit_heuristic_loop")
        print(f"calculated distances: {s} / {nPoints * (nPoints - 1)}")