SSE/_L2_ImplNumba.py at main · SELGroup/SSE · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
# taken from source code of paper CMS

from numba import jit
import numpy as np


@jit(nopython=True)
def _find_min_max_l2_impl(points):
    x_sq = 0
    for d in range(points.shape[1]):
        x_sq += (points[0, d] - points[1, d]) ** 2

    l2_min = x_sq
    l2_max = x_sq

    for i in range(len(points)):
        for j in range(i + 1, len(points)):
            x_sq = 0
            for d in range(points.shape[1]):
                x_sq += (points[i, d] - points[j, d]) ** 2

            if x_sq == 0:
                continue

            if x_sq < l2_min:
                l2_min = x_sq
            if x_sq > l2_max:
                l2_max = x_sq

    return l2_min, l2_max


@jit("void(float32[:,:], float32[:,:], float32[:,:])", nopython=True)
def _l2_squared_impl(left, right, result):
    for i in range(len(left)):
        for j in range(len(right)):
            result[i, j] = np.sum((left[i] - right[j]) ** 2)


@jit("void(float32[:,:], float32[:,:])", nopython=True)
def _l2_squared_self_impl(left, result):
    for i in range(len(left)):
        for j in range(i, len(left)):
            r = np.sum((left[i] - left[j]) ** 2)
            result[i, j] = r
            result[j, i] = r


def numba_l2_squared(left, right=None):
    assert len(left.shape) == 2

    if right is None:
        result = np.empty((len(left),) * 2, dtype=np.float32)
        _l2_squared_self_impl(left.astype(np.float32), result)
    else:
        assert len(right.shape) == 2
        assert right.shape[1] == left.shape[1]
        result = np.empty((len(left), len(right)), dtype=np.float32)
        _l2_squared_impl(left.astype(np.float32), right.astype(np.float32), result)
    return result