diff --git a/ica/ica_torch.py b/ica/ica_torch.py new file mode 100644 index 0000000..c2a80da --- /dev/null +++ b/ica/ica_torch.py @@ -0,0 +1,278 @@ +''' +Independent Component Analysis (ICA): +This script computes ICA using the INFOMAX criteria. +The preprocessing steps include demeaning and whitening. +''' +import torch +import torch.linalg as tla +import numpy as np + +# Global constants +EPS = 1e-16 +MAX_W = 1e8 +ANNEAL = 0.9 +MAX_STEP = 500 +MIN_LRATE = 1e-6 +W_STOP = 1e-6 + + +def norm(x): + """Computes the norm of a vector or the Frobenius norm of a + matrix_rank + + """ + return torch.norm(x.ravel()) + + +class ica: + + def __init__(self, n_components=10): + self.n_comp = n_components + + def fit(self, x2d): + x_white, self.white, self.dewhite\ + = pca_whiten(x2d, self.n_comp) + self.mix, self.sources, self.unmix\ + = infomax1(x_white, self.n_comp) + return self.mix, self.sources + + +def diagsqrts(w): + """ + Returns direct and inverse square root normalization matrices + """ + Di = torch.diag(1. / (torch.sqrt(w) + torch.finfo(float).eps)) + D = torch.diag(torch.sqrt(w)) + return D, Di + + +def pca_whiten(x2d, n_comp, verbose=True): + """ data Whitening + *Input + x2d : 2d data matrix of observations by variables + n_comp: Number of components to retain + *Output + Xwhite : Whitened X + white : whitening matrix (Xwhite = np.dot(white,X)) + dewhite : dewhitening matrix (X = np.dot(dewhite,Xwhite)) + """ + x2d_demean = x2d - x2d.mean(axis=1).reshape((-1, 1)) + NSUB, NVOX = x2d_demean.shape + if NSUB > NVOX: + cov = torch.matmul(x2d_demean.T, x2d_demean) / (NSUB - 1) + w, v = torch.eigh(cov, eigvals=(NVOX - n_comp, NVOX - 1)) + D, Di = diagsqrts(w) + u = torch.matmul(torch.matmul(x2d_demean, v), Di) + x_white = v.T + white = torch.matmul(Di, u.T) + dewhite = torch.matmul(u, D) + else: + cov = torch.matmul(x2d_demean, x2d_demean.T) / (NVOX - 1) + w, u = tla.eigh(cov) + w = w[(NSUB - n_comp):] + u = u[:, (NSUB - n_comp):] + D, Di = diagsqrts(w) + white = torch.matmul(Di, u.T) + x_white = torch.matmul(white, x2d_demean) + dewhite = torch.matmul(u, D) + return (x_white, white, dewhite) + + +def w_update(weights, x_white, bias1, lrate1): + """ Update rule for infomax + This function recieves parameters to update W1 + * Input + W1: unmixing matrix (must be a square matrix) + Xwhite1: whitened data + bias1: current estimated bias + lrate1: current learning rate + startW1: in case update blows up it will start again from startW1 + * Output + W1: updated mixing matrix + bias: updated bias + lrate1: updated learning rate + """ + device = weights.device + NCOMP, NVOX = x_white.shape + block1 = int(np.floor(np.sqrt(NVOX / 3))) + permute1 = torch.randperm(NVOX) + for start in range(0, NVOX, block1): + if start + block1 < NVOX: + tt2 = start + block1 + else: + tt2 = NVOX + block1 = NVOX - start + + unmixed = torch.matmul( + weights, x_white[:, permute1[start:tt2]]) + bias1 + logit = 1 - (2 / (1 + torch.exp(-unmixed))) + weights = weights + lrate1 * torch.matmul(block1 * torch.eye(NCOMP, device=device) + + torch.matmul(logit, unmixed.T), weights) + bias1 = bias1 + lrate1 * logit.sum(axis=1).reshape(bias1.shape) + # Checking if W blows up + if (torch.isnan(weights)).any() or torch.max(torch.abs(weights)) > MAX_W: + print("Numeric error! restarting with lower learning rate") + lrate1 = lrate1 * ANNEAL + weights = torch.eye(NCOMP, device=device) + bias1 = torch.zeros((NCOMP, 1), device=device) + error = 1 + + if lrate1 > 1e-6 and \ + tla.matrix_rank(x_white) < NCOMP: + print("Data 1 is rank defficient" + ". I cannot compute " + + str(NCOMP) + " components.") + return (None, None, None, 1) + + if lrate1 < 1e-6: + print("Weight matrix may" + " not be invertible...") + return (None, None, None, 1) + break + else: + error = 0 + + return (weights, bias1, lrate1, error) + +# infomax1: single modality infomax + + +def infomax1(x_white, verbose=False): + """Computes ICA infomax in whitened data + Decomposes x_white as x_white=AS + *Input + x_white: whitened data (Use PCAwhiten) + verbose: flag to print optimization updates + *Output + A : mixing matrix + S : source matrix + W : unmixing matrix + """ + device = x_white.device + NCOMP = x_white.shape[0] + # Initialization + weights = torch.eye(NCOMP, device=device) + old_weights = torch.eye(NCOMP, device=device) + d_weigths = torch.zeros(NCOMP, device=device) + old_d_weights = torch.zeros(NCOMP, device=device) + lrate = 0.005 / np.log(NCOMP) + bias = torch.zeros((NCOMP, 1), device=device) + change = 1 + angle_delta = 0 + if verbose: + print("Beginning ICA training...") + step = 1 + + while step < MAX_STEP and change > W_STOP: + + (weights, bias, lrate, error) = w_update(weights, x_white, bias, lrate) + + if error != 0: + step = 1 + error = 0 + lrate = lrate * ANNEAL + weights = torch.eye(NCOMP, device=device) + old_weights = torch.eye(NCOMP, device=device) + d_weigths = torch.zeros(NCOMP, device=device) + old_d_weights = torch.zeros(NCOMP, device=device) + bias = torch.zeros((NCOMP, 1), device=device) + else: + d_weigths = weights - old_weights + change = norm(d_weigths)**2 + + if step > 2: + angle_delta = torch.arccos( + torch.sum(d_weigths * old_d_weights) / + (norm(d_weigths) * norm(old_d_weights) + 1e-8) + ) * 180 / np.pi + + old_weights = weights.clone() + + if angle_delta > 60: + lrate = lrate * ANNEAL + old_d_weights = d_weigths.clone() + elif step == 1: + old_d_weights = d_weigths.clone() + + if verbose and change < W_STOP: + print("Step %d: Lrate %.1e," + "Wchange %.1e," + "Angle %.2f" % (step, lrate, + change, angle_delta)) + + step = step + 1 + + # A,S,W + return (tla.inv(weights), torch.matmul(weights, x_white), weights) + +# Single modality ICA + + +def ica1(x_raw, ncomp, verbose=False): + ''' + Single modality Independent Component Analysis + ''' + device = x_raw.device + if verbose: + print("Whitening data...") + x_white, _, dewhite = pca_whiten(x_raw, ncomp) + if verbose: + print('x_white shape: %d, %d' % x_white.shape) + print("Done.") + if verbose: + print("Running INFOMAX-ICA ...") + mixer, sources, unmixer = infomax1(x_white, verbose) + mixer = torch.matmul(dewhite, mixer) + + scale = sources.std(axis=1).reshape((-1, 1)) + sources = sources / scale + scale = scale.reshape((1, -1)) + mixer = mixer * scale + + if verbose: + print("Done.") + return (mixer, sources, unmixer) + + +def icax(x_raw, ncomp, verbose=True): + device = x_raw.device + if verbose: + print("Whitening data...") + x_white, _, dewhite = pca_whiten(x_raw, ncomp) + + mixer_list = [] + sources_list = [] + for it in range(10): + if verbose: + print('Run number %d' % it) + print("Running INFOMAX-ICA ...") + mixer, sources, unmix = infomax1(x_white, verbose) + mixer_list.append(mixer) + sources_list.append(sources) + + # Reorder all sources to the order of the first + S1 = sources_list[0] + for it in range(1, 10): + S2 = sources_list[it] + A2 = mixer_list[it] + cor_m = torch.corrcoef(S1, S2)[:ncomp, ncomp:] + idx = torch.argmax(np.abs(cor_m), axis=1) + S2 = S2[idx, :] + A2 = A2[:, idx] + cor_m = torch.corrcoef(S1, S2)[:ncomp, ncomp:] + S2 = S2 * torch.sign(torch.diag(cor_m)).reshape((ncomp, 1)) + A2 = A2 * torch.sign(torch.diag(cor_m)).reshape((1, ncomp)) + sources_list[it] = S2 + mixer_list[it] = A2 + + # Average sources + temp_sources = torch.zeros(sources.shape, device=device) + temp_mixer = torch.zeros(mixer.shape, device=device) + for sources, mixer in zip(sources_list, mixer_list): + temp_sources = temp_sources + sources + temp_mixer = temp_mixer + mixer + + temp_sources = temp_sources / 10.0 + temp_mixer = temp_mixer / 10.0 + + return (temp_mixer, temp_sources)