harmonization.solver

  1import numpy as np
  2
  3from time import time
  4from itertools import cycle, product
  5
  6from joblib import Parallel, delayed
  7from tqdm import tqdm
  8
  9from sklearn.utils import gen_batches
 10from harmonization.elastic_net import lasso_path, lasso_crossval
 11
 12
 13def update_D(D, A, B, X, niter=1, eps=1e-15, positivity=False):
 14
 15    u = np.zeros(D.shape)
 16    norm2 = np.zeros(D.shape[1])
 17    to_purge = np.zeros(D.shape[1], dtype=bool)
 18
 19    for _ in range(niter):
 20        # divide by zeros are replaced by 0, but still raise the warning, so we squelch it
 21        with np.errstate(divide='ignore', invalid='ignore'):
 22            u[:] = np.where(np.diag(A) != 0, (B - np.dot(D, A)) / np.diag(A) + D, 0)
 23
 24        norm2[:] = np.sqrt(np.sum(u**2, axis=0))
 25
 26        # purge unused atoms from D
 27        to_purge[:] = norm2 < eps
 28
 29        if np.any(to_purge):
 30            indices = np.random.randint(X.shape[0], size=np.sum(to_purge))
 31            u[:, to_purge] = X[indices].reshape(np.sum(to_purge), u.shape[0]).T
 32            norm2[to_purge] = np.sqrt(np.sum(u[:, to_purge]**2, axis=0))
 33
 34        D[:] = u / np.maximum(1, norm2)
 35
 36        if positivity:
 37            D[D < 0] = 0
 38
 39def lasso_path_parallel(D, X, nlambdas, positivity=False, variance=None, fit_intercept=True, standardize=True, use_crossval=False):
 40    N = X.shape[0]
 41    alpha = np.zeros((N, D.shape[1]), dtype=np.float32)
 42    intercept = np.zeros(N, dtype=np.float32)
 43    Xhat = np.zeros((N, D.shape[0]), dtype=np.float32)
 44    lbda = np.zeros(N, dtype=np.float32)
 45
 46    if use_crossval:
 47        # use_crossval holds the number of split, if not we default to 3 because reasons
 48        if use_crossval == 1:
 49            n_splits = 3
 50        else:
 51            n_splits = use_crossval
 52
 53        Xhat[i], alpha[i], intercept[i], lbda[i] = lasso_crossval(D, X[i], nlam=nlambdas, fit_intercept=fit_intercept, n_splits=n_splits,
 54                                                                  pos=positivity, standardize=standardize, penalty=None)
 55    else:
 56        for i in range(N):
 57            alpha[i], intercept[i], Xhat[i], lbda[i] = lasso_path(D, X[i], nlam=nlambdas, fit_intercept=fit_intercept, criterion='aicc',
 58                                                                  pos=positivity, standardize=standardize, penalty=None)
 59    return Xhat, alpha, intercept, lbda
 60
 61
 62def solve_l1(X, D, alpha=None, return_all=False, nlambdas=100, ncores=-1, positivity=False, variance=None, fit_intercept=True, standardize=True,
 63             progressbar=True, use_joblib=False, use_crossval=False, verbose=5):
 64
 65    if alpha is None:
 66        return_alpha = True
 67        alpha = np.zeros((D.shape[1], X.shape[0]))
 68    else:
 69        return_alpha = False
 70
 71    if variance is None:
 72        variance = [None] * alpha.shape[1]
 73
 74    Xhat = np.zeros((alpha.shape[1], D.shape[0]), dtype=np.float32)
 75    intercept = np.zeros((alpha.shape[1], 1), dtype=np.float32)
 76    lbda = np.zeros((alpha.shape[1], 1), dtype=np.float32)
 77
 78    step = 10
 79    arglist = ((D,
 80                X[i:i+step],
 81                nlambdas,
 82                positivity,
 83                variance[i:i+step],
 84                fit_intercept,
 85                standardize,
 86                use_crossval) for i in range(0, alpha.shape[1], step))
 87
 88    plauncher = Parallel(return_as='generator',
 89                         pre_dispatch='all',
 90                         backend='loky',
 91                         n_jobs=ncores)
 92
 93    if use_joblib:
 94        stuff = plauncher(delayed(lasso_path_parallel)(*args) for args in arglist)
 95
 96        if progressbar:
 97            stuff = tqdm(stuff, total=X.shape[0]//step+1)
 98    else:
 99        raise ValueError('Only joblib path is supported now.')
100
101    for i, content in enumerate(stuff):
102        slicer = np.index_exp[i * step:(i+1) * step]
103        slicer2d = np.index_exp[:, i * step:(i+1) * step]
104
105        Xhat[slicer], temp_alpha, temp_intercept, temp_lbda = content
106        alpha[slicer2d] = temp_alpha.T
107        intercept[slicer] = temp_intercept[:, None]
108        lbda[slicer] = temp_lbda[:, None]
109
110    if return_all:
111        return Xhat, alpha, intercept, lbda
112
113    if return_alpha:
114        return alpha
115
116
117def online_DL(X, D=None, n_atoms=None, niter=250, batchsize=128, rho=1.0, t0=1e-3, variance=None,
118              shuffle=True, fulldraw=False, positivity=False, fit_intercept=True, standardize=True, ncores=-1, nlambdas=100,
119              progressbar=True, use_joblib=True, eps=1e-6):
120
121    tt = time()
122    seen_patches = 0
123
124    if n_atoms is None:
125        n_atoms = 2 * np.prod(X.shape[1:])
126
127    if n_atoms > X.shape[0]:
128        n_atoms = X.shape[0]
129
130    if D is None:
131        # we pick n_atoms indexes along each dimensions randomly
132        indices = np.random.permutation(X.shape[0])[:n_atoms]
133        D = X[indices].reshape(n_atoms, np.prod(X.shape[1:])).T
134
135    # Put to unit l2 norm, will also copy D if we passed it to the function
136    D = D / np.sqrt(np.sum(D**2, axis=1, keepdims=True))
137
138    A = np.eye(D.shape[1]) * t0
139    B = t0 * D
140
141    A_prime = np.copy(A)
142    B_prime = np.copy(B)
143
144    x = np.zeros((batchsize, D.shape[0]))
145    alpha = np.zeros((n_atoms, batchsize))
146
147    alpha_alpha_T = np.zeros((n_atoms, n_atoms))
148    x_alpha_T = np.zeros((x.shape[1], n_atoms))
149
150    if shuffle:
151        np.random.shuffle(X)
152
153    if fulldraw:
154        # one iteration goes through all elements
155        batches = gen_batches(X.shape[0], batchsize)
156        iterator = product(range(1, niter + 1), batches)
157    else:
158        # one iteration goes through one draw of size batchsize from X
159        batches = cycle(gen_batches(X.shape[0], batchsize))
160        iterator = zip(range(1, niter + 1), batches)
161
162    for t, batch in tqdm(iterator, total=niter):
163
164        cutter = X[batch].shape[0]
165        if cutter < batchsize:
166            continue
167
168        x[:] = X[batch].reshape(batchsize, -1)
169        _, _, _, lbda = solve_l1(x, D, alpha, positivity=positivity, ncores=ncores, nlambdas=nlambdas, variance=variance, return_all=True,
170                                 fit_intercept=fit_intercept, standardize=standardize, use_joblib=use_joblib, progressbar=False, verbose=False)
171
172        np.dot(alpha, alpha.T, out=alpha_alpha_T)
173        # x is transposed with regards to original notation
174        np.dot(x.T, alpha.T, out=x_alpha_T)
175
176        beta = (1 - (1 / t))**rho
177
178        A *= beta
179        A += alpha_alpha_T / batchsize
180        A_prime *= beta
181        A_prime += alpha_alpha_T / batchsize
182
183        B *= beta
184        B += x_alpha_T
185        B_prime *= beta
186        B_prime += x_alpha_T
187
188        update_D(D, A, B, X, positivity=False)
189
190        # reset past information every two full epochs
191        if seen_patches > 2 * X.shape[0]:
192            A[:] = A_prime
193            B[:] = B_prime
194
195            A_prime[:] = 0
196            B_prime[:] = 0
197
198            seen_patches = 0
199        else:
200            seen_patches += batchsize
201
202    print(f'total time: {round((time() - tt) / 60, 2)} mins')
203    return D