classdef ParadigmPAL < ParadigmBaseSimplified
    % Experimental implementation of the Pattern Alignment Learning (PAL) method.
    %
    % This paradigm is not yet completed, please move on! :-)
    %
    % Name:
    %   Wave Alignment, work in progress
    %
    %                            Christian Kothe, Swartz Center for Computational Neuroscience, UCSD
    %                            2011-08-28

    methods

        function defaults = preprocessing_defaults(self)
            % define the default pre-processing parameters of this paradigm
            defaults = {'SpectralSelection',[0.1 35],'EpochExtraction',[-0.5 1.5], 'Resampling',100};
        end

        function model = calibrate_simple(self,varargin)
            % Override this function to implement your calibration code
            % Model = calibrate_simple(Signal,Arguments...)
            %
            % In:
            %   Signal : a single continuous EEGLAB data set (annotated with target markers; see
            %            set_targetmarkers for more info)
            %
            %   Arguments : further optional user arguments
            %
            % Out:
            %   Model : a model struct with a mandatory field .filter_graph and arbitrary other content
            %           * The .filter_graph filed is a 1x1 cell array that contains the desciption of
            %             filter steps that is to be applied to the data, and is usually passed as either
            %             the .tracking.online_expression field of the processed Signal or the processed
            %             Signal itself. If no signal processing is performed by this paradigm, the raw
            %             signal may be passed.
            %
            %           * May have optional fields .prediction_function, .prediction_window and
            %             .prediction_channels - though these are generally auto-deduced
            %
            % Notes:

            args = arg_define(varargin, ...
                arg_norep('signal'), ...
                arg_sub({'flt','SignalProcessing'}, self.preprocessing_defaults(), @flt_pipeline, 'Signal processing stages. These parameters control filter stages that run on the signal level; they can be enabled, disabled and configured for the given paradigm. The prediction operates on the outputs of this stage.'), ...
                arg({'lambdas','Lambdas'},[2.^(8:-0.125:-5)],[0 Inf],'Regularization parameter. Larger means stronger regularization, 1.0 is the Bayesian ad hoc choice.'), ...
                arg({'duration','WaveDuration'},0.8,[0 Inf],'Duration of Waveform. This is the length of the waveform that will be searched in the epoch.'), ...
                arg({'stepping','Stepping'},3,uint32([1 100]),'Relative jitter stepping. This is the precision (in samples) at which relative jitters will be considered.'), ...
                arg({'admm_iters','ADMMIterations'},30,uint32([1 1000]),'Max Iterations. Number of iterations for the overall ADMM solver.'), ...
                arg({'inner_iters','InnerIterations'},10,uint32([1 1000]),'Max Inner Iterations. Number of iterations for the inner ADMM solver that solves the sparsity vs trace norm.'), ...
                arg({'cg_iters','CgIterations'},30,uint32([1 1000]),'Max CG Iterations. Number of iterations for the inner conjugate-gradient solver.'), ...
                arg({'lanczos_iters','LanczosIterations'},10,uint32([1 1000]),'Max Lanczos Iterations. Number of iterations for the inner Lanczos solver.'), ...
                arg({'abs_tol','AbsoluteTolerance'},1e-6,[],'Absolute tolerance.'), ...
                arg({'rel_tol','RelativeTolerance'},1e-3,[],'Relative tolerance.'), ...
                arg({'alpha','OverRelaxation'},1,[1 1.8],'Over-relaxation parameter. Some references suggest that values between 1.5 and 1.8 can improve convergence.'), ...
                arg({'rho','AugmentedLagrangian'},1.0,[0 Inf],'Augmented Lagrangian parameter.'), ...
                arg({'rho_update','RhoUpdate'},true,[],'Update Rho. Whether to update rho dynamically according to 3.4.1 in [1]'), ...
                arg({'rho_cutoff','RhoUpdateThreshold'},10.0,[0 Inf],'Rho update threshold.'), ...
                arg({'rho_incr','RhoUpdateIncr'},2.0,[1 Inf],'Rho update increment factor.'), ...
                arg({'rho_decr','RhoUpdateDecr'},2.0,[1 Inf],'Rho update decrement factor.'), ...
                arg({'kappa','SparsityTraceTradeoff'},0.75,[0 1],'Sparsity (=1) vs. Trace norm (=0) tradeoff. Between 0 and 1, similar to the elastic net criterion.'), ...
                arg({'patterns','NumPatterns'},20,uint32([1 1000]),'Number of weight patterns. This is the maximum number of (possibly distinct) weight patterns to recover.'), ...
                arg({'verbose','Verbose'},'really',[],'Verbose output.'), ...
                arg({'scaling','Scaling'}, 'std', {'none','center','std','minmax','whiten'}, 'Pre-scaling of the data. For the regulariation to work best, the features should either be naturally scaled well, or be artificially scaled.'), ...
                arg({'fast','Fast'},true,[],'Fast mode.'), ...
                arg({'arg_dialogsel','ConfigLayout'},self.dialog_layout_defaults(),[],'Parameters displayed in the config dialog. Cell array of parameter names to display (dot-notation allowed); blanks are translated into empty rows in the dialog. Referring to a structure argument lists all parameters of that struture, except if it is a switchable structure - in this case, a pulldown menu with switch options is displayed.','type','cellstr','shape','row'));

            % first pre-process the data (symbolically)
            signal = flt_pipeline('signal',args.signal, args.flt); %#ok<*NODEF>

            % evaluate this in an optimized fashion
            signal = exp_eval_optimized(signal);

            % get trial labels
            labels = set_gettarget(signal);

            % build meta-data
            model.args = rmfield(args,'signal');
            model.classes = unique(labels,'rows');
            model.tracking.filter_graph = signal;

            if length(model.classes) > 2
                error('Currently, only two classes are supported here.'); end

            % extract time-shifted windows from each trial...
            duration = round(args.duration * signal.srate);
            offsets = 1 : args.stepping : (signal.pnts - duration);
            trials = cell(1,length(offsets));

            % get the entire epoch and figure out an appropriate rescaling
            alldata = reshape(permute(signal.data,[2 1 3]),[],signal.trials);
            sc_info = hlp_findscaling(alldata,args.scaling);

            % for each time shift...
            for o = offsets
                % determine the window to cut out of the data
                cutwindow = o + (1:duration);
                % extract from all trials
                data = signal.data(:,cutwindow,:);
                % vectorize into #trials x #features, and rescale
                data = permute(data,[2 1 3]);
                trials{o} = hlp_applyscaling(reshape(data,[],signal.trials),sc_info)';
            end

            T = signal.trials;
            fprintf('#trials=%.0f; #shifts=%.0f; #features=%.0f; #channels=%.0f; #windows=%.0f\n',T,length(offsets),signal.nbchan*length(cutwindow),signal.nbchan,length(cutwindow));

            % vertically concatenate all shifts (into #shifts*#trials x #features)
            A = vertcat(trials{:});

            % remap the labels to -1 / +1
            labels(labels == model.classes(1)) = -1;
            labels(labels == model.classes(2)) = +1;

            % and generate shift-replicated target matrix b
            b = repmat(labels,1,length(offsets));
            b = b(:);

            C0 = -[b A];   % (bias is first elem, followed by rest...)
            ccp = C0';

            % weight vectors
            [m,n] = size(A);
            Wz = zeros(n+1,m);          % consensus variable
            Wx = zeros(n+1,m);          % data weights variable, FxN; n+1 accounts for the added bias (this is actually just the initial conditions vector...)
            Wu = zeros(n+1,m);          % scaled dual parameter u for Wx

            alpha = args.alpha;         % over-relaxation parameter (if 1, OR turned off)
            rho = args.rho;             % augmented Lagrangian parameter; probably determines sort of the tightness of the constraints...
            kappa = args.kappa;         % blend factor between trialwise sparsity and trialwise agreement
            K = args.patterns;          % the maximum number of SV's to recover

            rescale = 1;                  % scale of the left-hand side term (could be m/length(offsets)
            for mu = args.lambdas       % regularization parameter (should be decreasing)
                fprintf('Now using lambda = %.3f\n',mu);
                if args.verbose
                    fprintf('%3s\t%10s\t%10s\t%10s\t%10s\t%10s\n', 'iter', ...
                        'r norm', 'eps pri', 's norm', 'eps dual', 'objective');
                end

                for k = 1:args.admm_iters
                    Wzold = Wz;
                    fprintf('Outer step #%.0f\n',k);

                    %% x-update
                    Wx = ParadigmPAL.minimize_batch(Wx, @ParadigmPAL.x_objective_batch, args.cg_iters, Wz, Wu, ccp, rho);
                    Wx_hat = alpha*Wx + (1-alpha)*Wzold;  % (optional over-relaxation)

                    %% z-update
                    Wz = Wx_hat + Wu;
                    [Wz,fuuu,U,V] = ParadigmPAL.prox_sum(Wz,rescale,mu,kappa,rho,K,T,m,args);

                    %% dual parameter update
                    fprintf('  updating u''s\n');
                    Wu = Wu + (Wx_hat - Wz);

                    %% diagnostics, reporting, termination checks
                    fprintf('  diagnostics...\n');
                    history.objval(k)  = 0; %ParadigmPAL.full_objective(ccp, mu, Wx, Wz);
                    history.r_norm(k)  = norm(Wx(:) - Wz(:));
                    history.s_norm(k)  = norm(rho*(Wz(:) - Wzold(:)));
                    history.eps_pri(k) = args.abs_tol + args.rel_tol*max(norm(Wx(:)), norm(Wz(:)));
                    history.eps_dual(k)= args.abs_tol + args.rel_tol*norm(rho*Wu(:));

                    if args.verbose
                        fprintf('%3d\t%10.4f\t%10.4f\t%10.4f\t%10.4f\t%10.2f\n', k, ...
                            history.r_norm(k), history.eps_pri(k), ...
                            history.s_norm(k), history.eps_dual(k), history.objval(k));
                        if strcmp(args.verbose,'really') && nnz(fuuu)
                            f=figure; set(f,'Position',get(f,'Position')-[0 500 0 0]);
                            imagesc(Wz); ylabel('features'); xlabel('shifted trials'); title(sprintf('ADMM iter #%.0f, lambda=%.3f',k,mu));
                            figure;
                            % weight patterns
                            N = nnz(fuuu);
                            cols = ceil(sqrt(N));
                            rows = ceil(N/cols);
                            for p=1:N
                                subplot(cols,rows,p);
                                plot((0:round(args.duration * signal.srate)-1)/signal.srate,reshape(U(2:end,p),length(cutwindow),signal.nbchan)); title(sprintf('weight array #%.0f',p)); xlabel('time (s)'); ylabel('weight');
                            end
                            % topoplots...
                            f=figure; set(f,'Position',get(f,'Position')-[550 0 0 0]);
                            N = nnz(fuuu);
                            cols = ceil(sqrt(N));
                            rows = ceil(N/cols);
                            for p=1:N
                                subplot(cols,rows,p);
                                mat = reshape(U(2:end,p),length(cutwindow),signal.nbchan);
                                mat_inner = mat(round(end*1/4):round(end*3/4),:);
                                [dummy,peakidx] = max(mean(abs(mat_inner),2)); % #ok<NASGU>
                                topoplot(mat_inner(peakidx,:),signal.chanlocs);
                            end
                            drawnow;
                        end
                    end

                    if history.r_norm(k) < history.eps_pri(k) && history.s_norm(k) < history.eps_dual(k)
                        break; end

                    %% update of rho
                    if args.rho_update
                        if history.r_norm(k) > args.rho_cutoff * history.s_norm(k)
                            disp('  incrementing rho.');
                            rho = rho .* args.rho_incr;
                            Wu = Wu ./ args.rho_incr;
                        elseif history.s_norm(k) > args.rho_cutoff * history.r_norm(k)
                            disp('  decrementing rho.');
                            rho = rho ./ args.rho_incr;
                            Wu = Wu .* args.rho_incr;
                        end
                    end
                end
            end
            1
            %% plotting...

%             % full weight matrix
%             figure;imagesc(Wz); title('full weight matrix'); xlabel('features'); ylabel('shifted trials');
%
%             % shift matrix (?)
%             figure; imagesc(reshape(V(:,1),length(offsets),[])); title('shift matrix'); xlabel('trials'); ylabel('shifts');
        end

        function outputs = predict_simple(self,signal,model)
            % Override this function to implement your prediction code
            % Outputs = predict_simple(Sginal,Model)
            %
            % In:
            %   Signal : a signal pre-processed according to the model's filter graph
            %
            %   Model : a predictive model as created by your calibrate_simple() function
            %
            % Out:
            %   Outputs : a prediction/estimate for the most recent time point in the data (or one for
            %             every epoch if the signal is epoched); see ml_predict for the allowed formats

            error('This paradigm implements no predict() function.');
        end

        function visualize(self,model)
            % Optionally override this function to implement your visualization code
            % visualize(Model)
            %
            % In:
            %   Model: a model as created by your calibrate() function;
            %          a plot or GUI will be produced to inspect the model
            %

            error('This paradigm implements no visualize() function.');
        end

        % --- implementation ----

        function model = calibrate(self,varargin)
            % Implementation of the calibrate() function (see ParadigmBase)
            % Model = calibrate(Collection,GoalIdentifier,Arguments...)
            %
            % This is essentially a wrapper around calibrate_simple() which peels off any additional
            % data structure (for multiple collections and stream bundles) before passing it on.
            %
            % In:
            %   Collection : a collection (cell array) of stream bundles
            %
            %   GoalIdentifier: the goal-identifier
            %
            %   Arguments... : further optional user arguments
            %
            % Out:
            %   Model : The returned model struct

            % get the collection argument and extract the signal argument
            collection = arg_extract(varargin,{'collection','Collection'},[],{});
            if ~isempty(collection)
                if length(collection) > 1
                    error('This paradigm does not support collections of more than one recording.'); end
                if length(collection{1}.streams) > 1
                    error('This paradigm does not support more than one stream in parallel.'); end
                signal = collection{1}.streams{1};
            else
                signal = {};
            end

            % call calibrate_simple with the signal passed in as additional argument
            model = self.calibrate_simple('signal',signal,varargin{:});
        end

        function outputs = predict(self,bundle,model)
            % Override this function to implement your prediction code
            % Outputs = predict(Bundle,Model)
            %
            % This is a wrapper around the predict_simple() function which peels off the stream bundle
            % representation and just passes on the contained signal.
            %
            % In:
            %   Bundle : stream bundle preprocessed by the model's filter graph
            %
            %   Model : a predictive model
            %
            % Out:
            %   Outputs : the outputs of calibrate_simple()

            if length(bundle.streams) > 1
                error('This paradigm does not support more than one stream in parallel.'); end
            outputs = self.predict_simple(bundle.streams{1},model);
        end

    end

    methods(Static)
        function obj = full_objective(ccp, mu, Wx, Wz)
            % this is the overall objective function of alignment learning
            obj = sum(log(1 + exp(sum(ccp.*Wx)))) + mu*(norm_nuc(Wz) + sum(norms(Wz))) * size(ccp,2);
        end

        function [val,grad] = x_objective(x, C, z, u, rho)
            % this is the objective function for the decoupled x criterion (basically l2 logreg)
            ecx = exp(C*x);
            val = sum(log(1 + ecx)) + (rho/2)*norm(x - z + u).^2;
            grad = C'*(ecx./(1 + ecx)) + rho*(x - z + u);
        end

        function [vals,grads] = x_objective_batch(Wx,Wz,Wu,ccp,rho)
            % this is the batch version of the above (for all samples & variables in parallel)
            ecx = exp(sum(ccp.*Wx));
            d = Wx - Wz + Wu;
            vals = log(1 + ecx) + (rho/2)*sum(d.*d);
            grads = bsxfun(@times,ccp,(ecx./(1 + ecx))) + rho*d;
        end

        % l1/l2 shrinkage operator
        function Wz = shrinkage(Wa, kappa)
            Wz = pos(1 - kappa/norm(Wa))*Wa;
        end

        % l1 shrinkage operator
        function z = shrinkage_l1(a, kappa)
            z = max(0, a-kappa) - max(0, -a-kappa);
        end

        % trace norm proximal operator
        function [W,fuuu,U,V] = prox_tr(W,rescale,mu,kappa,rho,K,args)
            fprintf('  trace norming r\n');
            [U,S,V] = pca(W,K,args.lanczos_iters);
            fuuu = ParadigmPAL.shrinkage_l1(diag(S),(rescale*mu*(1-kappa))/rho); nn = length(fuuu); % we shrink the SV spectrum using prox operator
            Ws = sparse(1:nn,1:nn,fuuu,size(S,1),size(S,2));
            % clamp minority of V's members at 0
            smv = sign(median(V));
            V(:,smv==-1) = min(0,V(:,smv==-1));
            V(:,smv==1) = max(0,V(:,smv==1));
            W = U*(Ws*V');

            if strcmp(args.verbose,'ultra')
                f=figure; set(f,'Position',get(f,'Position')-[0 500 0 0]);
                imagesc(W); ylabel('features'); xlabel('shifted trials'); title(sprintf('lambda=%.3f',mu));
                figure;
                % weight patterns
                N = nnz(spec);
                cols = ceil(sqrt(N));
                rows = ceil(N/cols);
                signal=args.signal;
                for p=1:N
                    subplot(cols,rows,p);
                    plot((0:(numel(U(2:end,p))/signal.nbchan)-1)/signal.srate,reshape(U(2:end,p),[],signal.nbchan)); title(sprintf('weight array #%.0f',p)); xlabel('time (s)'); ylabel('weight');
                end
                drawnow;
            end

        end

        % group-sparse proximal operator
        function W = prox_gl(W,rescale,mu,kappa,rho,T,m)
            fprintf('  sparsifying z\n');
            for i = 1:m
                W(2:end,i) = ParadigmPAL.shrinkage(W(2:end,i), (rescale*mu*kappa/10)/rho); end
        end

        % summed proximal operator (using some splitting scheme, see Bauscke, 2008)
        function [x,fuuu,U,V] = prox_sum(r,rescale,mu,kappa,rho,K,T,m,args)
            y = r;
            p = zeros(size(y));
            q = zeros(size(y));
            for n=1:args.inner_iters
                [x,fuuu,U,V] = ParadigmPAL.prox_tr(y+q,rescale,mu,kappa,rho,K,args);
                q = y+q-x;
                y = ParadigmPAL.prox_gl(x+p,rescale,mu,kappa,rho,T,m);
                p = x+p-y;
                if norm(x(:) - y(:)) < args.rel_tol*max(norm(x(:)),norm(y(:)))
                    break; end
            end
        end

        function X = minimize_batch(X, f, len, Wz,Wu,ccp,rho)
            % Copyright (C) 2001 - 2010 by Carl Edward Rasmussen, 2010-01-03

            INT = 0.1;    % don't reevaluate within 0.1 of the limit of the current bracket         % const
            EXT = 3.0;                  % extrapolate maximum 3 times the current step-size         % const
            MAX = 20;                         % max 20 function evaluations per line search         % const
            RATIO = 10;                                       % maximum allowed slope ratio         % const
            SIG = 0.1; RHO = SIG/2; % SIG and RHO are the constants controlling the Wolfe-          % const
            if max(size(len)) == 2, red=len(2); len=len(1); else red=1; end                         % scalar

            [f0 df0] = feval(f,X,Wz,Wu,ccp,rho);          % get function value and gradient
            N = length(f0);

            [A,B,d1,d2,d3,d4,f1,f2,f3,f4,x1,x2,x4] = deal(zeros(1,N)); % initialize a few variables...
            df3 = zeros(size(X));
            i = 0;                                            % zero the run length counter         % scalar
            ls_failed = false(1,N);                             % no previous line search has failed% boolean vector
            i = i + (len<0);                                            % count epochs?!            % scalar
            s = -df0;                                 % initial search direction (steepest)         % column MATRIX
            d0 = -sum(s.*s);                                                    % and slope         % row vector
            x3 = red./(1-d0);                                  % initial step is red/(|s|+1)        % row vector

            keep_going = true(1,N);                                                                 % mask for the outermost loop updates...
            while any(keep_going) && i < abs(len)                                      % while not finished            % scalar
                i = i + (len>0);                                      % count iterations?!          % scalar
                fprintf('  Cg step #%.0f\n',i);

                X0 = X;                                                                             % column MATRIX
                F0 = f0;                                                                            % row vector
                dF0 = df0;                   % make a copy of current values                        % column MATRIX
                if len>0, M = MAX; else M = min(MAX, -len-i); end                                   % scalar

                om = keep_going;                                                                    % mask for outer-loop updates
                while 1                             % keep extrapolating as long as necessary
                    x2(om) = 0;                                                                     % row vector
                    f2(om) = f0(om);                                                                % row vector
                    d2(om) = d0(om);                                                                % row vector
                    f3(om) = f0(om);                                                                % row vector
                    df3(:,om) = df0(:,om);                                                          % column matrix
                    ok(om) = false;                                                                 % boolean vector
                    while ~all(ok(om)) && M > 0                                                     % mask check
                        M = M - 1; i = i + (len<0);                         % count epochs?!        % scalar
                        nokom = ~ok&om;
                        [f3(nokom) df3(:,nokom)] = feval(f, X(:,nokom) + bsxfun(@times,x3(nokom),s(:,nokom)), Wz(:,nokom), Wu(:,nokom), ccp(:,nokom), rho); % masked array op
                        ok(om) = ok(om) | ~(isnan(f3(om)) | isinf(f3(om)) | any(isnan(df3(om))+isinf(df3(om))));            % mask update
                        x3(~ok&om) = (x2(~ok&om)+x3(~ok&om))/2;         % bisect and try again      % maked update
                    end
                    mask = f3<F0;                                       % keep best values          % mask generation
                    X0(:,mask&om) = X(:,mask&om) + bsxfun(@times,x3(mask&om),s(:,mask&om));         % masked updates
                    F0(mask&om) = f3(mask&om);
                    dF0(:,mask&om) = df3(:,mask&om);

                    d3(om) = sum(df3(:,om).*s(:,om));                                    % new slope% row vector

                    om(d3 > SIG*d0 | f3 > f0 + x3.*d0*RHO) = false;             % update outer mask
                    if ~any(om) || M == 0  % are we done extrapolating?
                        break; end

                    x1(om) = x2(om); f1(om) = f2(om); d1(om) = d2(om); % move point 2 to point 1    % row vectors
                    x2(om) = x3(om); f2(om) = f3(om); d2(om) = d3(om); % move point 3 to point 2    % row vectors
                    A(om) = 6*(f1(om)-f2(om))+3*(d2(om)+d1(om)).*(x2(om)-x1(om));% make cubic extrap.% row vectors
                    B(om) = 3*(f2(om)-f1(om))-(2*d1(om)+d2(om)).*(x2(om)-x1(om));
                    x3(om) = x1(om)-d1(om).*(x2(om)-x1(om)).^2./(B(om)+sqrt(B(om).*B(om)-A(om).*d1(om).*(x2(om)-x1(om)))); % num. error possible, ok!

                    % implement fixups: num prob | wrong sign | new point beyond extrapolation limit?
                    issue_mask = ~isreal(x3) | isnan(x3) | isinf(x3) | x3 < 0 | x3 > x2*EXT;
                    x3(om & issue_mask) = x2(om & issue_mask)*EXT;
                    % new point too close to previous point?
                    issue_mask = ~issue_mask & (x3 < x2+INT*(x2-x1));
                    x3(om & issue_mask) = x2(om & issue_mask) + INT*(x2(om & issue_mask) - x1(om & issue_mask));
                end                                                       % end extrapolation

                while 1
                    om = keep_going & (abs(d3) > -SIG*d0 | f3 > f0+x3.*d0*RHO);    % update outer mask
                    if ~(any(om) && M > 0)
                        break; end
                    tog = d3 > 0 | f3 > f0+x3.*d0*RHO;           % subinterval toggle
                    x4(om&tog) = x3(om&tog);
                    f4(om&tog) = f3(om&tog);
                    d4(om&tog) = d3(om&tog);
                    x2(om&~tog) = x3(om&~tog);
                    f2(om&~tog) = f3(om&~tog);
                    d2(om&~tog) = d3(om&~tog);
                    tog=f4>f0;                                   % interpolation toggle
                    x3(om&tog) = x2(om&tog) - (0.5*d2(om&tog).*(x4(om&tog)-x2(om&tog)).^2) ./ (f4(om&tog)-f2(om&tog)-d2(om&tog).*(x4(om&tog)-x2(om&tog)));  % quadratic interpolation
                    A(om&~tog) = 6*(f2(om&~tog)-f4(om&~tog))./(x4(om&~tog)-x2(om&~tog)) + 3*(d4(om&~tog)+d2(om&~tog));                    % cubic interpolation
                    B(om&~tog) = 3*(f4(om&~tog)-f2(om&~tog)) - (2*d2(om&~tog)+d4(om&~tog)).*(x4(om&~tog)-x2(om&~tog));
                    x3(om&~tog) = x2(om&~tog) + (sqrt(B(om&~tog).*B(om&~tog) - A(om&~tog).*d2(om&~tog).*(x4(om&~tog)-x2(om&~tog)).^2)-B(om&~tog))./A(om&~tog);  % num. error possible, ok!
                    numprob = isnan(x3) | isinf(x3);             % if we had a numerical problem then bisect
                    x3(om&numprob) = (x2(om&numprob)+x4(om&numprob))/2;

                    x3(om) = max(min(x3(om), x4(om)-INT*(x4(om)-x2(om))),x2(om)+INT*(x4(om)-x2(om)));  % don't accept too close
                    [f3(om) df3(:,om)] = feval(f, X(:,om) + bsxfun(@times,x3(om),s(:,om)), Wz(:,om), Wu(:,om), ccp(:,om), rho); % masked array op
                    mask = om & f3<F0;                                      % keep best values
                    X0(:,mask) = X(:,mask) + bsxfun(@times,x3(mask),s(:,mask));
                    F0(mask) = f3(mask);
                    dF0(:,mask) = df3(:,mask);
                    M = M - 1; i = i + (len<0);                             % count epochs?!
                    d3(om) = sum(df3(:,om).*s(:,om));                                    % new slope% row vector
                end

                % --- we are here ---

                okay = abs(d3) < -SIG*d0 & f3 < f0+x3.*d0*RHO;

                upd = okay & keep_going;
                X(:,upd) = X(:,upd) + bsxfun(@times,x3(upd),s(:,upd));
                f0(upd) = f3(upd);
                s(:,upd) = bsxfun(@times,sum(df3(:,upd).*df3(:,upd)) - sum(df0(:,upd).*df3(:,upd)) ./ sum(df0(:,upd).*df0(:,upd)),s(:,upd)) - df3(:,upd);
                df0(:,upd) = df3(:,upd);                                % swap derivatives
                d3(upd) = d0(upd);
                d0(upd) = sum(df0(:,upd).*s(:,upd));
                mask = upd & d0 > 0;                         % otherwise use steepest direction
                s(:,mask) = -df0(:,mask);
                d0(mask) = -sum(s(:,mask).*s(:,mask));
                x3(upd) = x3(upd) .* min(RATIO, d3(upd)./(d0(upd)-realmin));  % slope ratio but max RATIO
                ls_failed(upd) = false;                              % this line search did not fail

                upd = ~okay & keep_going;
                X(:,upd) = X0(:,upd);
                f0(upd) = F0(upd);
                df0(:,upd) = dF0(:,upd);

                keep_going(upd) = keep_going(upd) & ~ls_failed(upd); % line search failed twice in a row
                                                                     % so we give up on these guys
                upd = upd & keep_going;
                if i <= abs(len)
                    s(:,upd) = -df0(:,upd);
                    d0(:,upd) = -sum(s(:,upd).*s(:,upd));                             % try steepest
                    x3(upd) = 1./(1-d0(upd));
                    ls_failed(upd) = true;
                end
            end
        end

        function layout = dialog_layout_defaults(self)
            layout = {'SignalProcessing.Resampling.SamplingRate','SignalProcessing.SpectralSelection.FrequencySpecification', ...
                'SignalProcessing.EpochExtraction','','Lambdas','WaveDuration','Stepping','ADMMIterations','CgIterations',...
                'LanczosIterations','NumPatterns','SparsityTraceTradeoff'};
        end

    end
end

% (turn off a few editor warnings because some actual implementations are missing in this file)
%#ok<*INUSD,*STOUT,*MANU>


function MDP = DEM_demo_MDP_rule
% Demo of active inference for visual salience
%__________________________________________________________________________
%
% This routine simulates a crude form of consciousness using active
% inference and structure learning to vitiate ignorance or nescience. This
% entails learning the hyper parameters of causal structure generating
% outcomes and then using Bayesian model reduction (during sleep) to
% minimise complexity.
%
% We first set up an abstract problem in which an agent has to respond
% according to rules (identify the correct colour depending upon one of
% three rules that are specified by the colour of a cue in the centre of
% vision). If the rule is centre, the colour is always green; however, if
% the colour of the Centre cue is red, the correct colour is on the left
% (and on the right if the queue is blue). Simulations are provided when
% the agent knows the rules. This is then repeated in the absence
% (nescience) of any knowledge about the rules to see if the agent can
% learn causal structure through Bayesian belief updating of the likelihood
% array (A).
%
% We then consider the improvement in performance (in terms of variational
% free energy, its constituent parts and performance) following Bayesian
% model reduction of the likelihood model (heuristically, like slow wave
% sleep), followed by a restitution of posterior beliefs during fictive
% active inference (as in REM sleep). Finally, we address the communication
% of the implicit  structure learning to a conspecific or child to
% demonstrate the improvement under instruction.
%
% see also: DEM_demo_MDP_habits.m and spm_MPD_VB_X.m
%__________________________________________________________________________
% Copyright (C) 2005 Wellcome Trust Centre for Neuroimaging

% Karl Friston
% $Id: DEM_demo_MDP_rule.m 7679 2019-10-24 15:54:07Z spm $

% set up and preliminaries
%==========================================================================

% specify the (rule-based) generative process (environment)
%==========================================================================
rng('default')

% prior beliefs about initial states (in terms of counts_: D
%--------------------------------------------------------------------------
D{1} = [1 1 1]';        % rule:   {'left','centre','right'}
D{2} = [1 1 1]';        % what:   {'red','green','blue'}
D{3} = [0 0 0 1]';      % where:  {'left','centre','right','null'}
D{4} = [0 0 0 1]';      % report: {'red','green','blue','undecided'}

% probabilistic mapping from hidden states to outcomes: A
%--------------------------------------------------------------------------
Nf    = numel(D);
Ns    = zeros(1,Nf);
for f = 1:Nf
    Ns(f) = numel(D{f});
end
for f1 = 1:Ns(1)                     % rule
    for f2 = 1:Ns(2)                 % correct colour
        for f3 = 1:Ns(3)             % location of fixation
            for f4 = 1:Ns(3)         % decision

                % A{1} what: {'red','green','blue','null'}
                %==========================================================
                if f3 == 4, A{1}(4,f1,f2,f3,f4)  = 1;  end
                if f3 == 2, A{1}(f1,f1,f2,f3,f4) = 1;  end
                if f3 == 1
                    if f1 == 1
                        A{1}(f2,f1,f2,f3,f4) = 1;
                    else
                        A{1}(1:3,f1,f2,f3,f4) = 1/3;
                    end
                end
                if f3 == 3
                    if f1 == 3
                        A{1}(f2,f1,f2,f3,f4) = 1;
                    else
                        A{1}(1:3,f1,f2,f3,f4) = 1/3;
                    end
                end

                % A{2} where: {'left','centre','right','null'}
                %----------------------------------------------------------
                A{2}(f3,f1,f2,f3,f4) = 1;

                % A{3} feedback: {'null','right','wrong'}
                %----------------------------------------------------------
                if f4 == 4,
                    A{3}(1,f1,f2,f3,f4) = 1;                             % undecided
                else
                    if f1 == 2 && f4 == 2,  A{3}(2,f1,f2,f3,f4) = 1; end % right
                    if f1 == 2 && f4 ~= 2,  A{3}(3,f1,f2,f3,f4) = 1; end % wrong
                    if f1 == 1 && f4 == f2, A{3}(2,f1,f2,f3,f4) = 1; end % right
                    if f1 == 1 && f4 ~= f2, A{3}(3,f1,f2,f3,f4) = 1; end % wrong
                    if f1 == 3 && f4 == f2, A{3}(2,f1,f2,f3,f4) = 1; end % right
                    if f1 == 3 && f4 ~= f2, A{3}(3,f1,f2,f3,f4) = 1; end % wrong
                end
            end
        end
    end
end
Ng    = numel(A);
for g = 1:Ng
    No(g) = size(A{g},1);
    A{g}  = double(A{g});
end

% controlled transitions: B{f} for each factor
%--------------------------------------------------------------------------
for f = 1:Nf
    B{f} = eye(Ns(f));
end

% control states B(3): where {'left','centre','right','null'}
%--------------------------------------------------------------------------
for k = 1:Ns(3)
    B{3}(:,:,k) = 0;
    B{3}(k,:,k) = 1;
end

% control states B(4): report {'red','green','blue','undecided'}
%--------------------------------------------------------------------------
for k = 1:Ns(4)
    B{4}(:,:,k) = 0;
    B{4}(k,:,k) = 1;
end

% allowable policies (specified as the next action) U
%--------------------------------------------------------------------------
U(1,1,:)  = [1 1 1 4]';         % sample left
U(1,2,:)  = [1 1 2 4]';         % sample centre
U(1,3,:)  = [1 1 3 4]';         % sample right
U(1,4,:)  = [1 1 4 1]';         % return to centre and report red
U(1,5,:)  = [1 1 4 2]';         % return to centre and report green
U(1,6,:)  = [1 1 4 3]';         % return to centre and report blue

% priors: (utility) C; the agent expects to avoid mistakes
%--------------------------------------------------------------------------
for g = 1:Ng
    C{g}  = zeros(No(g),1);
end
% and expects itself to make a decision after the fifth observation
%--------------------------------------------------------------------------
C{3} = [ 0  0  0  0 -8 -8;
         0  0  0  0  0  0;
        -4 -4 -4 -4 -4 -4];

% MDP Structure
%--------------------------------------------------------------------------
mdp.T = size(C{3},2);           % number of moves
mdp.U = U;                      % allowable policies
mdp.A = A;                      % observation model
mdp.B = B;                      % transition probabilities
mdp.C = C;                      % preferred outcomes
mdp.D = D;                      % prior over initial states

mdp.Aname = {'what','where','feedback'};
mdp.Bname = {'rule','colour','where','decision'};

% illustrate a single trial
%==========================================================================
mdp  = spm_MDP_check(mdp);
MDP  = spm_MDP_VB_X(mdp);

% show belief updates (and behaviour)
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 1'); clf
spm_MDP_VB_trial(MDP);

% illustrate phase-precession and responses
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 2'); clf
spm_MDP_VB_LFP(MDP,[],1);

% illustrate behaviour
%--------------------------------------------------------------------------
subplot(3,2,3)
spm_MDP_rule_plot(MDP)

% illustrate a sequence of trials
%==========================================================================
clear MDP

% create structure array
%--------------------------------------------------------------------------
N     = 8;
for i = 1:N
    MDP(i) = mdp;
end

% Solve an example sequence under different initial states
%==========================================================================
MDP   = spm_MDP_VB_X(MDP);

% illustrate behavioural responses and neuronal correlates
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 3'); clf
spm_MDP_VB_game(MDP);

[F,Fu] = spm_MDP_F(MDP);

subplot(6,1,4), plot(1:N,F),  xlabel('trial'), spm_axis tight, title('Free energy','Fontsize',16)
subplot(6,1,5), plot(1:N,Fu), xlabel('trial'), spm_axis tight, title('Confidence','Fontsize',16)
subplot(6,1,6), spm_MDP_plot_moves(MDP),       spm_axis tight
title('Action ','Fontsize',16), legend({'saccades','hits'})

% illustrate phase-amplitude (theta-gamma) coupling
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 4'); clf
spm_MDP_VB_LFP(MDP);

% illustrate behaviour in more detail
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 5'); clf;
for i = 1:min(N,15)
    subplot(5,3,i), spm_MDP_rule_plot(MDP(i));
    axis square
end

% repeat with rule learning
%==========================================================================

% retain knowledge about the cue but remove knowledge about the rules
%--------------------------------------------------------------------------
for f1 = 1:Ns(1)                     % rule
    for f2 = 1:Ns(2)                 % correct colour
        for f3 = 1:Ns(3)             % location of fixation
            for f4 = 1:Ns(3)         % decision

                % A{1} what: {'red','green','blue','null'}
                %==========================================================
                if f3 == 4, a{1}(4,f1,f2,f3,f4)  = 128;  end
                if f3 == 2, a{1}(f1,f1,f2,f3,f4) = 128;  end
                if f3 == 1
                    a{1}(1:3,f1,f2,f3,f4) = 1;
                end
                if f3 == 3
                    a{1}(1:3,f1,f2,f3,f4) = 1;
                end
            end
        end
    end
end

a{2}   = A{2}*128;
a{3}   = A{3}*128;
mda    = mdp;
mda.a  = a;
mda.a0 = a;

% create structure array
%--------------------------------------------------------------------------
clear MDP
N     = 32;
for i = 1:N
    MDP(i) = mda;
end

% Solve - an example sequence
%==========================================================================
rng('default')
MDP  = spm_MDP_VB_X(MDP);

% show responses to a difficult (right) trial before and after learning
%--------------------------------------------------------------------------
for i = 1:N
    s(:,i) = MDP(i).s(:,1);
end
i = find(ismember(s(1:2,:)',s(1:2,1)','rows'));

spm_figure('GetWin','Figure 6 - before');
spm_MDP_VB_LFP(MDP(i(1)),[],2);
subplot(3,2,2), set(gca,'YLim',[-.1 1])
subplot(3,2,4), set(gca,'YLim',[-.1 .2])
subplot(3,2,3), spm_MDP_rule_plot(MDP(i(1)))
title(['trial ' num2str(i(1))],'FontSize',16)

spm_figure('GetWin','Figure 6 - after' );
spm_MDP_VB_LFP(MDP(i(end)),[],2);
subplot(3,2,2), set(gca,'YLim',[-.1 1]);
subplot(3,2,4), set(gca,'YLim',[-.1 .2])
subplot(3,2,3), spm_MDP_rule_plot(MDP(i(end)))
title(['trial ' num2str(i(end))],'FontSize',16)

% show learning effects in terms of underlying uncertainty
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 7 - before'); clf; spm_MDP_VB_trial(MDP(i(1)));
spm_figure('GetWin','Figure 7 - after' ); clf; spm_MDP_VB_trial(MDP(i(end)));

% illustrate phase-amplitude (theta-gamma) coupling
%--------------------------------------------------------------------------
n  = min(N,4);
spm_figure('GetWin','Figure 4a'); clf
spm_MDP_VB_LFP(MDP(1:n),[],2);
for i = 1:n
    subplot(4,n,i + 3*n), spm_MDP_rule_plot(MDP(i));
    axis square
end

% show trial-by-trial measures in terms of free energy terms
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 8'); clf;
spm_MDP_VB_game(MDP);

[F,Fu] = spm_MDP_F(MDP);

subplot(6,1,4), plot(1:N,F),  xlabel('trial'), spm_axis tight, title('Free energy','Fontsize',16)
subplot(6,1,5), plot(1:N,Fu), xlabel('trial'), spm_axis tight, title('Confidence','Fontsize',16)
subplot(6,1,6), spm_MDP_plot_moves(MDP),       spm_axis tight
title('Action ','Fontsize',16), legend({'saccades','hits'})

% illustrate Bayesian model reduction
%--------------------------------------------------------------------------
OPTIONS.g = 1;
OPTIONS.f = 2;
OPTIONS.T = 3;
OPTIONS.m = @(i,i1,i2,i3,i4) i == i2;

for n = 1:N
    sdp{n} = spm_MDP_VB_sleep(MDP(n),OPTIONS);
    cor(n) = corr(spm_vec(MDP(n).A{1} > 0),spm_vec(sdp{n}.a{1}) > 0);
end
[m,n] = max(cor);

spm_figure('GetWin','Figure 9'); clf; str = sprintf('Sleep (trial %d)',n);
for n = n
    subplot(3,2,1), spm_MDP_A_plot(MDP(n).A);
    subplot(3,2,2), spm_MDP_A_plot(MDP(n).a0), title('Before','Fontsize',16)
    subplot(3,2,3), spm_MDP_A_plot(MDP(n).a),  title('After', 'Fontsize',16)
    subplot(3,2,4), spm_MDP_A_plot(sdp{n}.a),  title(str,     'Fontsize',16)
end

% return % here for short demo

% Bayesian model reduction with dreaming
%--------------------------------------------------------------------------
OPTIONS.o = {MDP.o};
rdp       = spm_MDP_VB_sleep(MDP(n),OPTIONS);
sdp       = sdp{n};

% illustrate benefits of sleep (with and without dreaming)
%==========================================================================
RDP   = MDP(n:end);
RDP   = rmfield(RDP,{'o','u'});
for i = 1:length(RDP)
    RDP(i).s = RDP(i).s(:,1);
end
SDP    = RDP;
SDP(1).a = sdp.a;
RDP(1).a = rdp.a;
SDP    = spm_MDP_VB_X(SDP);
RDP    = spm_MDP_VB_X(RDP);

% plot performance and without sleep
%--------------------------------------------------------------------------
spm_figure('GetWin','Figure 10'); clf;

[F,Fu] = spm_MDP_F(MDP);

subplot(4,1,1),plot(1:N,F, ':'),  xlabel('trial'), spm_axis tight, hold on
subplot(4,1,2),plot(1:N,Fu,':'),  xlabel('trial'), spm_axis tight, hold on

% and after SWS sleep
%--------------------------------------------------------------------------
[F,Fu] = spm_MDP_F(SDP);

subplot(4,1,1),plot(n:N,F ,'-.'), xlabel('trial'), spm_axis tight, hold on
subplot(4,1,2),plot(n:N,Fu,'-.'), xlabel('trial'), spm_axis tight, hold on

% and after REM sleep
%--------------------------------------------------------------------------
[F,Fu] = spm_MDP_F(RDP);

subplot(4,1,1),plot(n:N,F ,'-'),  xlabel('trial'), spm_axis tight, hold on
title('Free energy (with sleep)','Fontsize',16)
subplot(4,1,2),plot(n:N,Fu,'-'),  xlabel('trial'), spm_axis tight, hold on
title('Confidence (with sleep)','Fontsize',16)

% and correct responses
%--------------------------------------------------------------------------
hit = @(MDP) any(MDP.o(3,:) == 2) & ~any(MDP.o(3,:) == 3);
k   = min(Fu);
hold on
for i = 1:numel(RDP)
    j = i + n - 1;
    if hit(SDP(i)), plot(j,k,'hr','MarkerSize',8); end
end
for i = 1:(n - 1)
    if hit(MDP(i)), plot(i,k,'hb','MarkerSize',8); end
end
hold off

% illustrate the benefits of instruction (communication of model)
%==========================================================================

% create instructed agent
%--------------------------------------------------------------------------
n     = 1:16;
NDP   = MDP(n);
NDP   = rmfield(NDP,{'o','u'});
for i = 1:length(NDP)
    NDP(i).s = NDP(i).s(:,1);
    NDP(i).a = sdp.a0;
end
NDP   = spm_MDP_VB_X(NDP);

% plot performance and without received wisdom
%--------------------------------------------------------------------------
[F,Fu] = spm_MDP_F(MDP(n));

subplot(4,2,5),plot(n,F ,':'), xlabel('trial'), spm_axis tight, hold on
subplot(4,2,6),plot(n,Fu,':'), xlabel('trial'), spm_axis tight, hold on

% and after instruction
%--------------------------------------------------------------------------
[F,Fu] = spm_MDP_F(NDP(n));
k      = min(Fu);

subplot(4,2,5),plot(n,F),       xlabel('trial'), spm_axis tight
title('Free energy (instruction)','Fontsize',16)
subplot(4,2,6),plot(n,Fu),      xlabel('trial'), spm_axis tight
title('Confidence (instruction)','Fontsize',16)

% and correct responses
%--------------------------------------------------------------------------
for i = n
    hold on
    if hit(NDP(i)); plot(i,k,'hr','MarkerSize',8), end
    if hit(MDP(i)); plot(i,k,'hb','MarkerSize',8), end
    hold off
end

% run multiple subjects
%==========================================================================
BMR.g   = 1;
BMR.f   = 2;
BMR.T   = 3;
BMR.m   = @(i,i1,i2,i3,i4) i == i2;
OPT.BMR = BMR;

N     = 32;
Ns    = 64;
for m = 1:Ns

    % create structure array and solve
    %----------------------------------------------------------------------
    clear MDP OPTIONS
    for i = 1:N
        MDP(i) = mda;
    end
    rng(m)
    MDP  = spm_MDP_VB_X(MDP);

    % free energy and confidence
    %----------------------------------------------------------------------
    [F,Fu]   = spm_MDP_F(MDP);
    Fm(:,m)  = F(:);
    Fum(:,m) = Fu(:);

    % find run of correct responses
    %----------------------------------------------------------------------
    for i = 1:N
        if hit(MDP(i)); h(i,m) = 1; else, h(i,m) = 0;  end
    end

    % repeat with BMR
    %----------------------------------------------------------------------
    RDP = MDP;
    RDP = rmfield(RDP,{'u','o'});
    for i = 1:N
        RDP(i).a  = mda.a;
        RDP(i).a0 = mda.a0;
        RDP(i).s  = RDP(i).s(:,1);
    end
    rng(m)
    RDP      = spm_MDP_VB_X(RDP,OPT);
    [F,Fu]   = spm_MDP_F(RDP);
    Rm(:,m)  = F(:);
    Rum(:,m) = Fu(:);

    % look for instances of BMR
    %----------------------------------------------------------------------
    vA    = spm_vec(A{1});
    for i = 1:N
       c(i,1) = corr(vA,(spm_vec(RDP(i).a{1}) > 0));
    end
    bmr(:,m) = diff(c);

    % preferred locations
    %----------------------------------------------------------------------
    for i = 1:N
        if hit(RDP(i)); r(i,m) = 1; else, r(i,m) = 0;  end
    end

    % find run of correct responses
    %----------------------------------------------------------------------
    for i = 1:N
        if hit(RDP(i)); r(i,m) = 1; else, r(i,m) = 0;  end
    end

    % find non-learners
    %----------------------------------------------------------------------
    R = spm_conv(r,2,0);
    H = spm_conv(h,2,0);

    % show results
    %----------------------------------------------------------------------
    spm_figure('GetWin','Figure 11');clf
    subplot(4,1,1)
    b = bar(mean(R(:,:),2)); set(b,'EdgeColor','w','FaceColor',[1 1 1]*.0),hold on
    b = bar(mean(H(:,:),2)); set(b,'EdgeColor','w','FaceColor',[1 1 1]*.8),hold off
    xlabel('trial'), ylabel('probability of correct'), axis([1/2 (N + 1/2) 1/3 1]);
    title('Average performance','Fontsize',16)

    subplot(4,1,2)
    spm_plot_ci(mean(Rm'),var(Rm')), hold on
    plot(mean(Fm'),'r'),              hold off
    xlabel('trial'), ylabel('free energy'); set(gca,'XLim',[1 N])
    title('Average free energy','Fontsize',16)

    subplot(4,1,3)
    spm_plot_ci(mean(Rum'),var(Rum')), hold on
    plot(mean(Fum'),'r'),              hold off
    xlabel('trial'), ylabel('confidence');  set(gca,'XLim',[1 N])
    title('Average confidence','Fontsize',16)

    % show individual performance
    %----------------------------------------------------------------------
    subplot(4,1,4), image(32*(r'))
    xlabel('trial'), ylabel('subject')
    title('Aha moments','Fontsize',16)

    % plot model updates
    %----------------------------------------------------------------------
    hold on
    for i = 1:m
        j = find(bmr(:,i) > 0) + 1;
        try, plot(j(1),i,'.m','MarkerSize',32), end
        try, plot(j(2),i,'.r','MarkerSize',32), end
        j = find(bmr(:,i) < 0);
        plot(j, (j - j + i),'.b','MarkerSize',32)
    end
    hold off, drawnow

    save paper

end

return

% confidence - negatively over policies
%--------------------------------------------------------------------------
for i = 1:numel(MDP)
    p     = MDP(i).R;
    Fu(i) = sum(sum(p.*log(p)));
end

return

function spm_MDP_plot_moves(MDP)
for i = 1:numel(MDP)
    m(i) = sum(~~diff(MDP(i).u(3,:)));
end
h = bar(m); set(h,'EdgeColor','w','FaceColor',[1 1 1]*.9), hold on

% and correct responses
%--------------------------------------------------------------------------
hit   = @(MDP) any(MDP.o(3,:) == 2) & ~any(MDP.o(3,:) == 3);
for i = 1:numel(MDP)
    if hit(MDP(i))
        h = plot(i,0,'hr','MarkerSize',8);
        set(h,'MarkerFaceColor',[1 1/2 0])
    end
end
return

function spm_MDP_A_plot(A)
% assemble key parts of the likelihood array
%==========================================================================
for i = 1:3
    for j = 1:3;
        a{i,j} = squeeze(A{1}(:,i,:,j,4));
        a{i,j} = a{i,j}*diag(1./sum(a{i,j}));
    end
end
a = spm_cat(a);
imagesc(a);
title( 'Sample: left - center - right', 'FontSize',16)
ylabel('Rule: left - center - right','FontSize',14)
xlabel('Correct color', 'FontSize',14)
set(gca,'XTick',1:9)
set(gca,'YTick',1:12)
set(gca,'XTicklabel',repmat(['r','g','b'],[1 3])')
set(gca,'YTicklabel',repmat(['r','g','b',' '],[1 3])')
axis image

return

function spm_MDP_rule_plot(MDP)
% illustrates visual search graphically
%==========================================================================

% locations
%--------------------------------------------------------------------------
x{1} = [-1  0; 0  1; 1  0; 0  0];
x{2} = [-1 -1; 0 -1; 1 -1; 0 -2]/2;
col  = {'r','g','b','c'};

% plot cues
%--------------------------------------------------------------------------
if strcmp('replace',get(gca,'Nextplot'))

    % plot cues
    %----------------------------------------------------------------------
    s     = MDP.s;hold off
    for i = 1:length(MDP.D{3})
        a = MDP.A{1}(:,s(1),s(2),i,1);
        j = find(rand < cumsum(a),1);
        plot(x{1}(i,1),x{1}(i,2),['.',col{j}],'MarkerSize',32), hold on
    end

    % plot choices
    %----------------------------------------------------------------------
    for i = 1:length(MDP.D{4})
        a = find(MDP.A{3}(:,s(1),s(2),4,i));
        if a == 2
            plot(x{2}(i,1),x{2}(i,2),['.m'],'MarkerSize',32), hold on
        end
        plot(x{2}(i,1),x{2}(i,2),['.',col{i}],'MarkerSize',16), hold on

    end
    axis([-2 2 -2 2]);

end

% Extract and plot eye movements and choice
%--------------------------------------------------------------------------
for i = 1:numel(MDP.o(2,:))
    X(i,:) = x{1}(MDP.o(2,i),:);
end
plot(X(:,1),X(:,2),'k')
for i = 1:numel(MDP.s(4,:))
    X(i,:) = x{2}(MDP.s(4,i),:);
end
plot(X(:,1),X(:,2),'k')
axis([-2 2 -2 2]);

return


function [At,b,c,K,prep,origcoeff] = pretransfo(At,b,c,K,pars)

% [At,b,c,K,prep] = pretransfo(At,b,c,K)
%
% PRETRANSFO  Checks data and then transforms into internal SeDuMi format.
%
% **********  INTERNAL FUNCTION OF SEDUMI **********
%
% See also sedumi

% Nearly complete rewrite for SeDumI 1.4
% Copyright (c) 2013 Michael C. Grant
%
% This file is part of SeDuMi 1.1 by Imre Polik and Oleksandr Romanko
% Copyright (C) 2005 McMaster University, Hamilton, CANADA  (since 1.1)
%
% Copyright (C) 2001 Jos F. Sturm (up to 1.05R5)
%   Dept. Econometrics & O.R., Tilburg University, the Netherlands.
%   Supported by the Netherlands Organization for Scientific Research (NWO).
%
% Affiliation SeDuMi 1.03 and 1.04Beta (2000):
%   Dept. Quantitative Economics, Maastricht University, the Netherlands.
%
% Affiliations up to SeDuMi 1.02 (AUG1998):
%   CRL, McMaster University, Canada.
%   Supported by the Netherlands Organization for Scientific Research (NWO).
%
% This program is free software; you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation; either version 2 of the License, or
% (at your option) any later version.
%
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with this program; if not, write to the Free Software
% Foundation, Inc.,  51 Franklin Street, Fifth Floor, Boston, MA
% 02110-1301, USA

% -------------------------------------------------------------------------
% Make sure that all fields exist in K, and verify that they are valid
% -------------------------------------------------------------------------

if ~isfield(K,'f') || isempty(K.f)
    K.f = 0;
elseif numel(K.f) ~= 1 || K.f ~= floor(K.f) || K.f < 0 || ~isreal(K.f)
    error('K.f should be nonnegative integer')
end
if ~isfield(K,'l') || isempty(K.l)
    K.l = 0;
elseif K.l ~= floor(K.l) || K.l < 0 || ~isreal(K.l)
    error('K.l should be nonnegative integer')
end
if ~isfield(K,'q') || ~nnz(K.q)
    K.q = zeros(1,0);
else
    K.q = K.q(:)';
    if any(K.q ~= floor(K.q)) || any(K.q<2) || ~isreal(K.q)
        error('K.q should contain only integers bigger than 1')
    end
end
if ~isfield(K,'r') || ~nnz(K.r)
    K.r = zeros(1,0);
else
    K.r = K.r(:)';
    if any(K.r ~= floor(K.r)) || any(K.r<3) || ~isreal(K.r)
        error('K.r should contain only integers bigger than 2')
    end
end
if ~isfield(K,'s') || ~nnz(K.s)
    K.s = zeros(1,0);
else
    K.s = K.s(:)';
    if any(K.s ~= floor(K.s)) || any(K.s<1) || ~isreal(K.s)
        error('K.s should contain only positive integers')
    end
end
% As an alternative to the 'scomplex' flag, I've added a 'z' parameter
% containing a list of Hermitian semidefinite cone sizes. For now, this
% is just translated to 'scomplex' for you. In the future, we may use
% this internally *instead* of scomplex or rsdpN.
if ~isfield(K,'z') || ~nnz(K.z)
    K.z = zeros(1,0);
else
    K.z = K.z(:)';
    if any(K.z ~= floor(K.z)) || any(K.z<1) || ~isreal(K.z)
        error('K.z should contain only positive integers')
    end
end

N_f    = K.f;
N_l    = K.l;
N_fl   = N_f + N_l;
L_q    = length(K.q);
N_q    = sum(K.q);
L_r    = length(K.r);
N_r    = sum(K.r);
N_qr   = N_q + N_r;
L_qr   = L_q + L_r;
L_s    = length(K.s);
L_z    = length(K.z);
L_sz   = L_s + L_z;
N_s    = sum((K.s).^2);
N_z    = sum((K.z).^2);
N_sz   = N_s + N_z;
L_qrsz = L_qr + L_sz;
N_flqr = N_fl + N_qr;
N      = N_flqr + N_sz;

if ~isfield(K,'ycomplex') || isempty(K.ycomplex)
    K.ycomplex = zeros(1,0);
else
    K.ycomplex = sort(K.ycomplex(:))';
    K.ycomplex(find(~diff(K.ycomplex))+1) = [];
    if any(K.ycomplex ~= floor(K.ycomplex)) || any(K.ycomplex<1)
        error('K.ycomplex should contain only positive integers');
    elseif any(K.ycomplex > numel(b))
        error('Elements of K.ycomplex are out of range');
    end
end
if ~isfield(K,'xcomplex') || isempty(K.xcomplex)
    K.xcomplex = zeros(1,0);
else
    K.xcomplex = sort(K.xcomplex(:))';
    K.xcomplex(find(~diff(K.xcomplex))+1) = [];
    if any(K.xcomplex ~= floor(K.xcomplex)) || any(K.xcomplex<1)
        error('K.xcomplex should contain only positive integers');
    elseif any(K.xcomplex>N_flqr)
        error('Elements of K.xcomplex are out of range');
    end
end
if ~isfield(K,'scomplex') || isempty(K.scomplex)
    K.scomplex = zeros(1,0);
else
    K.scomplex = sort(K.scomplex(:))';
    K.scomplex(find(~diff(K.scomplex))+1) = [];
    if any(K.scomplex~=floor(K.scomplex)) || any(K.scomplex<1)
        error('K.scomplex should contain only positive integers');
    elseif any(K.scomplex>L_s)
        error('Elements of K.xcomplex are out of range');
    end
end
if L_z,
    K.s = [ K.s, K.z ];
    K.scomplex = [ K.scomplex, L_s + 1 : L_sz ];
    K.z = zeros(1,0);
    L_s = L_s + L_z;
    N_s = N_s + N_z;
    L_z = 0; %#ok
    N_z = 0; %#ok
end

% -------------------------------------------------------------------------
% Verify the size and validity of At, b, and c
% N = # variables
% -------------------------------------------------------------------------

% SeDuMi assumes that if At is not consistent with K but At' is, that the
% user supplied the transpose of the coefficient matrix. This introduces a
% rare ambiguity in the case where At happens to be square. Past versions
% of SeDuMi would not have allowed this, instead rejecting the case where
% m >= N; however, with complex variables, the situation is not quite as
% clear, and there may technically be cases where m >= N is acceptable.

if ndims(At) > 2 %#ok
    error('A must be a matrix');
elseif nnz(isnan(At)) || nnz(isinf(At)),
    error('A contains NaN or Inf');
elseif size(At,1) == N,
    % nothing
elseif size(At,2) == N,
    At = At';
else
    error('(At,K) size mismatch');
end
if all(size(b)>1)
    error('Parameter b must be a vector');
elseif any(isnan(b)) || any(isinf(b)),
    error('b contains NaN or Inf');
elseif length(b) ~= size(At,2),
    error('(At,b) size mismatch');
else
    b = b(:);
end
if all(size(c)>1)
    error('Parameter c must be a vector');
elseif any(isnan(c)) || any(isinf(c))
    error('c contains NaN or Inf');
elseif length(c) ~= N
    error('(c,K) size mismatch');
else
    c = c(:);
end

% -------------------------------------------------------------------------
% Save the standardized data for further use if needed
% -------------------------------------------------------------------------

if isfield(pars,'errors') && pars.errors==1
    origcoeff.At=At;
    origcoeff.c=c;
    origcoeff.b=b;
    origcoeff.K=K;
else
    origcoeff=[];
end

% -------------------------------------------------------------------------
% Flag diagonal SDP blocks for removal
% -------------------------------------------------------------------------

% There is some serious MATLAB trickery here (if I do say so myself) that
% merits explanation. "spattern" contains the indices of symbolically
% nonzero elements of the dual variable z = c - At * y. This is a single
% vector across all LMI constraints, so "sblk" tells us which indices
% belong to which constraints. The "rem" statement is what determines if
% a particular element is off-diagonal. If there is even one off-diagonal
% element in an SDP, then its corresponding element of "sdiag" is false.
%
% This new version of the analysis replaces the old preprocessSDP(), and
% seems inexpensive enough to apply to all LMIs regardless of size/count.
% The previous version was applied more sparingly.
%
% In theory one could do a more complex analysis of the block structure of
% an LMI, potentially breaking a larger into smaller ones. But this would
% likely be significantly more expensive.

if L_s && ( ~isfield(pars,'sdp') || pars.sdp ),
    ssiz        = (K.s).^2;
    strt        = cumsum([1,ssiz(1:end-1)]);
    sblk        = zeros(1,N_s);
    sblk(strt)  = 1;
    sblk        = cumsum(sblk);
    spattern    = find(c(N_flqr+1:N)~=0|any(At(N_flqr+1:N,:),2))';
    sblk        = sblk(spattern);
    sblk        = sblk(rem(spattern-strt(sblk),K.s(sblk)+1)~=0);
    sdiag       = true(1,L_s);
    sdiag(sblk) = false;
else
    % Even if we disable SDP processing, we're going to move 1x1 SDPs into
    % the nonnegative variable block. It just doesn't make sense to deploy
    % all of that SDP machinery for nonnegative variables.
    sdiag = K.s == 1;
end

% -------------------------------------------------------------------------
% Handle K.ycomplex by splitting apart the complex constraints into pairs
% of real constraints
% -------------------------------------------------------------------------

if ~isempty(K.ycomplex)
    b  = [ real(b) ; imag(b(K.ycomplex)) ];
    At = [ At, 1j * At(:,K.ycomplex) ];
else
    b  = real(b);
end

% -------------------------------------------------------------------------
% Find the locations of the the complex data, so we can convert into
% SeDuMi's internal format, which uses only MATLAB's real representation.
% -------------------------------------------------------------------------
% Strictly speaking, nonnegative variables, the first variable in a Lorentz
% cone, and the first two variables in a rotated Lorentz cone are real. But
% But SeDuMi has allowed them all to be specified as complex; the imaginary
% portions are interpreted as free variables. We have kept that behavior.

% This code replaces whichcpx.c in its entirety. It actually fixes a bug in
% rotated Lorentz cone handling that was probably never exercised.
if isempty(K.xcomplex) && isempty(K.scomplex),
    K.fcplx = zeros(1,0);
    K.qcplx = zeros(1,0);
    K.rcplx = zeros(1,0);
    scplx   = false(1,L_s);
    sreal   = ~sdiag;
    K.rsdpN = L_s;
    N_fc = 0;
    K.cdim = 0;
else
    xc = K.xcomplex;
    tt = xc <= N_fl;
    K.fcplx = xc(tt);
    xc = xc(~tt) -  N_fl;
    tt = xc <= N_q;
    K.qcplx = xc(tt);
    xc = xc(~tt) - N_q;
    tt = xc <= N_r;
    K.rcplx = xc(tt);
    if ~isempty(K.qcplx),
        ndxs = cumsum([1,K.q(1:end-1)]);
        t2 = any(bsxfun(@eq,K.qcplx,ndxs'),1);
        K.fcplx = [ K.fcplx, K.qcplx(t2) + N_fl ];
        K.qcplx(t2) = [];
        t2 = sum(bsxfun(@gt,K.qcplx,ndxs'),1);
        K.q = K.q + full(sparse(1,t2,1,1,L_q));
        K.qcplx = K.qcplx + (1:length(K.qcplx));
        N_q = N_q + length(K.qcplx);
    end
    if ~isempty(K.rcplx),
        ndxs = cumsum([1,K.r(1:end-1)]);
        t2 = any(bsxfun(@eq,K.rcplx,[ndxs,ndxs+1]'),1);
        K.fcplx = [ K.fcplx, K.rcplx(t2) + N_fl + N_q ];
        K.rcplx(t2) = [];
        t2 = sum(bsxfun(@gt,K.rcplx,ndxs'),1);
        K.r = K.r + full(sparse(1,t2,1,1,L_r));
        % This 2*t2 offset is required because QR makes two accesses
        % each of the first two elements of a rotated Lorentz cone.
        K.rcplx = K.rcplx + (1:length(K.rcplx)) + 2 * t2;
        N_r = N_r + length(K.rcplx);
    end
    N_fc = length(K.fcplx);
    N_f  = N_f + N_fc;
    N_fl = N_f + N_l; %#ok
    N_qr = N_q + N_r;
    scplx = false(1,L_s);
    scplx(K.scomplex&~sdiag) = true;
    sreal = ~scplx & ~sdiag;
    K.rsdpN = nnz(sreal);
    K.cdim = length(K.xcomplex) + sum(K.s(scplx).^2);
end

% -------------------------------------------------------------------------
% We have significantly rewritten this section of the code. This section
% constructs a sparse matrix that represents the following transformations:
%   --- Free variables split into differences of nonnegative variables; OR
%   --- Free variables placed in a Lorentz-cone
%   --- Rotated lorentz cones translated to standard Lorentz cones
%   --- Lorentz cones rearranged to trace block + norm-bound blocks
%   --- Conversion of diagonal SDPs to nonnegative variables
%   --- SDP coefficients moved to the lower triangle for increased sparsity
% This code replaces the rotlorenz, qreshape, and vectril MEX files.
% -------------------------------------------------------------------------

newL = 0;
newQ = zeros(1,0);
ii = {}; jj = {}; vv = {};

% Split free variables into the difference of nonnegatives
if ~isfield( pars, 'free' ) || pars.free == 2 && L_qrsz,
    pars.free = 1;
end
if N_f && ~pars.free,
    jt = [ 1 : K.f, K.fcplx ; 1 : K.f, K.fcplx ];
    vt = [ ones(1,K.f), -1j*ones(1,N_fc) ; -ones(1,K.f), 1j*ones(1,N_fc) ];
    ii{end+1} = 1 : 2 * N_f;
    jj{end+1} = jt(:)';
    vv{end+1} = vt(:)';
    newL = 2 * N_f;
    prep.freeL = N_f;
end

% Copy nonnegative variables without change
if K.l,
    ii{end+1} = newL + 1 : newL + K.l;
    jj{end+1} = K.f + 1 : K.f + K.l;
    vv{end+1} = ones(1,K.l);
    newL = newL + K.l;
end

% Convert diagonal SDPs to nonnegative variables
if any(sdiag),
    dsize = K.s(sdiag);
    sdpL = sum(dsize);
    prep.sdiag = dsize;
    jstrt = cumsum([N_flqr+1,K.s(1:end-1).^2]);
    jstrt = jstrt(sdiag);
    istrt = cumsum([1,dsize(1:end-1)]);
    dsize = dsize + 1;
    dblks = cumsum(full(sparse(1,istrt,1,1,sdpL)));
    ii{end+1} = newL + 1 : newL + sdpL;
    jj{end+1} = jstrt(dblks) + dsize(dblks) .* ( ( 1 : sdpL ) - istrt(dblks) );
    vv{end+1} = ones(1,sdpL);
    newL = newL + sdpL;
end

% Stuff free variables into a Lorentz cone
tr_off = newL;
nb_off = newL + L_qr;
if N_f && pars.free,
    tr_off = tr_off + 1;
    nb_off = nb_off + 1;
    ii{end+1} = nb_off + 1 : nb_off + N_f;
    jj{end+1} = [ 1 : K.f, K.fcplx ];
    vv{end+1} = [ ones(1,K.f), -1j*ones(1,N_fc) ];
    nb_off = nb_off + K.f;
    newQ = N_f + 1;
end

% Rearrange Lorentz cones to trace block + norm-bound blocks
if N_q,
    ndxs      = cumsum([1,K.q(1:end-1)]);
    it        = zeros(1,N_q);
    it(ndxs)  = tr_off + 1 : tr_off + L_q;
    it(it==0) = nb_off + 1 : nb_off + ( N_q - L_q );
    jt = K.f + K.l + 1 : K.f + K.l + N_q;
    vt = ones(1,N_q);
    if ~isempty(K.qcplx),
        jt = jt - cumsum(full(sparse(1,K.qcplx,1,1,N_q)));
        vt(K.qcplx) = -1j;
    end
    ii{end+1} = it(:)';
    jj{end+1} = jt;
    vv{end+1} = vt;
    tr_off    = tr_off + L_q;
    nb_off    = nb_off + N_q - L_q;
end

% Transform rotated Lorentz cones to standard Lorentz cones, and rearrange
% to trace block + norm-bound blocks.
if N_r,
    ndxr       = cumsum([1,K.r(1:end-1)]);
    ndxp       = ndxr + 2*(0:L_r-1);
    it         = zeros(1,N_r+2*L_r);
    it(ndxp)   = tr_off + 1 : tr_off + L_r;
    it(ndxp+1) = -1;
    it(ndxp+2) = it(ndxp);
    it(it==0)  = nb_off + 1 : nb_off + ( N_r - L_r );
    it(ndxp+1) = it(ndxp+3);
    jt = [ K.f + K.l + N_q + 1, ones(1,N_r+2*L_r-1) ];
    jt([ndxp+1,ndxp+3]) = 0;
    vt = ones(1,N_r+2*L_r);
    vt([ndxp,ndxp+1,ndxp+2]) = sqrt(0.5);
    vt(ndxp+3) = -sqrt(0.5);
    if ~isempty(K.rcplx),
        jt(K.rcplx) = 0;
        vt(K.rcplx) = -1j;
    end
    ii{end+1} = it;
    jj{end+1} = cumsum(jt);
    vv{end+1} = vt;
    nb_off = nb_off + N_r - L_r;
end

% Replace non-diagonal real SDP coefficients with tril(X) + tril(X',-1).
% This cuts the number of nonzeros approximately in half.
if K.rsdpN,
    dsize = K.s(sreal);
    sdpL  = sum(dsize.^2);
    jstrt = cumsum([N_flqr+1,K.s(1:end-1).^2]);
    jstrt = jstrt(sreal);
    istrt = cumsum([1,dsize(1:end-1).^2]);
    dblks = cumsum(full(sparse(1,istrt,1,1,sdpL)));
    istrt = istrt + nb_off;
    dsize = dsize(dblks);
    istrt = istrt(dblks);
    jndxs = ( nb_off + 1 : nb_off + sdpL ) - istrt;
    cols  = floor(jndxs ./ dsize);
    rows  = jndxs - dsize .* cols;
    ii{end+1} = max(rows,cols) + min(rows,cols) .* dsize + istrt;
    jj{end+1} = jndxs + jstrt(dblks);
    vv{end+1} = ones(1,sdpL);
    nb_off = nb_off + sdpL;
    clear dsize jstrt istrt dblks jndxs rows cols
end

% Replace Hermitian SDP coefficients with tril(X) + tril(X',-1). This one's
% a bit trickier because we have the real and complex values interleaved,
% and the imaginary values along the diagonal are zero.
if K.rsdpN < length(K.s),
    dsize = K.s(scplx);
    jsize = dsize .^ 2;
    sdpL  = 2 * sum(jsize);
    jstrt = cumsum([N_flqr+1,K.s(1:end-1).^2]);
    jstrt = jstrt(scplx);
    bstrt = cumsum([1,2*jsize(1:end-1)]);
    dblks = cumsum(full(sparse(1,bstrt,1,1,sdpL)));
    istrt = bstrt + nb_off;
    dsize = dsize(dblks);
    istrt = istrt(dblks);
    bndxs = ( nb_off + 1 : nb_off + sdpL ) - istrt;
    cols  = floor( bndxs ./ dsize );
    rows  = bndxs - dsize .* cols;
    imgv  = cols >= dsize;
    cols  = cols - imgv .* dsize;
    indxs = max(rows,cols) + min(rows,cols) .* dsize + imgv .* jsize(dblks) + istrt;
    vals  = ( 1 - 2 * ( cols > rows ) ) .* ( 1 - ( 1 + 1j ) .* imgv );
    keep  = ~imgv | ( rows ~= cols );
    jndxs = rows + cols .* dsize + jstrt(dblks);
    ii{end+1} = indxs(keep);
    jj{end+1} = jndxs(keep);
    vv{end+1} = vals(keep);
    clear dsize jsize jstrt bstrt istrt bndxs rows cols vals imgv keep
end

% Update free, nonnegative, and Lorentz variable counts
K.f = 0;
K.l = newL;
K.q = [ newQ, K.q, K.r ];
K.r = zeros(0,1);
K.s = [ K.s(:,~scplx&~sdiag), K.s(scplx&~sdiag) ];
K.rsdpN = nnz(~scplx&~sdiag);
K.N = K.l + sum(K.q) + sum(K.s(1:K.rsdpN).^2)+2*sum(K.s(K.rsdpN+1:end).^2);

% Create the artificial (x0,z0) variable for the self-dual model by
% appending a zero row to At and c. This is accomplished in QR by adding
% 1 to all of the row indices created above.
K.N  = K.N + 1;
K.l  = K.l + 1;
K.m  = length(b);

% Transform At, c
% The transformation matrix QR does not satisfy QR'*QR=I. But it does, in
% fact, serve as a reverse transformation:
%    --- For Lorentz cones, QR applies a permutation; QR' reverses it.
%    --- For rotated Lorentz cones, QR applies a unitary, self-adjoint
%        rotation on the first two variables, so QR' reverses it.
%    --- For split free variables, QR creates the positive and negative
%        parts; QR' combines them back together.
%    --- For free variables placed in a Lorentz cone, QR moves them to the
%        Lorentz cone block, adding an extra epigraph variable. QR' moves
%        them back and drops the extra variable.
%    --- For semidefinte cones, QR adds the strict upper triangle to the
%        strict lower triangle; QR' copies the strict lower triangle to the
%        strict upper triangle, ensuring symmetry.
%    --- QR adds a row for the self-dual variable; QR' removes it.
[dummy,ndxs] = sort(cellfun(@(x) x(1),jj)); %#ok
QR = sparse( horzcat(ii{ndxs})+1, horzcat(jj{ndxs}), horzcat(vv{ndxs}), K.N, length(c) );
At = real( sparse( QR * At ) );
c  = real( sparse( QR * c  ) );
b  = sparse( b );
prep.QR = QR;
clear ii jj vv

% -------------------------------------------------------------------------
% Now K has field K.{l,q,s}
% Generate a more detailed description of cone K:
% Let K.blkstart(i):K.blkstart(i+1)-1 give the index range of block i.
% Compute maxN=max(order), len=sum(order) for LORENTZ, real PSD, herm PSD
% yields: K.{l,q,s,rsdpN,blkstart,rLen,hLen,qMaxn,rMaxn,hMaxn}
% -------------------------------------------------------------------------
Ksr = K.s(1:K.rsdpN);
Ksc = K.s(K.rsdpN+1:end);
K.blkstart = cumsum([K.l+1,length(K.q)+length(K.r),K.q-1,Ksr.^2,2*Ksc.^2]);
K.rLen = sum(Ksr);
K.hLen = sum(Ksc);
K.qMaxn = max([0,K.q]);
K.rMaxn = max([0,Ksr]);
K.hMaxn = max([0,Ksc]);
K.mainblks = K.blkstart(cumsum([1 1 length(K.q)]));
K.qblkstart = K.blkstart(2:2+length(K.q));  % Also include blkend
K.sblkstart = K.blkstart(2+length(K.q):end);
K.lq = K.mainblks(end)-1;