nmf.py

import numpy as np
import cvxpy as cp
import random

def is_reducible(X):
    def permutation(blocks, n):
        re = np.zeros((n, n))
        flat_compo = [c for compo in blocks for c in compo]
        for i, j in zip(flat_compo, range(n)):
            re[i, j] = 1
        return re
    def connected_components(X):
        n = len(X)
        visited = [False] * n
        subblocks = []
        for v in range(len(X)):
            if not visited[v]:
                queue = [v]
                visited[v] = True
                connected = []
                while queue:
                    node = queue.pop(-1)
                    connected.append(node)
                    for i in range(n):
                        if i != node and not visited[i] and abs(X[node][i]) > 1e-10:
                            queue.append(i)
                            visited[i] = True
                subblocks.append(connected)
        return subblocks

    clusters = connected_components(X)
    if len(clusters) > 1:
        sub_matrices = []
        blocksize = [len(block) for block in clusters]
        P = permutation(clusters, len(X))
        block_formed_X = P @ X @ P.T
        idx = 0
        for size in blocksize:
            sub_matrices.append(block_formed_X[idx:idx+size, idx:idx+size])
            idx += size
        return {'reducible':True, 'block_idx':clusters, 'P':P, 'blocks':sub_matrices}
    else:
        return {'reducible':False, 'block_idx':None, 'P':None, 'blocks':None}

def nonneg_biconvex(X, r, maxdiff, init=None, constraints=[], weights=[], eta=1/4, verbose=False):
    W0 = initialize_W(X, r, init)
    W = np.array(W0)
    Y = np.random.random((r, r))
    t = 0
    err = lambda W, Y: np.linalg.norm(W@Y@W.T - X, ord='fro')
    def objective(constraints, weights):
        # all loss/ constrainst are functions of W and Y; if no dependency None is returned
        loss = err
        DYp = lambda W, Y: W.T @ W @ Y @ W.T @ W + 1e-10
        DYm = lambda W, Y: W.T @ X @ W
        DWp = lambda W, Y: W @ Y @ W.T @ W @ Y.T + W @ Y.T @ W.T @ W @ Y + 1e-10
        DWm = lambda W, Y: X @ W @ Y.T + X.T @ W @ Y
        sum_function = lambda f, g, w: lambda x, y: f(x, y) + w * g(x, y)
        if len(constraints) > 0:
            for gen_con, w in zip(constraints, weights):
                con = gen_con()
                loss = sum_function(loss, con['loss'], w)
                if con['DYp'] is not None: DYp = sum_function(DYp, con['DYp'], w)
                if con['DYm'] is not None: DYm = sum_function(DYm, con['DYm'], w)
                if con['DWp'] is not None: DWp = sum_function(DWp, con['DWp'], w)
                if con['DWm'] is not None: DWm = sum_function(DWm, con['DWm'], w)
        return {'loss':loss, 'DYp':DYp, 'DYm':DYm, 'DWp':DWp, 'DWm':DWm}

    problem = objective(constraints, weights)
    dYp, dYm = problem['DYp'], problem['DYm']
    dWp, dWm = problem['DWp'], problem['DWm']
    diff = err(W, Y)
    while t < 1e5 and diff > min(1e-5, maxdiff):
        Y = Y * np.power(dYm(W, Y)/dYp(W, Y), eta)
        W = W * np.power(dWm(W, Y)/dWp(W, Y), eta)
        diff = err(W, Y)
        t += 1
    succ = diff < maxdiff
    if verbose:
       print("initialize W:  ", W0)
       print("element-wise difference: \n", X - W @ Y @ W.T)
       print("W:\n", W)
       print("Y:\n", Y)
       print("WYW.T:\n", W @ Y @ W.T)
       print("final update:\n", dWm(W, Y)/dWp(W, Y))

    return {'success':succ, 'W':W, 'Y':Y, 'eps_aprx':W @ Y @ W.T, 'err':diff}

#'nonneg-W' is always a constraint
def quadnmf(X, r, maxdiff, method=nonneg_biconvex, init=None, constraints=[], verbose=False):
    check = is_reducible(X)
    if check['reducible']:
        print('input eps is block diagonal')
        blocks = check['blocks']
        W = np.zeros((len(X), r))
        Y = np.zeros((r, r))
        success = True
        idxW, idxY = 0, 0
        for b in blocks:
            # make sure this tolerance is set to the same as the one in minimization routine
            br = np.linalg.matrix_rank(b, tol=1.e-3)
            sub_res = method(b, br, maxdiff, init, constraints, verbose=verbose)
            W[idxW:idxW+len(b), idxY:idxY+br] = sub_res['W']
            Y[idxY:idxY+br, idxY:idxY+br] = sub_res['Y']
            idxW += len(b)
            idxY += br
            success = success and sub_res['success']
        W = np.linalg.inv(check['P']) @ W
        res = {'success':success,'W':W, 'Y':Y, 'eps_aprx':W @ Y @ W.T, 'err':np.linalg.norm(W@Y@W.T-X, ord='fro')}
    else:
        res = method(X, r, maxdiff, init, constraints, verbose=verbose)

    return res

def solveY(X, W, options={}):
    '''
      used when there is no nonnegative constraint on the monomer-monomer interactions
    '''
    if len(X) == 2:
        return X/W**2
    WW_inv = np.linalg.pinv(W.T @ W)

    return WW_inv @ W.T @ X @ W @ WW_inv

def nonnegY_cvx(X, W):
    '''
      using CVXPY solving for nonnegative Y; slower than the iterative method below
    '''
    Y = cp.Variable((len(W[0]), len(W[0])))
    objective = cp.Minimize(cp.norm(cp.matmul(cp.matmul(W, Y), W.T)-X, 'fro'))
    constraints = [Y >= 0.]
    prob = cp.Problem(objective, constraints)
    prob.solve()
    if prob.status != cp.OPTIMAL:
        raise Exception("Solver did not converge!")

    return Y.value

def nonnegY(X, W, options={}):
    '''
      multiplicative update scheme for ensuring nonnegative Y; equivalent to weighted gradient descent
    '''
    dYp = lambda W, Y: W.T @ W @ Y @ W.T @ W + 1e-10
    dYm = lambda W, Y: W.T @ X @ W
    err = lambda W, Y: np.linalg.norm(W@Y@W.T - X, ord='fro')
    loss = err

    Y = np.ones((len(W[0]), len(W[0])))

    if options:
        sum_function = lambda f, g, w: lambda x, y: f(x, y) + w * g(x, y)
        constraints, weights = options['constraints'], options['weights']
        if len(constraints) > 0:
            for gen_con, w in zip(constraints, weights):
                con = gen_con()
                loss = sum_function(loss, con['loss'], w)
                if con['DYp'] is not None: dYp = sum_function(dYp, con['DYp'], w)
                if con['DYm'] is not None: dYm = sum_function(dYm, con['DYm'], w)

    eta = 1/2
    diff = float('inf')
    t = 0
    while t < 1e2 and diff > 1.e-5:
        Y = Y * np.power(dYm(W, Y)/dYp(W, Y), eta)
        diff = loss(W, Y)
        t += 1

    return Y


def sort_W(W):
    '''
      find the indices of the largest elements in each row and sort the rows
      according to the indices
    '''
    def enrichment(row):
        return row.index(max(row))
    W_rows = [list(row) for row in W]
    # according to the documentation, sorted() is stable
    sorted_W = sorted(W_rows, key=enrichment, reverse=False)

    return np.array(sorted_W, dtype=int)


### regularizations on W
def var_tot():
    DYp = lambda W, Y: 2./len(Y)**2 * Y
    DYm = lambda W, Y: 2./len(Y)**4 * np.sum(Y) * np.ones(Y.shape)
    loss = lambda W, Y: np.var(Y)

    return {'loss': loss, 'DYp':DYp, 'DYm':DYm, 'DWp':None, 'DWm':None}

def var_diag():
    DYp = lambda W, Y: 2./len(Y) * np.diag(np.diag(Y))
    DYm = lambda W, Y: 2./len(Y)**2 * np.sum(np.diag(Y)) * np.eye(Y.shape[0])
    loss = lambda W, Y: np.var(np.diag(Y))

    return {'loss': loss, 'DYp':DYp, 'DYm':DYm, 'DWp':None, 'DWm':None}