nmf_program.py

import math, argparse, gzip, pickle
import cvxpy as cp
import numpy as np
from contextlib import ExitStack
from nmf import *

def update_parameter(observations, rowsum):
    if len(observations) < 1:
        return
    else:
        data = observations[0]
        cnts = np.zeros(data.shape)
        for ob in observations:
            cnts = cnts + ob
        return cnts/len(observations)/rowsum[:, None]


def crossentropy(X, r, L, method_getY=solveY, err_tol=1e-4, \
	nsample=1e5, maxiter=50, kappa=0, prior_option={}, solver_options={}, \
        verbose=False, saveproc=None, random_seed=None):

    constraints, weights = solver_options['constraints'], solver_options['weights']
    err = lambda W, Y: np.linalg.norm(W@Y@W.T - X, ord='fro')
    loss = err
    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)

    nsample = int(nsample)
    N = len(X)
    if random_seed is not None:
        rng = np.random.default(random_seed)
    else:
        rng = np.random.default_rng()

    def solve(X, r, Ls):
        assert len(X) == len(Ls), "mismatched eps dimension and constraint dimension"
        if r <= 1:
            min_W = np.array(Ls)
            min_Y = method_getY(X, min_W, options=solver_options)
            return {'W' : min_W, 'Y' : min_Y, 'err' : 0, 'loss' : 0}

        W0 = np.ones((N, r))
        percentile = 0.001
        nelite = int(nsample * percentile)
        p0 = np.full((N, r+1), 1/(r+1)) # relaxing the equality constraint on L
        print('prior:', prior_option)
        if prior_option:
            distr, parameters, size = prior_option['distribution'], prior_option['args'], prior_option['size']
            for i in range(N):
                if parameters:
                    p0[i] = distr(parameters, size=size)
                else:
                    p0[i] = distr(size=size)
        ps = p0 / np.sum(p0, axis=1, keepdims=1)

        stp = 0
        min_W = np.array(W0)
        min_err = min_loss = min_delta = math.inf
        ptraj, sampletraj = [], []
        top_k_solutions = []
        print("p0 = \n", ps)
        while stp < maxiter and min_err > 1e-4:
            candidates, errs = [], []
            for t in range(nsample):
                W_dummy = np.zeros((N, r+1))
                for j in range(N):
                    W_dummy[j,:] = rng.multinomial(Ls[j], ps[j,:], size=1)
                W_dummy = sort_W(W_dummy)
                W = np.array(W_dummy[:, 0:r])
                Y = method_getY(X, W, options=solver_options)

                current_loss = loss(W, Y)
                candidates.append([current_loss, err(W, Y), stp, W_dummy])
                errs.append(current_loss)

            prev = np.array(ps)

            if type(saveproc) == str:
                ptraj.append(ps.flatten())
                sampletraj.append(list(errs))

            top_candidates = sorted(candidates, key=lambda x:x[0])[:nelite]
            top_W = [cand[-1] for cand in top_candidates]
            ps[:, :] = update_parameter(top_W, Ls)
            min_loss = min(errs)
            if min_loss < min_delta:
                min_delta = min_loss
                min_W = np.array(top_W[0][:, 0:r])
            stp += 1

            if max(list(abs(prev - ps).flatten())) < 0.005:
                print("stop with min change: ", max(list(abs(prev - ps).flatten())))
                break

            if verbose:
                print("top candidates: loss, err, regular, var(diag)", flush=True)
                for i in range(5):
                    print(i, top_candidates[i][0], top_candidates[i][1], top_candidates[i][0] - top_candidates[i][1], flush=True)
                print(ps)
                print('------------------ %d' %stp, flush=True)

        min_Y = method_getY(X, min_W, options=solver_options)
        min_err = err(min_W, min_Y)
        min_loss = loss(min_W, min_Y)
        if type(saveproc) == str:
            np.savetxt("%s_ps.out" %saveproc, np.array(ptraj))
            np.savetxt("%s_samples.out" %saveproc, np.array(sampletraj))

        return {'W' : min_W, 'Y' : min_Y, 'err' : min_err, 'loss' : min_loss}

    check = is_reducible(X)
    if check['reducible']:
        print('input eps is block diagonal')
        blocks = check['blocks']
        W = np.zeros((N, r))
        Y = np.zeros((r, r))
        idxW, idxY = 0, 0
        for b in blocks:
            br = np.linalg.matrix_rank(b, tol=1e-6)
            sub_res = solve(b, br, L[idxW:idxW+len(b)])
            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
        W = np.linalg.inv(check['P']) @ W
        res = {'W':W, 'Y':Y, 'err':np.linalg.norm(W@Y@W.T-X, ord='fro'), 'loss':loss(W, Y)}
    else:
        res = solve(X, r, L)

    return res


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('path', type=str, help="path to eps solutions of targets .p.gz file")
    parser.add_argument('--posY', action='store_true', help="solveY(no constraint) / nonnegY")
    parser.add_argument('--rank', type=int, default=0, help="prescribed rank (number of monomer species)")
    parser.add_argument('--reg', type=str, default=None, help="regularization option on the monomer interactions [None]")
    parser.add_argument('--kappa', type=float, default=0., help="regularization weight")
    parser.add_argument('--scankappa', action='store_true', help="scan regularization weight [False]")
    parser.add_argument('--verbose', action='store_true', help="turn on verbose output [False]")
    parser.add_argument('--prior', type=str, default=None, help="prior distribution for sampling W [None]")
    parser.add_argument('--epstol', type=float, default=1.e-3, help="tol on rank of epsilon")
    parser.add_argument('--saveproc', type=str, default=None, help="save the output of ce iterations")
    parser.add_argument('--output', type=str, default=None, help="output .p.gz path to save results [None]")

    clargs = parser.parse_args()
    with gzip.open(clargs.path, 'rb') as f:
        try:
            results = pickle.load(f)
        except EOFError:
            pass
        targets = results['targets']
        N = results['targets'].shape[1]
        x = np.zeros(results['targets'].shape)
        for alpha in range(results['targets'].shape[0]):
            x[alpha,:] = results['targets'][alpha,:] / results['targets'][alpha,:].sum()
        L = results['L'] * np.ones(N)
        try:
            eps_r = results['eps'].low_rank_approx(results['r'])[2]
            r = results['r']
            r = min(r, np.linalg.matrix_rank(eps_r, tol=clargs.epstol))
        except:
            try:
                eps_r = results['eps']
                r = results['r']
            except:
                eps_r = results['eps_r']
                r = np.linalg.matrix_rank(results['eps_r'], tol=1e-4)

        print("targets:\n", targets)
        print("loaded rank = ", r, "np rank = ", np.linalg.matrix_rank(eps_r, clargs.epstol))
        print("N = %d, r = %d" %(N, r))
        print("-eps_r = :\n", -eps_r)
        if clargs.rank != 0:
            print("rank(eps)=", r)
            print("rank to use for W =", clargs.rank)
            r = clargs.rank

        prior_options = {}
        distr_rng = np.random.default_rng()
        if clargs.prior == 'dirichlet':
            prior_options['distribution'] = distr_rng.dirichlet
            prior_options['args'] = tuple(np.ones(r+1) / (r+1))
            prior_options['size'] = 1

        if clargs.prior == 'random':
            prior_options['distribution'] = distr_rng.random
            prior_options['args'] = None
            prior_options['size'] = r + 1

        if 'diag' in clargs.reg:
            reg_option = var_diag
        elif 'tot' in clargs.reg:
            reg_option = var_tot
        else:
            reg_option = None

        list_kappa = [clargs.kappa]
        if clargs.scankappa:
            list_kappa = np.array(list_kappa + list(np.logspace(-4, 1, 10)))

        with ExitStack() as stack:
            if clargs.output is not None:
                foutput = stack.enter_context(gzip.open(clargs.output, 'wb'))

            for kappa in list_kappa:
                ops = {'constraints':[reg_option], 'weights':[kappa]}

                if clargs.posY:
                    eps_test = np.copy(eps_r)
                    eps_test[eps_test > 0] = -1.e-6
                    res = crossentropy(-eps_test, r, L, nonnegY, err_tol=1.e-3, nsample=1.e4,\
                          prior_option=prior_options, solver_options=ops, verbose=True, saveproc=clargs.saveproc)
                else:
                    res = crossentropy(-eps_r, r, L, solveY, err_tol=1e-3, nsample=1e4,\
                            prioe_options=prior_options, solver_options=ops, verbose=True, saveproc=clargs.saveproc)

                print('kappa: ', kappa, ' min err: ', res['err'], ' min loss: ', res['loss'])
                print('W:\n', res['W'])
                print('Y:\n', res['Y'])

                W, Y, error, loss =  res['W'], res['Y'], res['err'], res['loss']
                if clargs.output is not None:
                    pickle.dump({'eps_target': eps_r, 'W':W, 'Y':Y, 'err':error, 'loss': loss, 'kappa':kappa}, foutput)
                    foutput.flush()