function [A_opt,L_opt,Res] = varMicro(X,K_range,opts)
%VARMICRO Variational Microstates.
% Implementation of variational microstates, with optional smoothing of
% segments as described in [1].
% Free energy is used to choose amongst multiple restarts, for the same
% number of microstates (K). Free energy is also used to find the optimal
% number of microstates.
%
% [1] - (unpublished manuscript). Variational microstate analysis.
% [2] - Pascual-Marqui, R. D., Michel, C. M., & Lehmann, D. (1995).
% Segmentation of brain electrical activity into microstates: model
% estimation and validation. IEEE Transactions on Biomedical
% Engineering.
%
% Please cite this toolbox as:
% Poulsen, A. T., Pedroni, A., Langer, N., & Hansen, L. K. (2018).
% Microstate EEGlab toolbox: An introductionary guide. bioRxiv.
%
% Andreas Trier Poulsen, atpo@dtu.dk
% Technical University of Denmark, Cognitive systems - February 2017
%
% Inputs
% X - EEG (channels x samples).
% K_range - A priori number of microstates. Can be a vector.
%
% Optional inputs
% opts.
% sig2_0 - Prior variance of activations (default: average
% EEG channel variance).
% p0 - Probability for having same microstate as last
% timepoint (default: 0). Setting to zero turns
% smoothing off.
% reps - Number of random initialisations of algorithm
% (default: 10).
% max_iterations - Maximum number of iterations of algorithm
% (default: 1000).
% thresh - Threshold of convergence based on relative change
% in free energy (default: 1e-6).
% verbose - Print status messages to command window?
% (default:1).
%
% Outputs
% A_opt - Spatial distribution of microstates (channels x K).
% L_opt - Label of the most active microstate at each timepoint (1 x samples).
% Res - Struct containing additional results:
% .Z_all - Cell containing the microstate activations for each number of
% microstates defined in K_range. The dimensions of each cell
% is (K x samples).
% .A_all - Cell containing the spatial distribution for each number of
% microstates defined in K_range. The dimensions of each cell
% is (channels x K).
% .L_all - Cell containing the labels for each number of microstates
% defined in K_range. The dimensions of each cell is
% (1 x samples).
% .S_opt - Sparsity parameters (K x samples) of optimal solution.
% .beta - Noise precision for the best solution of each K in K_range.
% .sig2_z_opt - Variance of microstate activations (K x samples) of
% optimal solution.
% .K_act - Number of microstates for optimal solution.
% .Free_energy - Free energy for the best solution of each K in K_range.
% .sig2_modk - Noise variance as calculated in [2] for the best
% solution of each K in K_range.
% .sig2_modk_mcv - Modified predictive residual variance as calculated
% in [2] for the best solution of each K in K_range.
% .R2 - Explained variance as calculated in [2] for the best
% solution of each K in K_range.
% .MSE - Mean squared error for the best solution of each K in K_range.
% .opts - Settings used to run the algorithm.
%% Checking inputs and setting default settings, where needed.
if nargin<3; opts = []; end
if nargin<2; K_range = 3:8; end
if ~isfield(opts,'reps'), opts.reps = 10; end
if ~isfield(opts,'max_iterations'), opts.max_iterations = 1000; end
if ~isfield(opts,'verbose'), opts.verbose = 1; end
if ~isfield(opts,'thresh'), opts.thresh = 1e-6; end
if ~isfield(opts,'sig2_0'), opts.sig2_0 = mean(var(X,0,2)); end
if ~isfield(opts,'p0'), opts.p0 = 0; end
if opts.p0==0 % Checking if smoothing is requested
opts.smooth = 0;
else
opts.smooth = 1;
end
%% Preallocating
N_K = length(K_range);
A_all = cell(N_K,1);
L_all = cell(N_K,1);
Z_all = cell(N_K,1);
R2_all = nan(N_K,1);
sig2_modk_all = nan(N_K,1);
sig2_modk_mcv_all = nan(N_K,1);
MSE_all = nan(N_K,1);
beta_all = nan(N_K,1);
F_all = nan(N_K,1);
%% Looping over all K values in K_range
F_opt = inf;
K_ind = 0;
if opts.verbose
if opts.smooth
disp('Starting variational microstates with temporal smoothing')
else
disp('Starting variational microstates without temporal smoothing')
end
end
for K = K_range
K_ind = K_ind + 1;
if opts.verbose
fprintf('Analysis no. %i out of %i. Starting %i random initialisations for %i microstates.\n'...
,K_ind,N_K,opts.reps, K)
end
% Finding best fit amongst a given number of restarts
F_best = inf;
for r = 1:opts.reps
if opts.verbose
fprintf('Starting initialisations no. %i out of %i. '...
,r,opts.reps)
end
% Segmenting with variational microstates
[A,Z,sig2_z,S,beta,F,ind,sig2_modk,sig2_modk_mcv,R2,MSE] = segmentation(X,K,opts);
if opts.verbose
fprintf('Finished in %i iterations.\n',ind)
end
% Checking for best fit
if F < F_best
A_all{K_ind} = A;
Z_all{K_ind} = Z;
% Calculating label of active microstate for each timepoint for
% best solution
[~,L_all{K_ind}] = max(S);
sig2_z_best = sig2_z;
S_best = S;
beta_all(K_ind) = beta;
F_best = F;
sig2_modk_all(K_ind) = sig2_modk;
sig2_modk_mcv_all(K_ind) = sig2_modk_mcv;
R2_all(K_ind) = R2;
MSE_all(K_ind) = MSE;
end
end
F_all(K_ind) = F_best;
% Finding optimum solution amongst different values of K
if F_all(K_ind) < F_opt
A_opt = A_all{K_ind};
L_opt = L_all{K_ind};
sig2_z_opt = sig2_z_best;
S_opt = S_best;
K_act = K;
F_opt = F_all(K_ind);
end
end
%% Saving additional results and info to Res struct
Res.Z_all = Z_all;
Res.A_all = A_all;
Res.L_all = L_all;
Res.S_opt = S_opt;
Res.beta = beta_all;
Res.sig2_z_opt = sig2_z_opt;
Res.K_act = K_act;
Res.Free_energy = F_all;
Res.sig2_modk = sig2_modk_all;
Res.sig2_modk_mcv = sig2_modk_mcv_all;
Res.R2 = R2_all;
Res.MSE = MSE_all;
Res.opts = opts;
end
function [A,Z,sig2_Z,S,beta,F,ind,sig2_modk,sig2_modk_mcv,R2,MSE] = segmentation(X,K,opts)
% Updating variables for variational microstates in iterations.
% Dimensions of variables: X=(C x N), Z=(K x N), sig2_Z=(K x N), A=(C x K),
% S=(K x N).
%% Initialising
[C,N] = size(X);
F = inf;
F_old = 0;
% Reading settings
max_iterations = opts.max_iterations;
thresh = opts.thresh;
sig2_0 = opts.sig2_0;
p0 = opts.p0;
tmp = (1-p0)/(K-1);
smooth_const = log(1 + (p0 - tmp)/tmp);
% random initialisation of A, beta, and S
A = X(:,randperm(N,K)); % selecting K random timepoints to use as initial microstate maps
A = bsxfun(@rdivide,A,sqrt(sum(A.*A)));% normalising
beta = 1/mean(var(X,0,2));
S = rand(K,N)*0.1; % giving all states a non-zero probability
[~,L] = max((A'*X).^2); % finding most active microstates based on A
for n = 1:N, S(L(n),n) = 1; end % setting these to 1
S = bsxfun(@rdivide,S,sum(S,1));% normalising so each n sum to 1
%% Iterate over updates until convergence
ind = 0;
while abs(F_old-F) >= thresh*F && max_iterations>ind
ind = ind + 1;
F_old = F;
% Z (mu_z)
denom = 1/(sig2_0) + beta*bsxfun(@times, S,sum(A.^2)');
Z = beta*(S.*(A'*X)) ./ denom;
if isnan(Z), error('hov'), end
% sig2_z
sig2_Z = 1./denom;
if isnan(sig2_Z), error('hov'), end
% S (mu_s) with optional temporal smoothing
if opts.smooth
S_past = [zeros(K,1) S(:,1:end-1)];
S_future = [S(:,2:end) zeros(K,1)];
end
S = beta*( Z.*(A'*X) - 0.5*bsxfun(@times, Z.^2+sig2_Z, sum(A.^2)') );
if isnan(S), error('hov'), end
if opts.smooth
S = S + smooth_const*(S_past+S_future);
end
S = safe_softmax(S);
if isnan(S), error('hov'), end
% A
A = bsxfun(@rdivide, X*(Z.*S)', sum((Z.^2+sig2_Z).*S,2)');
if isnan(A), error('hov'), end
% beta
beta = C*N / sum(sum( X.^2 - 2*X.*(A*(Z.*S)) + A.^2*((Z.^2+sig2_Z).*S) ));
if isnan(beta), error('hov'), end
% Free energy
F = free_energy(Z,sig2_Z,sig2_0,S,beta,C,N,K);
end
%% Calculating measures of fit
% Noise variance and performance measures as calculated in Pascual-Marqui (1995)
const = sum(sum(X.^2));
[~,L] = max(Z.^2); %finding most active microstates
sig2_modk = (const - sum(sum(A(:,L).*X).^2)) / (N*(C-1));
sig2_modk_mcv = sig2_modk * ((C-1)^-1 * (C-1-K))^-2;
sig2_D = const / (N*(C-1));
R2 = 1 - sig2_modk/sig2_D;
% Mean squared error
MSE = mean(mean((X-A*(Z.*S)).^2));
end
function X = safe_softmax(X)
% Computes a safe softmax, to avoid troubles with numerical precision, by
% subtracting largest values before exponentiation using exp(x) =
% exp(x-a)*exp(a).
% Input
% X (K,N) log-probs
% Output
% X (K,N) softmax'ed values
% Subtracting max values for each n
A = max(X,[],1);
X = bsxfun(@minus,X,A);
% Softmaxing
X = exp(X);
X = bsxfun(@times,X,1./sum(X,1));
% Avoid NaNs when calculating log(X)
X(X==0) = eps;
end
function F = free_energy(Z,sig2,sig2_0,S,beta,C,N,K)
% Calculate free energy of solution. beta needs to be last variable
% calculated.
F1 = -0.5*(N*K*log(2*pi)+sum(sum(log(sig2)))) -N*K/2 + sum(sum(S.*log(S)));
F2 = 0.5*K*N*log(2*pi*sig2_0) + sum(sum(Z.^2+sig2))/(2*sig2_0) + N*log(K);
F3 = -C*N/2*log(beta/(2*pi)) + C*N/2; % using expression for beta in last part
F = F1+F2+F3;
if isnan(F), error('hov'), end
end