Search code examples
pythonnumpyscipysparse-matrixaffinity

Efficient ways of calculating affinity only for local neighbourhood in a large graph


I am relatively new to python and numpy and am trying to cluster a dense matrix with floating point numbers and having dimensions of 256x256 using spectral clustering. Since the affinity matrix will be of size 65536x65536, a full affinity matrix cannot be computed (due to memory limitations). As such, I am currently calculating the affinity between a given matrix entry and its 5x5 local neighbourhood, and build a sparse graph (in 3-tuple representation).

To do so, I am using for loops (basically, a sliding widow approach) which I think is not the most efficient way of doing so.

import numpy as np

def getAffinity(f1, f2):
    return np.exp(-np.linalg.norm(np.absolute(f1 - f2))/ 2.1)

G = np.arange(256*256).reshape((256,256))
dim1 = 256 # Dimension 1 of matrix
dim2 = 256 # Dimension 1 of matrix
values = np.zeros(1623076, dtype=np.float32) # To hold affinities
rows = np.zeros(1623076, dtype=np.int32) # To hold row index
cols = np.zeros(1623076, dtype=np.int32) # To hold column index
index = 0 # To hold column index

for i in range(dim1):
    for j in range(dim2):
        current = G[i, j]
        for k in range(np.maximum(0, i-2), np.minimum(dim1 , i+3)): # traverse rows
            for l in range(np.maximum(0, j-2), np.minimum(dim2 , j+3)): # traverse columns
                rows[index] = i*d1 + j
                cols[index] = k*d1 + l
                values[index] = getAffinity(current, G[k, l])
                index += 1

I was wondering whether there are any other efficient ways of achieving the same goal.


Solution

  • Here is a sparse matrix approach. It is >800x faster than the loopy code.

    import numpy as np
    from scipy import sparse
    from time import perf_counter as pc
    
    T = []
    T.append(pc())
    def getAffinity(f1, f2):
        return np.exp(-np.linalg.norm(np.absolute(f1 - f2))/ 2.1)
    
    G = 2*np.arange(256*256).reshape((256,256))
    dim1 = 256 # Dimension 1 of matrix
    dim2 = 256 # Dimension 1 of matrix
    values = np.zeros(1623076, dtype=np.float32) # To hold affinities
    rows = np.zeros(1623076, dtype=np.int32) # To hold row index
    cols = np.zeros(1623076, dtype=np.int32) # To hold column index
    index = 0 # To hold column index
    
    for i in range(dim1):
        for j in range(dim2):
            current = G[i, j]
            for k in range(np.maximum(0, i-2), np.minimum(dim1 , i+3)): # traverse rows
                for l in range(np.maximum(0, j-2), np.minimum(dim2 , j+3)): # traverse columns
                    rows[index] = i*dim1 + j
                    cols[index] = k*dim1 + l
                    values[index] = getAffinity(current, G[k, l])
                    index += 1
    
    T.append(pc())
    
    affs_OP = sparse.coo_matrix((values,(rows,cols))).tocsr()
    
    import scipy.sparse as sp
    
    def getAffinity(f1, f2): # similar to @PaulPanzer, I don't think OP is right
        return np.exp(-np.abs(f1 - f2)/ 2.1)
    
    
    def affinity_block(dim = 256, dist = 2):
        i = np.arange(-dist, dist+1)
        init_block = sp.dia_matrix((np.ones((i.size, dim)), i), (dim, dim))
        out = sp.kron(init_block, init_block).tocoo()
        out.data = getAffinity(Gf[out.row], Gf[out.col])
        return out
    
    T.append(pc())
    
    Gf = G.ravel()
    offsets = np.concatenate((np.mgrid[1:3,-2:3].reshape(2,-1).T,np.mgrid[:1,1:3].reshape(2,-1).T), axis=0)
    def make_diag(yo,xo):
        o = 256*yo+xo
        diag = np.exp(-np.abs(Gf[o:]-Gf[:-o])/2.1)
        if xo>0:
            diag[:xo-256].reshape(-1,256)[:,-xo:] = 0
        elif xo<0:
            diag[:xo].reshape(-1,256)[:,:-xo] = 0
            diag[xo:] = 0
        return diag
    diags = [make_diag(*o) for o in offsets]
    offsets = np.sum(offsets*[256,1], axis=1)
    affs_pp = sparse.diags([*diags,[np.ones(256*256)],*diags],np.concatenate([offsets,[0],-offsets]))
    
    T.append(pc())
    
    affs_df = affinity_block()
    
    T.append(pc())
    
    print("OP: {:.3f} s  convert OP to sparse matrix: {:.3f} s   pp {:.3f} s   df: {:.3f} s".format(*np.diff(T)))
    
    diff = affs_pp-affs_OP
    diff *= diff.sign()
    md = diff.max()
    
    print(f"max deviation pp-OP: {md}")
    print(f"number of different entries pp-df: {(affs_pp-affs_df).nnz}")
    

    Sample run:

    OP: 23.392 s  convert OP to sparse matrix: 0.020 s   pp 0.025 s   df: 0.093 s
    max deviation pp-OP: 2.0616356788405454e-08
    number of different entries pp-df: 0
    

    A bit of explanation, 1D first to keep it simple. Let's imagine an actually sliding window, so we can use time as an intuitive axis:

            space
      +-------------->       
      |           
    t |    xo...     x: window center
    i |    oxo..     o: window off center
    m |    .oxo.     .: non window
    e |    ..oxo
      |    ...ox
      v
    

    time here actually is equivalent to space because we move with constant speed. We can now see that all the window points can be described as three diagonals. Offsets are 0, 1 and -1 but note that because the affinities are symmetric and the one for 0 is trivial, we need only calculate them for 1.

    Now lets skip to 2D, the smallest example we can do is 3x3 window in 4x4 array. In row major this looks like.

    xo..oo..........
    oxo.ooo.........
    .oxo.ooo........
    ..ox..oo........
    oo..xo..oo......
    ooo.oxo.ooo.....
    .ooo.oxo.ooo....
    ..oo..ox..oo....
    ....oo..xo..oo..
    ....ooo.oxo.ooo.
    .....ooo.oxo.ooo
    ......oo..ox..oo
    ........oo..xo..
    ........ooo.oxo.
    .........ooo.oxo
    ..........oo..ox
    

    The relevant offsets are (0,1),(1,-1),(1,0),(1,1) or in row major 0x4+1 = 1, 1x4-1 = 3, 1x4+0 = 4, 1x4+1 = 5. Also note that most of these diagonals are not complete, the missing bits explained by row major wrapping around, i.e. at z = y,x x = 3 the right neighbor z+1 is not actually a right neighbor y,x+1 ; instead, because of line jump, it is y+1,0 The if-else clause in the code above blanks the right bits of each diagonal.

    @DanielF's strategy is similar but takes advantage of the block structure evident in the figure.

    xo.. oo.. .... ....
    oxo. ooo. .... ....
    .oxo .ooo .... ....
    ..ox ..oo .... ....
    
    oo.. xo.. oo.. ....
    ooo. oxo. ooo. ....
    .ooo .oxo .ooo ....
    ..oo ..ox ..oo ....
    
    .... oo.. xo.. oo..
    .... ooo. oxo. ooo.
    .... .ooo .oxo .ooo
    .... ..oo ..ox ..oo
    
    .... .... oo.. xo..
    .... .... ooo. oxo.
    .... .... .ooo .oxo
    .... .... ..oo ..ox