From 3ad1994a7f4db8ac9990779622d7b542a6eb4062 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <brudfors@gmail.com>
Date: Mon, 14 Mar 2022 17:37:29 +0000
Subject: [PATCH 1/5] FEAT: tissue weights (JA)

---
 spm_mb_appearance.m |   6 +-
 spm_mb_classes.m    | 157 +++++++++++++++-------------
 spm_mb_fit.m        |  36 ++-----
 spm_mb_gmm.m        |   7 +-
 spm_mb_init.m       |  17 ++-
 spm_mb_io.m         |   2 +-
 spm_mb_merge.m      |   2 +-
 spm_mb_output.m     |   4 +-
 spm_mb_shape.m      | 250 ++++++++++++++++++++++++++++++++------------
 tbx_cfg_mb.m        |  19 ++--
 10 files changed, 313 insertions(+), 187 deletions(-)

diff --git a/spm_mb_appearance.m b/spm_mb_appearance.m
index a0cb5da..bc4bf23 100644
--- a/spm_mb_appearance.m
+++ b/spm_mb_appearance.m
@@ -21,7 +21,7 @@
 for n=1:numel(dat)
     if isfield(dat(n).model,'gmm')
         p    = dat(n).model.gmm.pop;
-        same = all(sum(diff(sett.gmm(p).pr{1},1,2).^2,1))==0;
+        same = ~all(sum(diff(sett.gmm(p).pr{1},1,2).^2,1));
         if same
             % Intensity priors are identical for all clusters
             % so need to break the symmetry.
@@ -276,7 +276,7 @@
         W  = sett.gmm(p).pr{3};
         for k=1:size(W,3)
             S        = inv(W(:,:,k));
-            W(:,:,k) = inv(0.999*S + 0.001*mean(diag(S))*eye(size(S)));
+            W(:,:,k) = inv(S*(1-1e-9) + 1e-9*mean(diag(S))*eye(size(S)));
         end
         sett.gmm(p).pr{3} = W;
 
@@ -624,7 +624,7 @@ function debug_show(img,img_is,modality,fig_title,do_show)
     c   = 1;  % Channel to show
     img = img(:,:,:,c);
 elseif strcmp(img_is,'template')
-    img = spm_mb_classes('template_k1',img,4);
+    img = spm_mb_classes('template_k1',img,[],4);
 end
 clim = [-Inf Inf];
 if modality == 2
diff --git a/spm_mb_classes.m b/spm_mb_classes.m
index 35e26bc..d024de0 100644
--- a/spm_mb_classes.m
+++ b/spm_mb_classes.m
@@ -7,11 +7,9 @@
 % P    - Updated tissue classes
 %
 % FORMAT [dat,P] = spm_mb_classes('update_cat',dat,mu,sett)
-% FORMAT lab     = spm_mb_classes('get_labels',dat,K1)
-% FORMAT cm      = spm_mb_classes('get_label_conf_matrix',cm_map,w,K1)
-% FORMAT l       = spm_mb_classes('LSE',mu,ax)
-% FORMAT mu      = spm_mb_classes('template_k1',mu,ax)
+% FORMAT l       = spm_mb_classes('LSE0',mu,ax)
 % FORMAT l       = spm_mb_classes('LSE1',mu,ax)
+% FORMAT mu      = spm_mb_classes('template_k1',mu,delta)
 %__________________________________________________________________________
 % Copyright (C) 2019-2020 Wellcome Centre for Human Neuroimaging
 
@@ -27,29 +25,85 @@
 %==========================================================================
 function [P,dat] = get_classes(dat,mu,sett)
 
+% Memory hungry. Needs to be addressed later.
+mu = add_delta(mu,dat.delta);
+
 if isfield(dat.model,'cat')
     % Categorical model
     [dat,P] = update_cat(dat,mu);
 elseif isfield(dat.model,'gmm')
     % GMM model
 
-    % Expand mu to include the background class and combine with labels if
-    % required.
-    mu  = template_k1(mu);
-    lab = get_labels(dat,size(mu,4));
-    if numel(lab)>1
-        % Add labels to template
-   %    mu = mu + lab;
-    end
-    clear lab
+    % Expand mu to include the background class.
+    mu1 = template_k1(mu);
     if sett.gmm(dat.model.gmm.pop).nit_appear >0
-        [dat,P] = spm_mb_appearance('update',dat,mu,sett);
+        [dat,P] = spm_mb_appearance('update',dat,mu1,sett);
     else
-        P       = exp(bsxfun(@minus,mu(:,:,:,1:(size(mu,4)-1)),LSE1(mu,4)));
+        P       = exp(bsxfun(@minus,mu1(:,:,:,1:(size(mu1,4)-1)),LSE1(mu1,4)));
     end
 else
     error('This should not happen');
 end
+if ~isempty(dat.delta)
+    [dat.delta,tmp] = update_delta(dat.delta,mu,P,sett.del_settings);
+    dat.E(1) = dat.E(1)+tmp;
+end
+%==========================================================================
+
+%==========================================================================
+function [delta,dE] = update_delta(delta,mu,P,del_settings)
+K = size(mu,4);
+L = (eye(K)-1/(K+1))*del_settings;
+H = L;
+g = L*delta(:);
+for k=1:size(mu,3)
+    [g1,H1] = gradhess1(mu(:,:,k,:),P(:,:,k,:));
+    g = g + double(reshape(sum(sum(g1,1),2),[K 1]));
+    H = H + double(reshape(sum(sum(H1,1),2),[K K]));
+end
+dE       = 0.5*delta(:)'*L*delta(:);
+delta(:) = delta(:) - H\g;
+%==========================================================================
+
+%==========================================================================
+function [g,H] = gradhess1(mu,P,delta)
+dm    = size(mu);
+K     = size(mu,4);
+if nargin>=3 && ~isempty(delta)
+    delta = reshape(delta,[1 1 1 K]);
+    mu    = bsxfun(@plus,mu,delta);
+end
+H     = zeros([dm(1:3),K,K]);
+g     = zeros([dm(1:3),K,1]);
+sig   = softmax0(mu);
+msk   = ~(all(isfinite(sig),4) & all(isfinite(P),4));
+for k=1:K
+    sig_k        = sig(:,:,:,k);
+    tmp          = sig_k - P(:,:,:,k);
+    tmp(msk)     = 0;
+    g(:,:,:,k)   = tmp;
+
+    tmp          = sig_k - sig_k.^2;
+    tmp(msk)     = 0;
+    H(:,:,:,k,k) = tmp;
+    for k1=(k+1):K
+        tmp           = -sig_k.*sig(:,:,:,k1);
+        tmp(msk)      = 0;
+        H(:,:,:,k,k1) = tmp;
+        H(:,:,:,k1,k) = tmp;
+    end
+end
+%==========================================================================
+
+%==========================================================================
+function P = softmax0(mu,ax)
+% safe softmax function (matches LSE0)
+
+if nargin<2, ax = 4; end
+mx  = max(mu,[],ax);
+E   = exp(bsxfun(@minus,mu,mx));
+den = sum(E,ax)+exp(-mx);
+P   = bsxfun(@rdivide,E,den);
 %==========================================================================
 
 %==========================================================================
@@ -62,73 +116,40 @@
 % Compute subject-specific categorical cross-entropy loss between
 % segmentation and template
 msk      = all(isfinite(P),4) & all(isfinite(mu),4);
-tmp      =  sum(P.*mu,4) - LSE(mu,4);
+tmp      =  sum(P.*mu,4) - LSE0(mu,4);
 dat.E(1) = -sum(tmp(msk(:)));
 dat.nvox =  sum(msk(:));
 %==========================================================================
 
 %==========================================================================
-function lab = get_labels(dat, K1)
-lab=0; return;
-if isempty(dat.lab), lab = 0; return; end
-
-% Load labels
-lab    = spm_mb_io('get_data', dat.lab.f);
-sk     = dat.samp;
-lab    = lab(1:sk(1):end, 1:sk(2):end, 1:sk(3):end);
-dm     = [size(lab) 1 1];
-lab    = round(lab(:));
-cm_map = dat.lab.cm_map; % cell array that defines the confusion matrix
-lab(~isfinite(lab) | lab<1 | lab>numel(cm_map)) = numel(cm_map) + 1; % Prevent crash
-
-% Get confusion matrix that maps from label value to (log) probability value
-cm = get_label_conf_matrix(cm_map, dat.lab.w, K1);
-
-% Build "one-hot" representation using confusion matrix
-lab = reshape(cm(lab,:), [dm(1:3) K1]);
-%==========================================================================
-
-%==========================================================================
-function cm = get_label_conf_matrix(cm_map, w, K1)
-% FORMAT CM = get_label_conf_matrix(cm_map, w, K1)
-% cm_map - Defines the confusion matrix
-% w      - Weighting probability
-% K1     - Number of classes
-% cm     - confusion matrix
-%
-% Build Rater confusion matrix for one subject.
-% This matrix maps template classes to manually segmented classes.
-% Manual labels often do not follow the same convention as the Template,
-% and not all regions may be labelled. Therefore, a manual label may
-% correspond to several Template classes and, conversely, one Template
-% class may correspond to several manual labels.
-
-% Parse function settings
-w  = min(max(w, 1e-7), 1-1e-7);
-L  = numel(cm_map);              % Number of labels
-cm = zeros([L+1, K1], 'single'); % Allocate confusion matrix (including unknown)
-for l=1:L % Loop over labels
-    ix            = false(1, K1);
-    ix(cm_map{l}) = true;
-    cm(l, ix)     = log(w/sum( ix));
-    cm(l,~ix)     = log((1-w)/sum(~ix));
+function mu1 = add_delta(mu,delta)
+if isempty(delta)
+    mu1 = mu;
+else
+    mu1 = bsxfun(@plus,mu,reshape(delta,[1 1 1 size(mu,4)]));
 end
 %==========================================================================
 
 %==========================================================================
-function l = LSE(mu,ax)
-% log-sum-exp function
+function l = LSE0(mu,ax)
+% Strictly convex log-sum-exp function
+% https://en.wikipedia.org/wiki/LogSumExp#A_strictly_convex_log-sum-exp_type_function
 if nargin<2, ax = 4; end
 mx = max(max(mu,[],ax),0);
 l  = log(exp(-mx) + sum(exp(bsxfun(@minus,mu,mx)),ax)) + mx;
 %==========================================================================
 
 %==========================================================================
-function mu = template_k1(mu,ax)
+function mu1 = template_k1(mu,delta,ax)
 % Expand a template to include the implicit background class
-if nargin<2, ax = 4; end
-lse  = LSE(mu,ax);
-mu   = cat(ax,bsxfun(@minus,mu,lse), -lse);
+if nargin>=2
+    mu1  = add_delta(mu,delta);
+else
+    mu1  = mu;
+end
+if nargin<3,ax=4; end
+lse  = LSE0(mu1,ax);
+mu1  = cat(ax,bsxfun(@minus,mu1,lse), -lse);
 %==========================================================================
 
 %==========================================================================
@@ -141,7 +162,3 @@
 
 %==========================================================================
 
-%==========================================================================
-
-%==========================================================================
-
diff --git a/spm_mb_fit.m b/spm_mb_fit.m
index 0d33aba..5aa585b 100644
--- a/spm_mb_fit.m
+++ b/spm_mb_fit.m
@@ -10,7 +10,7 @@
 %__________________________________________________________________________
 % Copyright (C) 2020 Wellcome Centre for Human Neuroimaging
 
-% $Id: spm_mb_fit.m 8196 2021-12-16 15:18:25Z john $
+% $Id: spm_mb_fit.m 8226 2022-02-24 10:44:46Z john $
 
 
 % Repeatable random numbers
@@ -74,7 +74,7 @@
 
 % Update affine only
 %--------------------------------------------------------------------------
-fprintf('Rigid (zoom=%d): %d x %d x %d\n',2^(numel(sz)-1),sett.ms.d);
+fprintf('Rigid (zoom=1/%d): %d x %d x %d\n',2^(numel(sz)-1),sett.ms.d);
 spm_plot_convergence('Init','Rigid Alignment','Objective','Iteration');
 E      = Inf;
 for it0=1:nit_aff
@@ -116,38 +116,14 @@
         countdown = 6;
     end
 end
-
-% Finish affine registration of any subjects that need a few more iterations
-for it0=1:3
-
-    if updt_mu
-        [mu,sett,dat,te,E] = iterate_mean(mu,sett,dat,te,nit_mu);
-    end
-
-    for n=1:numel(dat)
-        En = Inf;
-        for it1=1:nit_aff
-            oEn    = En;
-            dat(n) = spm_mb_shape('update_simple_affines',dat(n),mu,sett);
-            En     = sum(dat(n).E)/nvox(dat(n));
-            if abs(oEn-En) < sett.tol*0.2, break; end
-        end
-    end
-
-    E   = sum(sum(cat(2,dat.E),2),1) + te;  % Cost function after previous update
-    fprintf('%8.4f\n', E/nvox(dat));
-    spm_plot_convergence('Set',E/nvox(dat));
-    do_save(mu,sett,dat);
-end
-
 spm_plot_convergence('Clear');
-nit_mu = 2;
+nit_mu = 4;
 
 % Update affine and diffeo (iteratively decreases the template resolution)
 %--------------------------------------------------------------------------
 spm_plot_convergence('Init','Diffeomorphic Alignment','Objective','Iteration');
 for zm=numel(sz):-1:1 % loop over zoom levels
-    fprintf('\nzoom=%d: %d x %d x %d\n', 2^(zm-1), sett.ms.d);
+    fprintf('\nzoom=1/%d: %d x %d x %d\n', 2^(zm-1), sett.ms.d);
 
     if updt_mu
         dat = spm_mb_appearance('restart',dat,sett);
@@ -163,7 +139,7 @@
     if ~updt_mu
         mu = spm_mb_shape('shrink_template',mu0,Mmu,sett);
     else
-        [mu,sett,dat,te,E] = iterate_mean(mu,sett,dat,te,nit_mu);
+        [mu,sett,dat,te] = iterate_mean(mu,sett,dat,te,nit_mu);
     end
 
     if updt_aff
@@ -176,7 +152,7 @@
     end
     fprintf('\n');
 
-    nit_max = nit_zm0 + (zm - 1)*3;
+    nit_max = nit_zm0 + (zm - 1)*2;
     for it0=1:nit_max
 
         oE  = E/nvox(dat);
diff --git a/spm_mb_gmm.m b/spm_mb_gmm.m
index 51103e0..ebeaced 100644
--- a/spm_mb_gmm.m
+++ b/spm_mb_gmm.m
@@ -220,7 +220,9 @@
         muo     = mu(io,k);
         mum     = mu(im,k);
 
-        iAkmm   = inv(Ak(im,im));
+        iAkmm   = Ak(im,im);
+        iAkmm   = iAkmm + eye(size(iAkmm))*(1e-6*max(diag(iAkmm))+1e-30);
+        iAkmm   = inv(iAkmm);
         SA      = iAkmm*Ak(im,io);
 
         % 1) observed
@@ -830,7 +832,6 @@
 eta   = N  + eta0;         % Number of observations
 
 % Starting estimates (working with log(alpha0))
-la0    = log(a0)*ones(K,1);
 la0    = log(mean(Alpha,2));
 alpha0 = exp(la0);
 for it=1:100
@@ -848,7 +849,7 @@
     la0    = la0 - H\g;
     alpha0 = exp(la0);
 
-    if norm(la0-la0o).^2 <= norm(la0).^2*1e-9, break; end
+    if norm(la0-la0o)^2 <= norm(la0)^2*1e-9, break; end
 end
 alpha0 = alpha0 + eps;
 lb  = [(-eta0*(sum(gammaln(alpha0)) - gammaln(sum(alpha0))) - v0'*alpha0)
diff --git a/spm_mb_init.m b/spm_mb_init.m
index b7aa38c..d7c9363 100644
--- a/spm_mb_init.m
+++ b/spm_mb_init.m
@@ -5,7 +5,7 @@
 % Copyright (C) 2018-2020 Wellcome Centre for Human Neuroimaging
 
 
-% $Id: spm_mb_init.m 8086 2021-04-01 09:13:20Z john $
+% $Id: spm_mb_init.m 8220 2022-02-09 12:21:22Z john $
 
 [dat,sett] = mb_init1(cfg);
 
@@ -43,7 +43,7 @@
 if ~isempty(sett.aff)
     B   = spm_mb_shape('affine_bases',sett.aff);
 else
-    B   = zeros([3 3 0]);
+    B   = zeros([4 4 0]);
 end
 sett.B = B;
 sett   = rmfield(sett,'aff');
@@ -77,7 +77,7 @@
 
 cl  = cell(N,1);
 dat = struct('dm',cl, 'Mat',cl, 'samp',[1 1 1], 'onam','', 'odir','',...
-             'q',cl, 'v',cl, 'psi',cl, 'model',cl, 'lab',cl, 'E',cl,'nvox',cl);
+             'q',cl, 'v',cl, 'delta', zeros(1,K), 'psi',cl, 'model',cl, 'lab',cl, 'E',cl,'nvox',cl);
 n   = 0;
 
 % Process categorical data
@@ -113,6 +113,12 @@
         dat(n).odir = sett.odir;
         dat(n).v    = fullfile(dat(n).odir,['v_' dat(n).onam '.nii']);
         dat(n).psi  = fullfile(dat(n).odir,['y_' dat(n).onam '.nii']);
+        if isfinite(sett.del_settings)
+            dat(n).delta = zeros(1,K);
+        else
+            dat(n).delta = [];
+        end
+
 
         Kn = 0;
         for c=1:Nc
@@ -212,6 +218,11 @@
             dat(n).v    = fullfile(dat(n).odir,['v_' dat(n).onam '.nii']);
             dat(n).psi  = fullfile(dat(n).odir,['y_' dat(n).onam '.nii']);
 
+            if isfinite(sett.del_settings)
+                dat(n).delta = zeros(1,K);
+            else
+                dat(n).delta = [];
+            end
             cf = zeros(Nc,1);
             for c=1:Nc
                 cf(c) = size(f(1).dat,4);
diff --git a/spm_mb_io.m b/spm_mb_io.m
index a9166de..9d952b0 100644
--- a/spm_mb_io.m
+++ b/spm_mb_io.m
@@ -183,7 +183,7 @@ function save_template(mu,sett)
 
 if true
     % Softmax
-    mu  = spm_mb_shape('softmax',mu,4);
+    mu  = spm_mb_shape('softmax0',mu,4);
     d   = [size(mu) 1 1];
     [pth,nam,ext] = fileparts(sett.mu.create.mu);
     nam      = ['softmax' nam(3:end)];
diff --git a/spm_mb_merge.m b/spm_mb_merge.m
index fa777aa..3d6198d 100644
--- a/spm_mb_merge.m
+++ b/spm_mb_merge.m
@@ -129,7 +129,7 @@
 dm0       = dm0(1:3);
 msk       = ~isfinite(bb(1,:));
 bb(1,msk) = 1;
-bb1(1,:)  = round(max(bb(1,:),1));
+bb(1,:)   = round(max(bb(1,:),1));
 msk       = ~isfinite(bb(2,:));
 bb(2,msk) = dm0(msk);
 bb(2,:)   = round(min(bb(2,:),dm0));
diff --git a/spm_mb_output.m b/spm_mb_output.m
index 1cc4288..2d94266 100644
--- a/spm_mb_output.m
+++ b/spm_mb_output.m
@@ -194,7 +194,7 @@
 
     % Get warped tissue priors
     mun    = spm_mb_shape('pull1',mu,psi);
-    mun    = spm_mb_classes('template_k1',mun,4);
+    mun    = spm_mb_classes('template_k1',mun,datn.delta,4);
 
     % Integrate use of multiple Gaussians per tissue
     gam  = gmm.gam;
@@ -343,7 +343,7 @@
         % native space.
         mun = spm_mb_shape('pull1',mu,psi);
         clear mu
-        mun = spm_mb_shape('softmax',mun,4);
+        mun = spm_mb_shape('softmax0',mun,4);
         mun = cat(4,mun,max(1 - sum(mun,4),0));
     end
 
diff --git a/spm_mb_shape.m b/spm_mb_shape.m
index 7fdb0c1..1372572 100644
--- a/spm_mb_shape.m
+++ b/spm_mb_shape.m
@@ -6,13 +6,13 @@
 % FORMAT psi       = spm_mb_shape('compose',psi1,psi0)
 % FORMAT id        = spm_mb_shape('identity',d)
 % FORMAT dat       = spm_mb_shape('init_def',dat,sett.ms)
-% FORMAT l         = spm_mb_shape('LSE',mu,ax)
+% FORMAT l         = spm_mb_shape('LSE0',mu,ax)
 % FORMAT a1        = spm_mb_shape('pull1',a0,psi,r)
 % FORMAT [f1,w1]   = spm_mb_shape('push1',f,psi,d,r)
 % FORMAT sd        = spm_mb_shape('samp_dens',Mmu,Mn)
 % FORMAT varargout = spm_mb_shape('shoot',v0,kernel,args)
 % FORMAT mu1       = spm_mb_shape('shrink_template',mu,oMmu,sett)
-% FORMAT P         = spm_mb_shape('softmax',mu,ax)
+% FORMAT P         = spm_mb_shape('softmax0',mu,ax)
 % FORMAT E         = spm_mb_shape('template_energy',mu,sett)
 % FORMAT dat       = spm_mb_shape('update_affines',dat,mu,sett)
 % FORMAT [mu,dat]  = spm_mb_shape('update_mean',dat, mu, sett)
@@ -25,7 +25,7 @@
 %__________________________________________________________________________
 % Copyright (C) 2019-2020 Wellcome Centre for Human Neuroimaging
 
-% $Id: spm_mb_shape.m 8126 2021-07-22 17:58:11Z john $
+% $Id: spm_mb_shape.m 8226 2022-02-24 10:44:46Z john $
 [varargout{1:nargout}] = spm_subfun(localfunctions,varargin{:});
 %==========================================================================
 
@@ -45,7 +45,7 @@
 %==========================================================================
 function B = affine_bases(code)
 g     = regexpi(code,'(?<code>\w*)\((?<dim>\d*)\)','names');
-g.dim = str2num(g.dim);
+g.dim = str2double(g.dim);
 if numel(g.dim)~=1 || (g.dim ~=0 && g.dim~=2 && g.dim~=3)
     error('Can not use size');
 end
@@ -154,8 +154,9 @@
 %==========================================================================
 
 %==========================================================================
-function l = LSE(mu,ax)
-% log-sum-exp function
+function l = LSE0(mu,ax)
+% Strictly convex log-sum-exp function
+% https://en.wikipedia.org/wiki/LogSumExp#A_strictly_convex_log-sum-exp_type_function
 
 if nargin<2, ax = 4; end
 mx = max(mu,[],ax);
@@ -312,7 +313,7 @@
 Mzoom = Mmu\oMmu;
 if any(d0~=d) || norm(Mzoom-eye(4))>1e-4
     y      = affine(d0,Mzoom);
-    mu     = exp(bsxfun(@minus,mu,LSE(mu,4)));
+    mu     = exp(bsxfun(@minus,mu,LSE0(mu,4)));
     [mu,c] = push1(mu,y,d,1);
     e      = eps('single');
     mu     = bsxfun(@rdivide,mu,max(c,e));
@@ -321,8 +322,21 @@
 %==========================================================================
 
 %==========================================================================
-function P = softmax(mu,ax)
-% safe softmax function
+function P = softmax0_delta(mu,delta)
+if nargin>=2 && ~isempty(delta)
+    P = zeros(size(mu),'like',mu);
+    delta = reshape(delta,[1 1 1 numel(delta)]);
+    for k=1:size(mu,3) % Save memory by looping over planes
+        P(:,:,k,:) = softmax0(bsxfun(@plus,mu(:,:,k,:),delta));
+    end
+else
+    P = softmax0(mu,4);
+end
+%==========================================================================
+
+%==========================================================================
+function P = softmax0(mu,ax)
+% safe softmax function (matches LSE0)
 
 if nargin<2, ax = 4; end
 mx  = max(mu,[],ax);
@@ -390,9 +404,9 @@
 d     = sett.ms.d;
 Mmu   = sett.ms.Mmu;
 
-q    = double(datn.q);
-Mn   = datn.Mat;
-samp = datn.samp;
+q     = double(datn.q);
+Mn    = datn.Mat;
+samp  = datn.samp;
 [Mr,dM3] = spm_dexpm(q,B);
 dM   = zeros(12,size(B,3));
 for m=1:size(B,3)
@@ -440,8 +454,8 @@
 
 msk    = all(isfinite(f),4) & all(isfinite(mu1),4);
 mu1(~isfinite(mu1)) = 0;
-a      = mask(f - softmax(mu1,4),msk);
-[H,g]  = affine_hessian(mu1,G,a,single(msk),accel,samp);
+a      = mask(f - softmax0_delta(mu1,datn.delta),msk);
+[H,g]  = affine_hessian(mu1,G,a,single(msk),datn.delta,accel,samp);
 g      = double(dM'*g);
 H      = dM'*H*dM;
 H      = H + eye(numel(q))*(norm(H)*1e-6 + 0.001);
@@ -450,7 +464,7 @@
 %==========================================================================
 
 %==========================================================================
-function [H,g] = affine_hessian(mu,G,a,w,accel,samp)
+function [H,g] = affine_hessian(mu,G,a,w,delta,accel,samp)
 if nargin<6, samp = [1 1 1]; end
 d = [size(mu,1),size(mu,2),size(mu,3)];
 I = horder(3);
@@ -460,7 +474,7 @@
 for i=1:d(3)
     x{3} = x{3}*0+(i-1)*samp(3)+1;
     gv   = reshape(sum(bsxfun(@times,a(:,:,i,:),G(:,:,i,:,:)),4),[d(1:2) 1 3]);
-    Hv   = bsxfun(@times,w(:,:,i),velocity_hessian(mu(:,:,i,:),G(:,:,i,:,:),accel));
+    Hv   = bsxfun(@times,w(:,:,i),velocity_hessian(mu(:,:,i,:),G(:,:,i,:,:),delta,accel));
     for i1=1:12
         k1g   = rem(i1-1,3)+1;
         k1x   = floor((i1-1)/3)+1;
@@ -477,6 +491,8 @@
 %==========================================================================
 function [mu,dat] = update_mean(dat, mu, sett)
 
+[dat,mu] = adjust_delta(dat,mu);
+
 % Parse function settings
 accel       = sett.accel;
 mu_settings = sett.ms.mu_settings;
@@ -509,46 +525,135 @@
 %==========================================================================
 
 %==========================================================================
-function mu = gn_mu_update(mu,g,w,mu_settings,accel)
+function [dat,mu] = adjust_delta(dat,mu)
+% Zero-mean the deltas over subjects, and make a corresponding
+% adjustment to mu.
+mean_delta = mean(cat(1,dat.delta),1);
+if ~isempty(mean_delta)
+    for n=1:numel(dat)
+        dat(n).delta = dat(n).delta - mean_delta;
+    end
+    mean_delta = reshape(mean_delta,[1 1 1 size(mu,4)]);
+    mu         = bsxfun(@plus,mu,mean_delta);
+end
+%==========================================================================
+
+%==========================================================================
+function mu = gn_mu_update(mu,g,w,mu_settings,accel,nit)
+% Solve the problem
+% x = (L + H)\g
+% Regularisation L = kron(eye(K)-1/(K+1),L0)
+% Hessian H is constructed on the fly using mu and w.
+if nargin<6, nit   = [1 16 1]; end
+if nargin<5, accel = 0.8;      end
+s = softmax0(mu);
+x = zeros(size(s),'single');
+for it=1:nit(1)
+    for subit=1:nit(2), x = relax_mu1(x,s,g,w,mu_settings,accel); end
+    for subit=1:nit(3), x = relax_mu2(x,s,g,w,mu_settings,accel); end
+end
+mu = mu - x;
+%==========================================================================
+
+%==========================================================================
+function x = relax_mu2(x,s,g,w,mu_settings,accel)
 % Use Gauss-Seidel method for computing updates
 % https://en.wikipedia.org/wiki/Gauss%E2%80%93Seidel_method
+% This approach is for when the regularisation is high relative
+% to the Hessian of the data.
+
+if isempty(x), x = zeros(size(g),'single'); end
 spm_field('boundary',1);
-K   = size(mu,4);
-update_settings = [mu_settings(1:3) mu_settings(4:end)*(1-1/(K+1)) 2 2];
-if accel>0, s   = softmax(mu); end
-dmu = zeros(size(mu),'single');
-for it=1:10
-    for k=1:K
+d = size(g);
+K = d(4);
+for k=1:K
 
-        % Diagonal elements of Hessian
-        if accel==0
-            h_kk = (0.5*(1-1/(K+1)) + 1e-5)*w;
-        else
-            h_kk = (accel*(s(:,:,:,k)-s(:,:,:,k).^2) + (1-accel)*0.5*(1-1/(K+1)) + 1e-5).*w;
-        end
+    % Diagonal elements of Hessian
+    h_kk = hessel(k,k,K,accel,s).*w;
 
-        % Dot product betwen off-diagonals of likelihood Hessian and dmu
-        g_k = zeros([size(mu,1),size(mu,2),size(mu,3)],'single');
-        for k1=1:K
-            if k1~=k
-                % Off-diaginal elements of Hessian
-                if accel==0
-                    g_k = g_k - dmu(:,:,:,k1).*(0.5/(K+1));
-                else
-                    g_k = g_k - dmu(:,:,:,k1).*(accel*(s(:,:,:,k).*s(:,:,:,k1)) + ((1-accel)*0.5/(K+1)));
-                end
-            end
+    % Dot product betwen off-diagonals of likelihood Hessian and x
+    g_k = zeros(d(1:3),'single');
+    for k1=1:K
+        if k1~=k, g_k = g_k + x(:,:,:,k1).*hessel(k,k1,K,accel,s); end
+    end
+    g_k = g_k.*w;
+
+    % Dot product between off-diagonals of regularisation Hessian and x
+    g_k = g_k - spm_field('vel2mom', (sum(x,4)-x(:,:,:,k))/(K+1), mu_settings);
+
+    % Gauss-Seidel update
+    x(:,:,:,k) = spm_field(h_kk, g(:,:,:,k) - g_k, [mu_settings(1:3) mu_settings(4:end)*(1-1/(K+1)) 2 2]);
+end
+%==========================================================================
+
+%==========================================================================
+function x = relax_mu1(x,s,g,w,mu_settings,accel)
+% https://en.wikipedia.org/wiki/Jacobi_method
+% This approach is for when the regularisation is small relative
+% to the Hessian of the data.
+
+if isempty(x), x = zeros(size(g),'single'); end
+L  = operator(mu_settings);
+dc = L(1,1,1);
+d  = size(g);
+K  = d(4);
+spm_field('boundary',1);
+
+% Multiply with "off diagonal" part of matrix and compute the residual
+x  = spm_field('vel2mom',x,mu_settings) - x*dc;
+sx = sum(x,4)/(K+1);
+for k=1:K, x(:,:,:,k) = g(:,:,:,k) - (x(:,:,:,k) - sx); end
+
+% Divide by "diagonal" part
+H = zeros([d(1:2) 1 d(4)*(d(4)+1)/2],'single');
+for i=1:d(3)
+
+    % Construct Hessian for this slice, adding the diagona term from the regularisation
+    si = s(:,:,i,:);
+    wi = w(:,:,i);
+
+    % Diagonal components
+    for k=1:K, H(:,:,1,k) = hessel(k, k, K, accel, si).*wi + dc*(1-1/(K+1)); end
+
+    % Off-diagonal components
+    kk = K+1;
+    for k=1:K
+        for k1=(k+1):K
+            H(:,:,1,kk) = hessel(k, k1, K, accel, si).*wi - dc/(K+1);
+            kk = kk + 1;
         end
-        g_k = g_k.*w;
+    end
 
-        % Dot product between off-diagonals of regularisation Hessian and dmu
-        g_k = g_k - spm_field('vel2mom', (sum(dmu,4)-dmu(:,:,:,k))/(K+1), mu_settings);
+    % Update with only diagonal part of regularisation
+    x(:,:,i,:) = spm_field(H,x(:,:,i,:),[mu_settings(1:3) 0 0 0  1 1]);
+end
+%==========================================================================
 
-        % Gauss-Seidel update
-        dmu(:,:,:,k) = spm_field(h_kk, g(:,:,:,k) - g_k, update_settings);
+%==========================================================================
+function h = hessel(k,k1,K,accel,s)
+% Element of Hessian
+if k==k1
+    h = 0.5*(1-1/(K+1)) + 1e-5;
+    if accel>0
+        sk = s(:,:,:,k);
+        h  = accel*(sk-sk.^2) + (1-accel)*h;
+    end
+else
+    h = -0.5/(K+1);
+    if accel>0
+        h = -accel*s(:,:,:,k).*s(:,:,:,k1) + (1-accel)*h;
     end
 end
-mu = mu - dmu;
+%==========================================================================
+
+%==========================================================================
+function L = operator(mu_settings)
+b        = spm_field('boundary');
+spm_field('boundary',0);
+L        = zeros([5 5 5],'single');
+L(1,1,1) = 1;
+L        = spm_field('vel2mom',L,mu_settings);
+spm_field('boundary',b);
 %==========================================================================
 
 %==========================================================================
@@ -575,7 +680,7 @@
 spm_diffeo('boundary',1); % Neumann
 mu       = pull1(mu,psi);
 [f,datn] = spm_mb_classes(datn,mu,sett);
-[g,w]    = push1(softmax(mu,4) - f,psi,d,1);
+[g,w]    = push1(softmax0_delta(mu,datn.delta) - f,psi,d,1);
 %==========================================================================
 
 %==========================================================================
@@ -590,7 +695,7 @@
     % Update the affine parameters
     spm_diffeo('boundary',1);
     G  = spm_diffeo('grad',mu);
-    H0 = velocity_hessian(mu,G,accel);
+    H0 = velocity_hessian(mu,G,[],accel); % May need to use subject-specific deltas
 
     nw = get_num_workers(sett,dat,3*sett.K+4,3*sett.K+5);
     if nw > 1 && numel(dat) > 1 % PARFOR
@@ -616,7 +721,6 @@
         for n=1:numel(dat)
             dat(n).q = dat(n).q - mq;
         end
-         
     end
 
     % Update orientations in deformation headers when appropriate
@@ -629,6 +733,21 @@
 
 %==========================================================================
 function datn = update_simple_affines_sub(datn,mu,G,H0,sett)
+if numel(datn.E) >=1
+    eprev = datn.E(1)/datn.nvox;
+else
+    eprev = Inf;
+end
+
+for it=1:8
+    datn = update_simple_affines_sub1(datn,mu,G,H0,sett);
+    if eprev - datn.E(1)/datn.nvox < sett.tol, break; end
+    eprev = datn.E(1)/datn.nvox;
+end
+%==========================================================================
+
+%==========================================================================
+function datn = update_simple_affines_sub1(datn,mu,G,H0,sett)
 
 % Parse function settings
 B    = sett.B;
@@ -649,7 +768,7 @@
 psi      = affine(df,Mmu\Mr*Mn,samp);
 mu1      = pull1(mu,psi);
 [f,datn] = spm_mb_classes(datn,mu1,sett);
-[a,w]    = push1(f - softmax(mu1,4),psi,d,1);
+[a,w]    = push1(f - softmax0_delta(mu1,datn.delta),psi,d,1);
 clear mu1 psi f
 
 [H,g]    = simple_affine_hessian(mu,G,H0,a,w);
@@ -688,7 +807,6 @@
 function dat = update_velocities(dat,mu,sett)
 
 % Parse function settings
-accel     = sett.accel;
 groupwise = isa(sett.mu,'struct') && isfield(sett.mu,'create');
 if groupwise
     scal = 1-1/numel(dat);
@@ -699,22 +817,17 @@
 
 spm_diffeo('boundary',1); % Neumann boundary
 G  = spm_diffeo('grad',mu);
-H0 = velocity_hessian(mu,G,accel);
-if size(G,3) == 1
-    % Data is 2D -> add some regularisation
-    H0(:,:,:,3) = H0(:,:,:,3) + mean(reshape(H0(:,:,:,[1 2]),[],1));
-end
 
 nw = get_num_workers(sett,dat,4*sett.K+28,3*sett.K+7);
 if nw > 1 && numel(dat) > 1 % PARFOR
-    parfor(n=1:numel(dat),nw), dat(n) = update_velocities_sub(dat(n),mu,G,H0,sett,scal); end
+    parfor(n=1:numel(dat),nw), dat(n) = update_velocities_sub(dat(n),mu,G,sett,scal); end
 else % FOR
-    for n=1:numel(dat), dat(n) = update_velocities_sub(dat(n),mu,G,H0,sett,scal); end
+    for n=1:numel(dat), dat(n) = update_velocities_sub(dat(n),mu,G,sett,scal); end
 end
 %==========================================================================
 
 %==========================================================================
-function datn = update_velocities_sub(datn,mu,G,H0,sett,scal)
+function datn = update_velocities_sub(datn,mu,G,sett,scal)
 
 % Parse function settings
 B          = sett.B;
@@ -733,13 +846,18 @@
 Mat      = Mmu\Mr*Mn;
 df       = datn.dm;
 psi      = compose(get_def(datn,Mmu),affine(df,Mat,samp));
-mu       = pull1(mu,psi);
-[f,datn] = spm_mb_classes(datn,mu,sett);
-[a,w]    = push1(f - softmax(mu,4),psi,d,1);
-clear psi f mu
+mu1      = pull1(mu,psi);
+[f,datn] = spm_mb_classes(datn,mu1,sett);
+[a,w]    = push1(f - softmax0_delta(mu1,datn.delta),psi,d,1);
+clear psi f mu1
 
 g         = reshape(sum(bsxfun(@times,a,G),4),[d 3]);
-H         = bsxfun(@times,w,H0);
+H         = velocity_hessian(mu,G,datn.delta,sett.accel);
+H         = bsxfun(@times,w,H);
+if size(G,3) == 1
+    % Data is 2D -> add some regularisation
+    H(:,:,:,3) = H(:,:,:,3) + mean(reshape(H(:,:,:,[1 2]),[],1));
+end
 clear a w
 
 spm_diffeo('boundary',0); % Circulant boundary conditions
@@ -756,13 +874,13 @@
 %==========================================================================
 
 %==========================================================================
-function H = velocity_hessian(mu,G,accel)
+function H = velocity_hessian(mu,G,delta,accel)
 d  = [size(mu,1),size(mu,2),size(mu,3)];
 M  = size(mu,4);
 Ab = 0.5*(eye(M)-1/(M+1)); % See Bohning's paper
 H  = zeros([d 6],'single');
 for i=1:d(3)
-    if accel>0, s = softmax(mu(:,:,i,:),4); end
+    if accel>0, s = softmax0_delta(mu(:,:,i,:),delta); end
     H11 = zeros(d(1:2));
     H22 = H11;
     H33 = H11;
diff --git a/tbx_cfg_mb.m b/tbx_cfg_mb.m
index dee35eb..7acb939 100644
--- a/tbx_cfg_mb.m
+++ b/tbx_cfg_mb.m
@@ -3,7 +3,7 @@
 %__________________________________________________________________________
 % Copyright (C) 2019-2020 Wellcome Centre for Human Neuroimaging
 
-% $Id: tbx_cfg_mb.m 8087 2021-04-01 09:35:27Z mikael $
+% $Id: tbx_cfg_mb.m 8219 2022-02-09 09:42:10Z john $
 
 
 if ~isdeployed, addpath(fileparts(mfilename('fullpath'))); end
@@ -179,7 +179,7 @@
 pr_upd.labels  = {'Yes','No'};
 pr_upd.values = {{'b0_priors',{0.01,0.01}}, []};
 %pr_upd.values  = {{}, []};
-pr_upd.val     = {pr_upd.values{1}};
+pr_upd.val     = {pr_upd.values{2}};
 pr_upd.help    = {['Specify whether the Gaussian-Wishart priors should be updated at each iteration. ' ...
                    'Enabling this can slow down convergence if there are small numbers of subjects. ' ...
                    'If only one subject is to be modelled (using a pre-computed template), then ' ...
@@ -200,7 +200,7 @@
 pop.tag   = 'gmm';
 pop.name  = 'Pop. of scans';
 pop.val   = {chans, has_labels, pr,...
-             const('tol_gmm', 0.0005), const('nit_gmm_miss',32), const('nit_gmm',8), const('nit_appear', 4), const('mg_ix', [])};
+             const('tol_gmm', 0.0005), const('nit_gmm_miss',32), const('nit_gmm',8), const('nit_appear', 8), const('mg_ix', [])};
 pop.check = @check_pop;
 pop.help  = {'Information about a population of subjects that all have the same set of scans.',''};
 % ---------------------------------------------------------------------
@@ -330,14 +330,17 @@
 aff             = cfg_menu;
 aff.tag         = 'aff';
 aff.name        = 'Affine';
-aff.labels      = {'None', 'Translations', 'Rigid'};
-aff.values      = {'', 'T(3)', 'SE(3)'};
+aff.labels      = {'None', 'Translations', 'Rigid', 'Affine'};
+aff.values      = {'', 'T(3)', 'SE(3)', 'Aff(3)'};
 aff.val         = {'SE(3)'};
 aff.help        = {[...
 'Type of affine transform to use in the model, which may be either ' ...
 'none, translations only (T(3)) or rigid body (SE(3)). The fitting ' ...
 'begins with affine registration, before continuing by interleaving ' ...
-'affine and diffeomorphic registrations over multiple spatial scales.'],''};
+'affine and diffeomorphic registrations over multiple spatial scales.' ...
+'Note that the ``Affine'''' option is likely to throw up lots of ' ...
+'warnings about ``QFORM0 representation has been rounded'''', which ' ...
+'can mostly be ignored.'],''};
 % ---------------------------------------------------------------------
 
 % ---------------------------------------------------------------------
@@ -346,7 +349,7 @@
 dff.name        = 'Shape regularisation';
 dff.strtype     = 'e';
 dff.num         = [1 5];
-dff.val         = {[0.00001 0 0.4 0.1 0.4]};
+dff.val         = {[0.0001 0 0.4 0.1 0.4]};
 dff.help        = {[...
 'Specify the regularisation settings for the diffeomorphic registration. ' ...
 'These consist of a vector of five values, which penalise different ' ...
@@ -397,7 +400,7 @@
 mb             = cfg_exbranch;
 mb.tag         = 'run';
 mb.name        = 'Fit Multi-Brain model';
-mb.val         = {mu_prov, aff, dff, onam, odir, segs, pops,...
+mb.val         = {mu_prov, aff, dff, const('del_settings',Inf), onam, odir, segs, pops,...
                    const('accel',0.8), const('min_dim', 8), const('tol',0.001),...
                    const('sampdens',2),const('save',true),const('nworker',0)};
 mb.prog        = @run_mb;

From 0334b7732677c3e7d61b2d47a391848fc56d18e8 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <brudfors@gmail.com>
Date: Tue, 15 Mar 2022 15:08:56 +0000
Subject: [PATCH 2/5] FEAT: clean_gwc

---
 spm_mb_output.m | 107 ++++++++++++++++++++++++++++++++++++++++++------
 tbx_cfg_mb.m    |  18 ++++----
 2 files changed, 102 insertions(+), 23 deletions(-)

diff --git a/spm_mb_output.m b/spm_mb_output.m
index 2d94266..f2ebef8 100644
--- a/spm_mb_output.m
+++ b/spm_mb_output.m
@@ -51,7 +51,7 @@
              'bb',cfg.bb,...
              'odir',cfg.odir,...
              'fwhm',cfg.fwhm);
-opt.proc_zn = cfg.proc_zn;
+opt.clean_gwc = cfg.clean_gwc;
 
 if nw > 1 && numel(dat) > 1 % PARFOR
     fprintf('Write output: ');
@@ -77,8 +77,8 @@
 %==========================================================================
 
 %==========================================================================
-% PostProcMRF()
-function zn = PostProcMRF(zn,Mn,strength,nit)
+% mrf_apply()
+function zn = mrf_apply(zn,Mn,strength,nit)
 if nargin < 4, nit = 10; end
 P   = zeros(size(zn),'uint8');
 G   = ones([size(zn,4),1],'single')*strength;
@@ -113,7 +113,7 @@
 fwhm        = opt.fwhm;      % FWHM for smoothing of warped tissues
 vx          = opt.vx;        % Template space voxel size
 bb          = opt.bb;        % Template space bounding box
-proc_zn     = opt.proc_zn;   % Function for processing native space responsibilities
+clean_gwc   = opt.clean_gwc; % Settings for cleaning up tissue classes
 
 cl   = cell(1,1);
 resn = struct('inu',cl,'i',cl,'mi',cl,'c',cl,'wi',cl, ...
@@ -220,16 +220,12 @@
 
     if mrf > 0
         % Ad-hoc MRF clean-up of segmentation
-        zn = PostProcMRF(zn,Mn,mrf);
+        zn = mrf_apply(zn,Mn,mrf);
     end
 
-    if iscell(proc_zn) && ~isempty(proc_zn) && isa(proc_zn{1},'function_handle')
-        % Applies a function that processes the native space responsibilities
-        try
-            zn = proc_zn{1}(zn);
-        catch
-            warning('Incorrect definition of out.proc_zn, no processing performed.')
-        end
+    if clean_gwc.do == true
+        % Ad-hoc clean-up of GM, WM and CSF
+        zn = do_clean_gwc(zn, clean_gwc.gm, clean_gwc.wm, clean_gwc.csf, clean_gwc.level);
     end
 
     if any(write_tc(:,1) == true)
@@ -467,3 +463,90 @@ function write_nii(f,img,M,descrip,typ)
 create(Nii);
 Nii.dat(:,:,:,:,:,:) = img;
 %==========================================================================
+
+%==========================================================================
+function zn = do_clean_gwc(zn,gm,wm,csf,level)
+if nargin < 2, gm = 1; end
+if nargin < 3, wm = 2; end
+if nargin < 4, csf = 3; end
+if nargin < 5, level = 1; end
+
+ixt = struct('gm',gm,'wm',wm,'csf',csf);
+b = sum(zn(:,:,:,ixt.wm),4);
+
+% Build a 3x3x3 seperable smoothing kernel
+kx=[0.75 1 0.75];
+ky=[0.75 1 0.75];
+kz=[0.75 1 0.75];
+sm=sum(kron(kron(kz,ky),kx))^(1/3);
+kx=kx/sm; ky=ky/sm; kz=kz/sm;
+
+% Erosions and conditional dilations
+th1 = 0.15;
+if level==2, th1 = 0.2; end
+niter  = 32;
+niter2 = 32;
+for j=1:niter
+    if j>2
+        th       = th1;
+    else
+        th       = 0.6;
+    end  % Dilate after two its of erosion
+    for i=1:size(b,3)
+        gp       = double(sum(zn(:,:,i,ixt.gm),4));
+        wp       = double(sum(zn(:,:,i,ixt.wm),4));
+        bp       = double(b(:,:,i));
+        bp       = (bp>th).*(wp+gp);
+        b(:,:,i) = bp;
+    end
+    spm_conv_vol(b,b,kx,ky,kz,-[1 1 1]);
+end
+
+% Also clean up the CSF.
+if niter2 > 0
+    c = b;
+    for j=1:niter2
+        for i=1:size(b,3)
+            gp       = double(sum(zn(:,:,i,ixt.gm),4));
+            wp       = double(sum(zn(:,:,i,ixt.wm),4));
+            cp       = double(sum(zn(:,:,i,ixt.csf),4));
+            bp       = double(c(:,:,i));
+            bp       = (bp>th).*(wp+gp+cp);
+            c(:,:,i) = bp;
+        end
+        spm_conv_vol(c,c,kx,ky,kz,-[1 1 1]);
+    end
+end
+
+th = 0.05;
+for i=1:size(b,3)
+    slices = cell(1,size(zn,4));
+    for k1=1:size(zn,4)
+        slices{k1} = double(zn(:,:,i,k1));
+    end
+    bp           = double(b(:,:,i));
+    bp           = ((bp>th).*(sum(cat(3,slices{ixt.gm}),3)+sum(cat(3,slices{ixt.wm}),3)))>th;
+    for i1=1:numel(ixt.gm)
+        slices{ixt.gm(i1)} = slices{ixt.gm(i1)}.*bp;
+    end
+    for i1=1:numel(ixt.wm)
+        slices{ixt.wm(i1)} = slices{ixt.wm(i1)}.*bp;
+    end
+
+    if niter2>0
+        cp           = double(c(:,:,i));
+        cp           = ((cp>th).*(sum(cat(3,slices{ixt.gm}),3)+sum(cat(3,slices{ixt.wm}),3)+sum(cat(3,slices{ixt.csf}),3)))>th;
+
+        for i1=1:numel(ixt.csf)
+            slices{ixt.csf(i1)} = slices{ixt.csf(i1)}.*cp;
+        end
+    end
+    tot       = zeros(size(bp))+eps;
+    for k1=1:size(zn,4)
+        tot   = tot + slices{k1};
+    end
+    for k1=1:size(zn,4)
+        zn(:,:,i,k1) = slices{k1}./tot;
+    end
+end
+%==========================================================================
\ No newline at end of file
diff --git a/tbx_cfg_mb.m b/tbx_cfg_mb.m
index 7acb939..6097fbf 100644
--- a/tbx_cfg_mb.m
+++ b/tbx_cfg_mb.m
@@ -629,16 +629,12 @@
 % ---------------------------------------------------------------------
 
 % ---------------------------------------------------------------------
-proc_zn        = cfg_entry;
-proc_zn.tag    = 'proc_zn';
-proc_zn.name   = 'Process responsibilities';
-proc_zn.help   = {'Function for processing native space responsibilities, ' ...
-                  'given as a function handle @(x) foo(x). The argument (x) is of ' ...
-                  'size(x) = [1, 4], where the first three dimensions are the size ' ...
-                  'of the image and the last dimension is the number of segmentation ' ...
-                  'classes (K + 1).'};
-proc_zn.val    = {{}};
-proc_zn.hidden = true;
+clean_gwc        = cfg_entry;
+clean_gwc.tag    = 'clean_gwc';
+clean_gwc.name   = 'GWC clean';
+clean_gwc.help   = {'Ad-hoc clean-up of GM, WM and CSF.'};
+clean_gwc.val    = {struct('do',false,'gm',1,'wm',2,'csf',3,'level',1)};
+clean_gwc.hidden = true;
 % ---------------------------------------------------------------------
 
 % ---------------------------------------------------------------------
@@ -656,7 +652,7 @@
 out      = cfg_exbranch;
 out.tag  = 'out';
 out.name = 'Output';
-out.val  = {res_file, i, mi, wi, wmi, inu, c, wc, mwc, sm, mrf, fwhm, bb, vox, proc_zn, odir};
+out.val  = {res_file, i, mi, wi, wmi, inu, c, wc, mwc, sm, mrf, fwhm, bb, vox, clean_gwc, odir};
 out.prog = @spm_mb_output;
 out.help = {[...
 'When ``Fit Multi-Brain model'' is run, the resulting model fit contains ' ...

From 78de0590b7595c281f38da2bc4360edd46317693 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <brudfors@gmail.com>
Date: Tue, 15 Mar 2022 15:09:09 +0000
Subject: [PATCH 3/5] FIX(affine_basis)

---
 spm_mb_shape.m | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/spm_mb_shape.m b/spm_mb_shape.m
index 1372572..fa32622 100644
--- a/spm_mb_shape.m
+++ b/spm_mb_shape.m
@@ -104,7 +104,7 @@
         B(3,2,6) = -1;
     case 'Aff'
         % Aff(3) - Affine
-        B        = zeros(4,4,6);
+        B        = zeros(4,4,12);
         B(1,1,1) =  1;
         B(2,1,2) =  1;
         B(3,1,3) =  1;

From 27748bf9274cbb8d3928dfac1a594cd0ce1d1a19 Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <brudfors@gmail.com>
Date: Tue, 15 Mar 2022 16:18:43 +0000
Subject: [PATCH 4/5] FEAT: Bohnings bound on the Hessian of tw update

---
 spm_mb_classes.m | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/spm_mb_classes.m b/spm_mb_classes.m
index d024de0..38cf2a7 100644
--- a/spm_mb_classes.m
+++ b/spm_mb_classes.m
@@ -45,19 +45,19 @@
     error('This should not happen');
 end
 if ~isempty(dat.delta)
-    [dat.delta,tmp] = update_delta(dat.delta,mu,P,sett.del_settings);
+    [dat.delta,tmp] = update_delta(dat.delta,mu,P,sett.del_settings,sett.accel);
     dat.E(1) = dat.E(1)+tmp;
 end
 %==========================================================================
 
 %==========================================================================
-function [delta,dE] = update_delta(delta,mu,P,del_settings)
+function [delta,dE] = update_delta(delta,mu,P,del_settings,accel)
 K = size(mu,4);
 L = (eye(K)-1/(K+1))*del_settings;
 H = L;
 g = L*delta(:);
 for k=1:size(mu,3)
-    [g1,H1] = gradhess1(mu(:,:,k,:),P(:,:,k,:));
+    [g1,H1] = gradhess1(mu(:,:,k,:),P(:,:,k,:),accel);
     g = g + double(reshape(sum(sum(g1,1),2),[K 1]));
     H = H + double(reshape(sum(sum(H1,1),2),[K K]));
 end
@@ -66,9 +66,10 @@
 %==========================================================================
 
 %==========================================================================
-function [g,H] = gradhess1(mu,P,delta)
+function [g,H] = gradhess1(mu,P,delta,accel)
 dm    = size(mu);
 K     = size(mu,4);
+Ab    = 0.5*(eye(K)-1/(K+1)); % Bohnings bound on the Hessian
 if nargin>=3 && ~isempty(delta)
     delta = reshape(delta,[1 1 1 K]);
     mu    = bsxfun(@plus,mu,delta);
@@ -82,12 +83,11 @@
     tmp          = sig_k - P(:,:,:,k);
     tmp(msk)     = 0;
     g(:,:,:,k)   = tmp;
-
-    tmp          = sig_k - sig_k.^2;
+    tmp          = (sig_k - sig_k.^2)*accel + (1-accel)*Ab(k,k);
     tmp(msk)     = 0;
     H(:,:,:,k,k) = tmp;
     for k1=(k+1):K
-        tmp           = -sig_k.*sig(:,:,:,k1);
+        tmp           = (-sig_k.*sig(:,:,:,k1))*accel + (1-accel)*Ab(k,k1);
         tmp(msk)      = 0;
         H(:,:,:,k,k1) = tmp;
         H(:,:,:,k1,k) = tmp;

From 50b3d9e765768db8c80c5f6894879a86a46a55ac Mon Sep 17 00:00:00 2001
From: Mikael Brudfors <brudfors@gmail.com>
Date: Wed, 16 Mar 2022 15:58:38 +0000
Subject: [PATCH 5/5] FIX(tissue_weight)

---
 spm_mb_classes.m | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/spm_mb_classes.m b/spm_mb_classes.m
index 38cf2a7..fcdcaa8 100644
--- a/spm_mb_classes.m
+++ b/spm_mb_classes.m
@@ -52,12 +52,13 @@
 
 %==========================================================================
 function [delta,dE] = update_delta(delta,mu,P,del_settings,accel)
+% disp(exp(delta)/sum(exp(delta)))
 K = size(mu,4);
 L = (eye(K)-1/(K+1))*del_settings;
 H = L;
 g = L*delta(:);
 for k=1:size(mu,3)
-    [g1,H1] = gradhess1(mu(:,:,k,:),P(:,:,k,:),accel);
+    [g1,H1] = gradhess1(mu(:,:,k,:),P(:,:,k,:),delta,accel);
     g = g + double(reshape(sum(sum(g1,1),2),[K 1]));
     H = H + double(reshape(sum(sum(H1,1),2),[K K]));
 end