/
spm_mb_classes.m
165 lines (147 loc) · 5.61 KB
/
spm_mb_classes.m
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
function varargout = spm_mb_classes(varargin)
% Get tissue classes
% FORMAT [P,dat] = spm_mb_classes(dat,mu,sett)
% dat - Data structure for a subject
% mu - Warped template data
% sett - Settings
% P - Updated tissue classes
%
% FORMAT [dat,P] = spm_mb_classes('update_cat',dat,mu,sett)
% 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
% $Id: spm_mb_classes.m 8057 2021-02-09 18:41:58Z john $
if isa(varargin{1},'char')
[varargout{1:nargout}] = spm_subfun(localfunctions,varargin{:});
else
[varargout{1:nargout}] = get_classes(varargin{:});
end
%==========================================================================
%==========================================================================
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.
mu1 = template_k1(mu);
if sett.gmm(dat.model.gmm.pop).nit_appear >0
[dat,P] = spm_mb_appearance('update',dat,mu1,sett);
else
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,sett.accel);
dat.E(1) = dat.E(1)+tmp;
end
%==========================================================================
%==========================================================================
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,:),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
dE = 0.5*delta(:)'*L*delta(:);
delta(:) = delta(:) - H\g;
%==========================================================================
%==========================================================================
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);
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)*accel + (1-accel)*Ab(k,k);
tmp(msk) = 0;
H(:,:,:,k,k) = tmp;
for k1=(k+1):K
tmp = (-sig_k.*sig(:,:,:,k1))*accel + (1-accel)*Ab(k,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);
%==========================================================================
%==========================================================================
function [dat,P] = update_cat(dat,mu)
% Categorical model
P = spm_mb_io('get_data',dat.model.cat.f);
sk = dat.samp;
P = P(1:sk(1):end,1:sk(2):end,1:sk(3):end,:);
% 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) - LSE0(mu,4);
dat.E(1) = -sum(tmp(msk(:)));
dat.nvox = sum(msk(:));
%==========================================================================
%==========================================================================
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 = 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 mu1 = template_k1(mu,delta,ax)
% Expand a template to include the implicit background class
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);
%==========================================================================
%==========================================================================
function l = LSE1(mu,ax)
% log-sum-exp (including final class) function
if nargin<2, ax = 4; end
mx = max(mu,[],ax);
l = log(sum(exp(bsxfun(@minus,mu,mx)),ax)) + mx;
%==========================================================================
%==========================================================================