 function pot = potential_fun(ptype, delta, param, varargin)
%function pot = potential_fun(ptype, delta, param, varargin)
%|
%| Define roughness penalty potential functions (as strum).
%|
%| The penalty will have the form
%|	R(x) = sum_k w_k * potential_k([Cx]_k, delta_k)
%| where w_k is provided elsewhere, not here!
%|
%| in
%|	ptype		quad broken huber hyper2 hyper3 cauchy qgg2 genhub
%|			lange1 lange3 (fair) geman&mcclure gf1 li98cfs
%|			Recommended: 'hyper3'
%|			To list all possible choices, use:
%|			potential_fun('list') for "smooth" options
%|			potential_fun('list1') for "non-smooth" options
%|	delta		scalar, or image-sized array;
%|			"cutoff" parameter for edge-preserving regularization
%|	param		optional additional parameter(s) for some choices:
%|				'gf1' (generalized Fair) : [a b]
%|				'qgg2' : q
%|				'genhub' & 'stevenson94dpr' : [p q]
%|				'table1' : [dz, dpot([1:K] * dz)]
%|
%| option
%|	'dummy'		include extra dummy argument for backward compatibility
%|			e.g., pot.potk([], C*x). default: 0
%|
%| out
%|	pot		strum object, with data: delta and param
%|	methods:
%|		pot.potk(C*x)	potential function value
%|		pot.wpot(C*x)	potential 'weights' (aka half-quad. curvatures)
%|		pot.dpot(C*x)	potential derivative
%|		pot.shrink(b, reg)	proximal (shrinkage) operation:
%|					argmin_z 1/2 |z - b|^2 + reg * pot(z)
%|		pot.plot()	plot all of the above functions
%|
%| trick: for ptype 'gf1-fit', the param argument should be:
%|	{'name of potential to fit', points, param}
%| and this function returns the [a b] parameters needed for a subsequent
%| call with ptype 'gf1'
%|
%| Copyright 2004-5-18, Jeff Fessler, University of Michigan

if nargin < 1, ir_usage, end
if streq(ptype, 'test') ir_potential_fun_test, return, end
if streq(ptype, 'test_lpp') ir_potential_fun_test_lpp, return, end

if streq(ptype, 'list') % return list of all "smooth" choices
	pot = {'quad', 'huber', 'huber2', ...
	'hyper2', 'hyper3', 'cauchy', ...
	'qgg2', 'genhub', 'lange1', 'lange3', 'geman&mcclure', 'gf1', ...
	'table0', 'table1', 'li98cfs'};
%	'stevenson94dpr' omit this one; use 'genhub' instead
return
end

if streq(ptype, 'list1') % return list of 'l1' type choices (mostly non-smooth)
	pot = {'l0', 'l1', 'lpp', 'tav', 'broken', 'fair-l1'};
return
end

if ~isvar('delta'), delta = []; end
if ~isvar('param'), param = []; end

arg.dummy = false;
arg.scale = 1;
arg = vararg_pair(arg, varargin);

if streq(ptype, 'gf1-fit') % trick: just return parameters for this case
	pot = ir_potential_fun_gf1_fit(param{:});
return
end

[pot.type, pot.delta, pot.param, potk, wpot, dpot, shrink] = ...
	ir_potential_fun_parse(ptype, delta(:), param(:));

if arg.scale ~= 1
	potk = @(pot, z) arg.scale * potk(pot, z);
	dpot = @(pot, z) arg.scale * dpot(pot, z);
	wpot = @(pot, z) arg.scale * wpot(pot, z);
end

if arg.dummy
	potk = @(pot, dum, z) potk(pot, z);
	dpot = @(pot, dum, z) dpot(pot, z);
	wpot = @(pot, dum, z) wpot(pot, z);
end

meth = {'potk', potk, 'wpot', wpot, 'dpot', dpot, 'shrink', shrink, ...
	'plot', @ir_potential_fun_plot};
pot = strum(pot, meth);

end % potential_fun()

% ir_potential_fun_gf1_fit()
% trick: gf1-fit means match gf1 to named potential at given points [0, s1, s2]
% where the given point values are s = z / delta, i.e., in "delta" units
% param: {'name of potential to fit', points, param}
function ab = ir_potential_fun_gf1_fit(ptype, sv, param)
sv = sv(:);
pot = potential_fun(ptype, 1, param); % trick: delta=1 here
pt = pot.wpot(sv);

% ab = [1 1; -pt']' \ ((pt-1) ./ sv);

s1 = sv(1);
s2 = sv(2);
w1 = pt(1);
w2 = pt(2);
ab(1) = (w2 * s2 * (1 - w1) - w1 * s1 * (1 - w2)) ...
		/ ((w1 - w2) * s1 * s2);
ab(2) = (s2 * (1 - w1) - s1 * (1 - w2)) ...
		/ ((w1 - w2) * s1 * s2);
end % ir_potential_fun_gf1_fit()

% ir_potential_fun_parse()
function [ptype, delta, param, potk, wpot, dpot, shrink] = ...
	ir_potential_fun_parse(ptype, delta, param)

dpot = [];
shrink = [];

% trick: huber2 is just a huber with delta / 2
% so that weighting function drops to 1/2 at delta, like hyper3 etc.
if streq(ptype, 'huber2')
	ptype = 'huber';
	delta = delta / 2;
end

% trick: hyper3 is just a hyperbola with delta scaled by sqrt(3)
% to approximately "match" the properties of 'cauchy' (old erroneous 'hyper')
if streq(ptype, 'hyper3')
	ptype = 'hyper2';
	delta = delta / sqrt(3);
end

if streq(ptype, 'gf1') && param(1) == 0
	if param(2) ~= 1
		fail 'only b=1 makes sense for gf1 with a=0'
	end
	ptype = 'lange3'; param = []; % trick: gf1 with a=0 and b=1 is lange3
end

switch ptype

% quadratic potential function
case 'quad'
	potk = @(pot, z) (abs(z).^2) / 2;
	wpot = @(pot, z) ones(size(z), class(z));
	dpot = @(pot, z) z;
	shrink = @(pot, b, reg) b ./ (1 + reg);

% broken parabola
case 'broken'
	potk = @(pot, z) min(z.^2, pot.delta.^2)/2;
	wpot = @(pot, z) ir_tonumeric(abs(z) < pot.delta, z);
	dpot = @(pot, z) z .* (abs(z) < pot.delta);
	shrink = @(pot, b, reg) ir_broken_shrink(b, reg, pot.delta);

% l0 (hard threshold)
case 'l0'
	potk = @(pot, z) ir_tonumeric(z ~= 0, z);
	wpot = @(pot, z) nan(size(z), class(z));
	dpot = @(pot, z) nan(size(z), class(z));
	shrink = @(pot, b, reg) b .* (abs(b) > sqrt(2*reg));

% absolute value (i.e., l1, for soft threshold)
case {'abs', 'l1'}
	potk = @(pot, z) abs(z);
	wpot = @(pot, z) 1 ./ abs(z); % invalid at z=0
	dpot = @(pot, z) sign(z);
	shrink = @(pot, b, reg) sign(b) .* max(abs(b) - reg, 0);

% l_p^p aka generalized gaussian (gg)
case 'lpp'
	potk = @(pot, z) abs(z).^param;
	dpot = @(pot, z) param * sign(z) .* abs(z).^(param-1); % invalid at z=0
	wpot = @(pot, z) param .* abs(z).^(param-2); % invalid at z=0
	switch param
	case 0
		potk = @(pot, z) ir_tonumeric(z ~= 0, z);
		dpot = @(pot, z) zeros(size(z), class(z)); % meaningless
		shrink = @(pot, b, reg) b .* (abs(b) > sqrt(2*reg));
	case 1
		potk = @(pot, z) abs(z);
		dpot = @(pot, z) sign(z); % invalid at t=0
		shrink = @(pot, b, reg) sign(b) .* max(abs(b) - reg, 0);
	case 0.5 % p=1/2
		shrink = @(pot, b, reg) ir_lpp_1_2_shrink(b, reg);
	case 4/3. % p=4/3
		shrink = @(pot, b, reg) ir_lpp_4_3_shrink(b, reg);
	case 1.5 % p=3/2
		shrink = @(pot, b, reg) sign(b) .* ...
			( sqrt((3/4*reg).^2 + abs(b)) - 3/4*reg ).^2;
		% from eqn. (4.5) in chaux:07:avf
	%	z + 9 * reg.^2 * sign(z) .* (1 - sqrt(1 + 16*abs(z)/(9*reg^2))) / 8;

	otherwise
		warn('no shrink for p=%g', param)
	end

% truncated absolute value
case 'tav'
	potk = @(pot, z) min(abs(z), pot.delta);
	wpot = @(pot, z) (1 ./ abs(z)) .* (abs(z) < pot.delta); % bad at t=0
	dpot = @(pot, z) sign(z) .* (abs(z) < pot.delta);
	shrink = @(pot, b, reg) ir_tav_shrink(b, reg, pot.delta);

% huber potential function
case 'huber'
	potk = @(pot, z) huber_pot(z, pot.delta);
	wpot = @(pot, z) huber_wpot(z, pot.delta);
	dpot = @(pot, z) huber_dpot(z, pot.delta);
	shrink = @(pot, b, reg) ir_huber_shrink(b, reg, pot.delta);

% cauchy penalty: d^2 / 2 * log(1 + (t/d)^2) (not convex!)
case 'cauchy'
	potk = @(pot, z) pot.delta.^2 / 2 .* log(1 + abs(z ./ pot.delta).^2);
	wpot = @(pot, z) 1 ./ (1 + abs(z ./ pot.delta).^2);
	dpot = @(pot, z) z ./ (1 + abs(z ./ pot.delta).^2);
	shrink = @(pot, b, reg) cauchy_shrink(b, reg, pot.delta);

% Geman&McClure penalty: d^2 / 2 * |z/d|^2 / (1 + |z/d|^2)
% Not convex!
case 'geman&mcclure'
	potk = @(pot, z) pot.delta.^2 / 2 .* abs(z/pot.delta).^2 ./ (1 + abs(z ./ pot.delta).^2);
	wpot = @(pot, z) 1 ./ (1 + abs(z ./ pot.delta).^2).^2;
	dpot = @(pot, z) z ./ (1 + abs(z ./ pot.delta).^2).^2;

% gf1: Generalized Fair 1st-order
% wpot(z) = (1 + a * |z/d|) / (1 + b * |z/d|)
case 'gf1'
	potk = @(pot, z) gf1_potk(z, pot.delta, pot.param(1), pot.param(2));
	wpot = @(pot, z) (1 + pot.param(1) .* abs(z ./ pot.delta)) ...
		./ (1 + pot.param(2) .* abs(z ./ pot.delta));
	shrink = @(pot, b, reg) ...
		ir_gf1_shrink(b, reg, pot.delta, pot.param(1), pot.param(2));

% hyperbola penalty: d^2 * [ sqrt(1 + (z/d)^2) - 1 ]
case 'hyper2'
	potk = @(pot, z) pot.delta.^2 .* (sqrt(1 + abs(z ./ pot.delta).^2) - 1);
	wpot = @(pot, z) 1 ./ sqrt(1 + abs(z ./ pot.delta).^2);
	dpot = @(pot, z) z ./ sqrt(1 + abs(z ./ pot.delta).^2);

case 'hyper'
	error 'use "cauchy" or "hyper3" not "hyper" now'

% Lange1 penalty
case 'lange1'
	potk = @(pot, z) abs(z).^2 / 2 ./ (1+abs(z./pot.delta));
	wpot = @(pot, z) (1 + abs(z ./ pot.delta) / 2) ./ (1 + abs(z ./ pot.delta)).^2;

% Lange3 penalty
case {'lange3', 'fair'}
	potk = @(pot, z) pot.delta.^2 .* (abs(z./pot.delta) - log(1+abs(z./pot.delta)));
	wpot = @(pot, z) 1 ./ (1 + abs(z ./ pot.delta));
	dpot = @(pot, z) z ./ (1 + abs(z ./ pot.delta));
	% caution: no built-in shink, use 'fair-l1' if you want shrink!

% Fair potential "rounded corner" approximation to l1
case 'fair-l1'
	potk = @(pot, z) abs(z) - pot.delta .* log(1+abs(z./pot.delta));
	wpot = @(pot, z) 1 ./ (pot.delta + abs(z));
	dpot = @(pot, z) z ./ (pot.delta + abs(z));
	shrink = @(pot, b, reg) fair_l1_shrink(b, reg, pot.delta);

% li98cfs
case 'li98cfs'
	% f = @(x) atan(x) / x - 0.5; fsolve(f, 2.3)
	delta = delta / 2.3311;
	potk = @(pot, z) ir_li98cfs_potk(z, pot.delta);
	wpot = @(pot, z) ir_li98cfs_wpot(z, pot.delta);

% qgg2: q-generalized gaussian for p=2, due to Thibault, Sauer, Bouman
% q = "param", same as lange1 when q=1
case 'qgg2'
	potk = @(pot, z) z.^2 / 2 ./ (1+abs(z./pot.delta).^(2-pot.param));
	wpot = @(pot, z) (1 + abs(z ./ pot.delta).^(2-pot.param) * pot.param ...
		 / 2) ./ (1 + abs(z ./ pot.delta).^(2-pot.param)).^2;

% genhub : generalized Huber (switch between two generalized gaussians)
% same as Huber when p=2 and q=1
% p is power near 0, q is asymptotic power
case 'genhub'
	potk = @(pot, z) ir_genhub_potk(z, pot.delta, pot.param(1), pot.param(2));
	wpot = @(pot, z) ir_genhub_wpot(z, pot.delta, pot.param(1), pot.param(2));

% from stevenson:94:dpr
% p = param(1), q = param(2), same as Huber when p=2 and q=1 ???
case 'stevenson94dpr'
	warn('Potential %s does not work well; use "genhub" instead', ptype)
	potk = @(pot, z) ir_stevenson94dpr_potk(z, pot.delta, pot.param(1), pot.param(2));
	wpot = @(pot, z) ones(size(z), class(z)); % bogus attempt at upper bound

% tabulate derivative of pot at z_k = dz * k for k=1,...,K
% param is samples of derivative dpot(z_k)
% dpot(0) = 0, and use sample-and-hold interpolation of dpot()
% except use linear interpolation over [0 dz]
case 'table0'
	if ~isnumeric(param) || numel(param) < 10
%	if ~iscell(param) || numel(param) ~= 2
%		fail 'for table0 param must be {dz, dpot([1:K]*dz}'
		fail 'for table0 param must be [dz, dpot([1:K]*dz)]'
	end
%	param = ir_table0_setup(param{1}, col(param{2}));
	param = ir_table0_setup(param(1), col(param(2:end)));
	potk = @(pot, z) ir_table0_potk(z, pot.param);
	wpot = @(pot, z) ir_table0_wpot(z, pot.param);
	shrink = @(pot, b, reg) ir_table0_shrink(b, reg, pot.param);

% tabulate derivative of pot at z_k = dz * k for k=1,...,K
% param is samples of derivative dpot(z_k)
% dpot(0 = 0, and use linear interpolation of dpot(t)
case 'table1'
	if ~isnumeric(param) || numel(param) < 10
%	if ~iscell(param) || numel(param) ~= 2
%		fail 'for table1 param must be {dz, dpot([1:K]*dz}'
		fail 'for table1 param must be [dz, dpot([1:K]*dz)]'
	end
%	param = ir_table1_setup(param{1}, col(param{2}));
	param = ir_table1_setup(param(1), col(param(2:end)));
	potk = @(pot, z) ir_table1_potk(z, pot.param);
	wpot = @(pot, z) ir_table1_wpot(z, pot.param);
	shrink = @(pot, b, reg) ir_table1_shrink(b, reg, pot.param);

otherwise
	fail('Unknown potential "%s"', ptype)
end

if isempty(dpot) % default is z * wpot(z)
	dpot = @(pot, z) z .* wpot(pot, z);
end
if isempty(shrink)
	shrink = @ir_potential_fun_shrink;
end

end % ir_potential_fun_parse()

% ir_potential_fun_shrink()
% find argmin_z 1/2 |z - b|^2 + reg * pot(z)
% default method uses fzero() which will be slow!
function out = ir_potential_fun_shrink(pot, b, reg)
out = zeros(size(b), class(b));
opt = optimset('tolx', 1e-7);
for ii=1:numel(b)
	a = abs(b(ii));
	s = sign(b(ii));
	cost = @(z) 0.5 * (z - a).^2 + reg * pot.potk(z);
%	dfun = @(z) z - a + reg * pot.dpot(z);
	try
	%	out(ii) = s * fzero(dfun, a);
	%	out(ii) = s * fminsearch(cost, a, opt);
		z0 = s * fminsearch(cost, 0, opt);
		z1 = s * fminsearch(cost, a, opt);
		if cost(z0) < cost(z1)
			out(ii) = z0;
		else
			out(ii) = z1;
		end
	catch
		whos
		keyboard
	end
end % for
end % ir_potential_fun_shrink

% ir_table0_setup()
% dz	[1]	z spacing
% dk	[K]	dpot([1:K] * dz)
function out = ir_table0_setup(dz, dk)
sk = dk(1)*dz/2 + dz * [0; cumsum(dk(1:end-1))]; % [K]
out.dz = dz;
out.dk = dk; % [1:K] samples of dpot
out.sk = sk; % [K] cumulative sums for pot(z)
K = numel(dk);
out.K = K;
end % ir_table0_setup

% ir_table0_potk()
function out = ir_table0_potk(z, param)
dz = param.dz;
dk = param.dk;
sk = param.sk;
z = abs(z);
k = floor(z / dz);
k = min(k, param.K); % at most K
big = z > dz;
out = zeros(size(z), class(z));
out(~big) = 0.5 * dk(1) / dz * z(~big).^2;
k = k(big);
out(big) = sk(k) + dk(k) .* (z(big) - k * dz);
end % ir_table0_potk

% ir_table0_wpot()
function out = ir_table0_wpot(z, param)
dz = param.dz;
dk = param.dk; % [K]
z = abs(z);
k = floor(z / dz);
k = min(k, param.K); % at most K
out = repmat(dk(1) / dz, size(z));
big = k > 0;
k = k(big);
z = z(big);
out(big) = dk(k) ./ z;
end % ir_table0_wpot

% ir_table0_shrink()
function out = ir_table0_shrink(b, reg, param)
K = param.K;
dk = param.dk;
dz = param.dz;
zk = (1:K)' * dz;
bk = zk + reg * dk; % must be done here because depends on reg
ck = bk + dz;
tmp = [0; col([bk ck]')];
out = [zk zk+dz];
out = [0; col([zk zk+dz]')];
out = sign(b) .* interp1(tmp, out, abs(b), 'linear', 'extrap');
end % ir_table0_shrink

% ir_table1_setup()
% dz	[1]	z spacing
% dk	[K]	dpot([1:K] * dz)
function out = ir_table1_setup(dz, dk)
K = numel(dk);
k = [0:K]'; % [K+1]
zk = dz * k; % [K+1]
dk0 = [0; dk]; % [K+1] prepend sample at zero
ck = [diff(dk0); 0] / dz; % [K+1] curvatures 0:K
tmp = (dk0 - zk .* ck) * dz ...
	+ dz^2/2 * ck .* ((k+1).^2 - k.^2); % [K+1]
sk = cumsum( [0; tmp] );
out.dz = dz;
out.dk = dk; % [1:K] samples of dpot
out.ck = ck; % [K+1] curvatures 0:K
out.sk = sk; % [K] cumulative sums for pot(t)
out.K = K;
end % ir_table1_setup

% ir_table1_potk()
function out = ir_table1_potk(z, param)
sk = param.sk;
dz = param.dz;
ck = param.ck;
dk0 = [0; param.dk]; % [K+1] prepend sample at zero
z = abs(z);
k = floor(z / dz);
k = min(k, param.K); % at most K
sk = sk(1+k); % matlab indexing
ck = ck(1+k); % matlab indexing
dk = dk0(1+k);
zk = k * dz;
ck = reshape(ck, size(z));
dk = reshape(dk, size(z));
sk = reshape(sk, size(z));
out = sk + (dk - zk .* ck) .* (z - zk) + ck / 2 .* (z.^2 - zk.^2);
end % ir_table1_potk

% ir_table1_wpot()
function out = ir_table1_wpot(z, param)
dz = param.dz;
ck = param.ck;
dk = param.dk;
z = abs(z);
k = floor(z / dz);
k = min(k, param.K); % at most K
out = repmat(dk(1) / dz, size(z));
big = k > 0;
k = k(big);
z = z(big);
dk = reshape(dk(k), size(z));
ck = reshape(ck(k+1), size(z));
out(big) = (dk + (z - k * dz) .* ck) ./ z;
end % ir_table1_wpot

% ir_table1_shrink()
% unfortunately this routine works only for a single scalar reg value.
function out = ir_table1_shrink(z, reg, param)
K = param.K;
zk = (1:K)' * param.dz;
bk = zk + reg * param.dk; % must be done here because depends on reg
out = sign(z) .* interp1([0; bk], [0; zk], abs(z), 'linear', 'extrap');
end % ir_table1_shrink

% gf1_potk()
% gf1: generalized fair 1st-order potential
function pot = gf1_potk(z, delta, a, b)
atd = abs(z ./ delta);
pot = delta.^2 ./ (2 * b.^3) * ...
	(2 * b.^2 .* atd + a .* b.^2 .* atd.^2 ...
	- 2 * a .* b .* atd + 2 * (a-b) .* log(1 + b .* atd));

if 0 % symbolic check
	syms x
	syms a positive
	syms b positive
	syms z positive
	int(x*(1+a*x) / (1+b*x), x, 0, z)
end
end % gf1_potk

% ir_broken_shrink()
function out = ir_broken_shrink(z, reg, delta)
out = z ./ (1 + reg);
big = delta * (1 + reg) < abs(z);
out(big) = z(big);
end % ir_broken_shrink

% cauchy_shrink()
function out = cauchy_shrink(z, reg, delta)
z = z(:);
coef = ones(numel(z),1);
coef = [1/delta^2*coef -z/delta^2 (1+reg)*coef -z];
out = zeros(size(z));
for ii=1:numel(z)
	tmp = roots(coef(ii,:));
	pick = tmp == real(tmp); % empirically, 3rd root is often real
	if sum(pick) ~= 1 % if multiple real roots, empirically pick largest
		pick = imax(abs(tmp));
	end
	out(ii) = tmp(pick);
end
end % cauchy_shrink

% ir_huber_shrink()
function out = ir_huber_shrink(z, reg, delta)
out = z ./ (1 + reg);
big = delta .* (1 + reg) < abs(z);
if numel(reg) > 1, reg = reg(big); end
if numel(delta) > 1, delta = delta(big); end
z = z(big);
out(big) = z .* (1 - reg .* delta ./ abs(z));
end % ir_huber_shrink

% fair_l1_shrink()
function out = fair_l1_shrink(z, reg, delta)
out = sign(z) .* (abs(z) - (delta + reg) ...
	+ sqrt( (delta + reg - abs(z)).^2 + 4 * delta .* abs(z) )) ./ 2;
end % fair_l1_shrink

% ir_gf1_shrink()
function out = ir_gf1_shrink(z, reg, delta, a, b)
u = a / delta;
v = b / delta;
out = sign(z) .* (v .* abs(z) - (1+reg) ...
	+ sqrt( (1 + reg - v.*abs(z)).^2 + 4 * (v + reg .* u) .* abs(z) )) ...
	./ (2 * (v + reg .* u));
end % ir_gf1_shrink

% ir_lpp_1_2_shrink()
% for l_p with p=1/2
% for a > 0 and z = t^2 the minimizer solves t^3 - a t + reg/2 = 0
function out = ir_lpp_1_2_shrink(b, reg)
sb = sign(b); a = abs(b);
%reg = 1;
%a = 50; roots([1 0 -a reg/2])
%a = linspace(0, 4, 401);
out = zeros(size(b), class(b));
%big = true(size(b));
%{
% loop method
for ib=1:numel(b)
	tmp = roots([1 0 -a(iz) reg/2])'
	out(ib) = max(tmp).^2;
end
%}
%{
% hyperbolic method for one real root:
big = 27 * (reg/2)^2 > 4 * a.^3;
t0 = -2 * sqrt(a/3) .* cosh(1/3 * acosh(3/2 * (reg/2) ./ a .* sqrt(3./a)));
%}
%{
% https://en.wikipedia.org/wiki/Cubic_function#Vieta.27s_substitution
% vieta solves t^3 + p t + q = 0 using t = w - p / (3w), i.e. p=-a and q=reg/2
% for which w^6 + q w^3 - p^3/27 = 0
% i.e. (w^3)^2 + reg/2 (w^3) + a^3/27 = 0
w3 = (-reg/2 + sqrt((reg/2).^2 - 4 * a.^3/27)) / 2; % w^3 solution by quad form
w1 = w3 .^ (1/3);
t0 = w1 + a ./ (3*w1);
out(big) = t0(big).^2;
%}
% trigonometric method for three real roots:
% https://en.wikipedia.org/wiki/Cubic_function#Trigonometric_method_for_three_real_roots
t0 = 2 * sqrt(a/3) .* cos(1/3 * acos(3*(reg/2)/2./(-a) .* sqrt(3./a)));
a_1_2 = 3/2 * reg^(2/3);
big = a > a_1_2;
out(big) = t0(big).^2;
% plot(z, out, '-o', z, z, ':', z, z-reg, '--')
out = out .* sb;
end % ir_lpp_1_2_shrink

% ir_lpp_4_3_shrink()
% for l_p with p=4/3
% from eqn. (4.5) in chaux:07:avf
function out = ir_lpp_4_3_shrink(b, reg)
sb = sign(b); a = abs(b);
x = sqrt(a.^2 + 256 * reg^3 / 729);
out = sb .* (a + 4 * reg / (3 * 2^1/3) * ((x - a).^1/3 - (x + a).^1/3));
end % ir_lpp_4_3_shrink

% ir_tav_shrink()
function out = ir_tav_shrink(z, reg, delta)
out = zeros(size(z), class(z));
big = reg < abs(z) & abs(z) < reg + delta;
out(big) = z(big) .* (1 - reg ./ abs(z(big)));
big = reg + delta <= abs(z);
out(big) = z(big);
end % ir_tav_shrink

% ir_li98cfs_potk()
function pot = ir_li98cfs_potk(z, d)
pot = d.^2 .* ((z ./ d) .* atan(z ./ d) - 0.5 * log(1 + (z ./ d).^2));
end % ir_li98cfs_potk

% ir_genhub_potk()
function pot = ir_genhub_potk(z, d, p, q)
pot = 0.5 * abs(z) .^ p .* (abs(z) <= d) ...
 + 0.5 * (p ./ q .* d .^ (p-q) .* abs(z) .^ q ...
 + (1 - p ./ q) .* d .^ p) .* (abs(z) > d);
end % ir_genhub_potk

% ir_genhub_wpot()
function pot = ir_genhub_wpot(z, d, p, q)
pot = p / 2 .* (d .^ (p-q)) .* (abs(z) .^ (q-2)) .* (abs(z) > d);
ii = abs(z) <= d;
pot(ii) = p / 2 .* (abs(z(ii)) .^ (p-2));
%pot = p / 2 .* abs(t) .^ (p-2) .* (abs(t) <= d) ...
% + p / 2 .* d .^ (p-q) .* abs(t) .^ (q-2) .* (abs(t) > d);
end % ir_genhub_wpot

% ir_stevenson94dpr_potk()
function pot = ir_stevenson94dpr_potk(z, d, p, q)
% pr [d p q]
tmp1 = (0.5 * abs(z) .^ p) .* (abs(z) <= d);
tmp2 = 0.5 * ( (p .* (d .^ (p-1)) .* abs(z) - p .* (d .^ p) ...
 + (1 ./ q) .^ (1 ./ (q-1)) ) .^ q ...
 + d .^ p - (1 ./ q) .^ (q ./ (q-1)) ) .* (abs(z) > d);
pot = tmp1 + tmp2;
end % ir_stevenson94dpr_potk

% ir_tonumeric()
% convert x to type of y
function out = ir_tonumeric(x, y)
switch class(y)
case 'double'
	out = double(x);
case 'single'
	out = single(x);
otherwise
	fail('unknown type %s', class(y))
end
end % ir_tonumeric

% ir_potential_fun_plot()
function dummy = ir_potential_fun_plot(pot)
z = linspace(-1,1,101)*2;
if isvar('pot.delta')
	z = z * pot.delta;
end
if ~im, return, end
im plc 2 2
im subplot 1
plot(z, pot.potk(z), '-', z, z.^2/2, ':', z, abs(z), ':')
axis([min(z) max(z) minmax(pot.potk(z))'])

im subplot 2
plot(z, pot.dpot(z), '-', z, z, ':')
axis([min(z) max(z) minmax(pot.dpot(z))'])

im subplot 3
plot(z, pot.wpot(z), '-', z, 1+0*z, ':')
axis([min(z) max(z) 0 1])

im subplot 4
reg = 1;
tmp = pot.shrink(z, reg);
plot(z, tmp, '-', z, z, ':')
axis([min(z) max(z) minmax(tmp)'])

dummy = [];
end % ir_potential_fun_plot

% ir_potential_fun_test()
% test routine
% examine potential functions after rescaling.
function ir_potential_fun_test

delta = 10; zmax = 4 * delta; reg = 1.5 * delta;
plist = potential_fun('list');
%plist = {'l1', 'fair-l1'}; delta = 0.5;
%plist = {'quad', 'li98cfs', 'hyper3', 'huber2'}; % show li98cfs roughly hyper3
%plist = {'quad', 'genhub', 'huber', 'stevenson94dpr'};
%plist = {'genhub'}
%plist = {'hyper3', 'qgg2', 'huber'}; delta = 10; zmax = 50;
%plist = {'lange3', 'qgg2'}; zmax = 200;
%plist = {'qgg2', 'gf1-fit'};
%plist = potential_fun('list1'); delta = 0.5;
%plist = {plist{:}, 'quad', 'gf1', 'huber', 'broken'}; % todo: more!
%plist = {'qgg2', 'table1', 'table0'};
%plist = {'qgg2', 'table1'}; fname = 'fig_reg_pot_table1_qgg2';
%plist = {'qgg2', 'table0'}; fname = 'fig_reg_pot_table0_qgg2';
%plist = {'qgg2', 'table1', 'gf1'}; fname = 'fig_reg_pot_table0_gf1_qgg2';
%plist = {'cauchy', 'l1'};
zz = zmax * linspace(-1, 1, 2001)';
bb = 3 * max(delta + reg, delta * (1+reg)) * linspace(-1, 1, 301)';
ps = [];
lshrink = {};
for ii=1:numel(plist)
	ptype = plist{ii}; % pr ptype
	if streq(ptype, 'quad')
		leg{ii} = ptype;
	else
		leg{ii} = [ptype sprintf(', $\\delta = %g$', delta)];
	end

	switch ptype
	case {'gf1', 'gf1-fit'}
		param = potential_fun('gf1-fit', nan, {'qgg2', [1 10], 1.2});
		ptype = 'gf1'; % trick
		leg{ii} = [leg{ii} sprintf(' %.3g %.4f', param(1), param(2))];
	case 'qgg2'
		param = 1.2;
		leg{ii} = [leg{ii} sprintf(', $q = %g$', param)];
	case 'genhub'
		param = [2.0 1.2];
		leg{ii} = [leg{ii} sprintf(', $p=%g, q=%g$', param(1), param(2))];
	case 'stevenson94dpr'
		param = [2 1.2];
		leg{ii} = [leg{ii} sprintf(', $p=%g, q=%g$', param(1), param(2))];
	case {'table0', 'table1'}
		dz = delta / 20;
		tmp = potential_fun('qgg2', delta, 1.2);
%		param = {dz, tmp.dpot([1:1e4] * dz)};
		param = [dz, tmp.dpot([1:1e4] * dz)];
%		tmp = sprintf(' K = %d dz = %g', numel(param{2}), param{1});
		tmp = sprintf(', $K = %d$, $\\Delta t = %g$', numel(param)-1, param(1));
		leg{ii} = [leg{ii} tmp];
	otherwise
		param = [];
	end

	pot = potential_fun(ptype, delta, param);
	pp(:,ii) = pot.potk(zz);
	pw(:,ii) = pot.wpot(zz);
	pd(:,ii) = pot.dpot(zz);

	% replace 0 with 1 to make (slow) figure showing fzero-based shrinkage
	if 0 || ~streq(func2str(pot.meth.shrink), 'ir_potential_fun_shrink')
		tmp = pot.shrink(bb, reg);
		if any(tmp ~= 0)
			ps(:,end+1) = pot.shrink(bb, reg);
			lshrink{end+1} = leg{ii};
		end
	end

	if 0 % test vs old
		try
			opot = potential_func(ptype, delta, param);
			opp = opot.potk(opot, zz);
			opw = opot.wpot(opot, zz);
			opd = opot.dpot(opot, zz);
		catch
			printm('%s no old version', ptype)
		%	keyboard % problem with private scope of ir_li98cfs_wpot
			continue
		end
		if ~isequal(opp, pp(:,ii)), 'p bug', ptype, keyboard, end
		if ~isequal(opw, pw(:,ii)), 'w bug', ptype, keyboard, end
		if ~isequal(opd, pd(:,ii)), 'd bug', ptype, keyboard, end
	end
end

if im
%	set(0,'DefaultAxesLineStyleOrder', '-|:')
	clf, pl = @(i) subplot(410 + i);
	pl(1), plot(zz, pp), title 'potk'
	axis tight, axisy([-0.0 2.5] * delta^2)
	ir_legend(leg, 'location', 'north')
	pl(2), plot(zz, pw), title 'wpot'
	axis tight, axisy(0, 1.1), ytick([0 1])
	pl(3), plot(zz, pd), title 'dpot'
	axis tight, axisy([-1 1] * 1.3 * delta)
	xlabelf '$z$'
%	ir_savefig eps_c fname
	% check derivatives
	pl(4)
	plot(zz, pd)
	title 'dpot check'
	hold on
	d = diffc(pp) / (zz(2)-zz(1));
	plot(zz(1:end-1), d(1:end-1,:), '--', ...
		zz(1:end-1), d(1:end-1,:)-pd(1:end-1,:), ':')
	hold off
	axis tight, axisy([-1 1] * 1.3 * delta)
end

if 1 && im
	if ~isempty(lshrink)
		prompt
		clf
		plot(bb, ps, '-', bb, bb, '-')
		axis equal, axis square
		axis([-1 1 -1 1]*400)
		xtick([-1 0 1] * 400)
		ytick([-1 0 1] * 400)
		ir_legend(lshrink, 'location', 'southeast')
		xlabelf '$c$'
%		ylabelf 'xhat(z)'
		ylabelf '$\hat{z}(c)$'
		grid
%		ir_savefig cw fig_reg_pot_table0_qgg2_shrink

		if 0 % figure for book
			plot(bb, ps(:,[3 2]) - repmat(ps(:,1), [1 2]), 'o')
			plot(	bb, ps(:,3) - ps(:,1), 'bo', ...
				bb, ps(:,2) - ps(:,1), 'gx')
			axis([0 500 -0.05 0.5]), ytick([0 0.05 0.1 0.5])
			xtick([0 500])
			ir_fontsize label 18
			ir_fontsize text 18
			xlabelf 'c', ylabelf 'shrinkage error for QGG2'
			ir_legend(lshrink([3 2]), 1)
%			ir_savefig cw fig_reg_pot_table01_qgg2_shrinker
			keyboard
		end
	end
end

% check 'scale' option
pot1 = potential_fun('lange3', 10, []);
pot2 = potential_fun('lange3', 10, [], 'scale', 2);
jf_equal(2 * pot1.potk(zz), pot2.potk(zz))
jf_equal(2 * pot1.dpot(zz), pot2.dpot(zz))
jf_equal(2 * pot1.wpot(zz), pot2.wpot(zz))

end % ir_potential_fun_test

% ir_potential_fun_test_lpp()
function ir_potential_fun_test_lpp
params = [0 0.01 0.1 1/2 3/4 1 4/3];
params = [0 1/2 1];
reg = 4*2^4; zmax = 2 * reg;
zz = zmax * linspace(-1, 1, 2001)';
bb = 2 * reg * linspace(-1, 1, 201);
%a_1_2 = 3 * (reg/4)^(2/3) % saddle point for p=1/2
a_1_2 = 3/2 * reg^(2/3); % breakpoint for p=1/2
if 0
	a = a_1_2;
	pot = potential_fun('lpp', nan, 1/2);
	z0 = ir_potential_fun_shrink(pot, a+1e-9, reg)
	cost = @(z) 1/2 * (z - a).^2 + reg * pot.potk(z);
	t0 = sqrt(z0)
	t0^3 - 2*a*t0 + 2*reg % should be 0
	t0^3 - a*t0 + reg/2 % should be 0
	3/2*reg - a * t0 % should be 0
	(3^(3/2) * reg)/(4*a^(3/2)) % should be 1/sqrt(2)
	cost(0)
	cost(z0)
%	syms x; solve(3 * acos(x) - acos(-x), x) % 1/sqrt(2)
return
end
tmp = [sqrt(2*reg) reg a_1_2];
tmp = outer_sum(tmp, [-1 +1]*1e-5);
bb = sort([bb col(tmp)' -col(tmp)'])'; % trick
ps = [];
lshrink = {};
if 0
	pot = potential_fun('lpp', nan, 1/2);
	potk = @(z) pot.potk(z);

%	y = sqrt(2*reg); % l0
%	y = reg; % l1
	y = a_1_2 + 1e-9; % p=1/2
	if 0
		ir_potential_fun_shrink(pot, y, reg)
		ir_potential_fun_shrink(pot, y, reg)
	return
	end

	if 0
		z1 = pot.shrink(y, reg)
		d1 = z1 - y + reg * pot.dpot(z1) % 1st derivative
		d2 = 1 + reg * (1/2)*(-1/2) * z1^(-3/2) % 2nd derivative > 0?

		dfit = 1/2 * (y - zz).^2;
		cost = dfit + reg * potk(zz);
		plot(zz, dfit, '--', zz, reg*potk(zz), '-', zz, cost, '-')
		tmp = minmax(zz(cost <= 3*min(cost)));
		tmp(1) = min(tmp(1), -10);
		tmp(2) = max(tmp(2), reg/3);
		xlim(tmp), ylim([0.0 2*1.2] * min(cost)), grid
	%	ytick([0 round(min(cost))])
		ytick([0 min(cost)])
		xtick([0 z1 a_1_2]), xlabelf '$z$', ylabelf '$\Psi(z)$'
	return
	end

	zs = ir_potential_fun_shrink(pot, bb, reg);
	zp = pot.shrink(bb, reg);

	plot(bb, zs, '-o', bb, zp, '.-', bb, bb, ':')
	legend({'fmin', 'trig'}, 'location', 'northwest')
	xlabelf '$b$'
	ylabelf '$\hat{z}(b)$'
	axis equal, axis square
	axis([-1 1 -1 1]*1.1*reg)
	tmp = [sqrt(2*reg) a_1_2 reg]; % l0 and l1 breakpoints
	xtick([-tmp 0 tmp])
	ytick([-1 0 1] * reg)
	grid
%	ir_savefig cw fig_?
return
end

for ip=1:length(params)
	pot = potential_fun('lpp', 0, params(ip));

	% replace 0 with 1 to make (slow) figure showing fzero-based shrinkage
	if 0 || ~streq(func2str(pot.meth.shrink), 'ir_potential_fun_shrink')
		tmp = pot.shrink(bb, reg);
		if any(tmp ~= 0)
			ps(:,end+1) = pot.shrink(bb, reg);
			lshrink{end+1} = sprintf('p=%g', params(ip));
		end
	end

	clf
	plot(bb, ps, '-', bb, bb, ':')
	legend(lshrink{:}, 'location', 'southeast')
	xlabelf '$b$'
	ylabelf '$\hat{z}(b)$'
	axis equal, axis square
	axis([-1 1 -1 1]*1.5*reg)
	tmp = [sqrt(2*reg) a_1_2 reg]; % l0 and l1 breakpoints
	xtick([-tmp 0 tmp])
	ytick([-1 0 1] * reg)
	grid
%	ir_savefig cw fig_?
end

end % ir_potential_fun_test_lpp()


% [INPUT]
% chi = A float n-by-n-by-t matrix [0,1] representing the time-varying Chi coefficients.
%
% [OUTPUT]
% achi = A row vector of floats [0,1] of length t representing the Average Chi.
% adr = A row vector of floats [0,1] of length t representing the Asymptotic Dependence Rate.

function [achi,adr] = chi_metrics(varargin)

    persistent ip;

    if (isempty(ip))
        ip = inputParser();
        ip.addRequired('chi',@(x)validateattributes(x,{'double'},{'real' '3d' 'nonempty'}));
    end

    ip.parse(varargin{:});

    ipr = ip.Results;
    chi = validate_input(ipr.chi);

    nargoutchk(2,2);

    [achi,adr] = chi_metrics_internal(chi);

end

function [achi,adr] = chi_metrics_internal(chi)

    up = isempty(getCurrentTask());

    [n1,n2,t] = size(chi);
    n = min(n1,n2);

    achi = zeros(t,1);
    adr = zeros(t,1);

    if (up)
        parfor k = 1:t
            chi_k = chi(:,:,k);

            adr_num = 0;
            chi_sum = 0;
            den = 0;

            for i = 1:n
                for j = 1:n
                    if (i == j)
                        continue;
                    end

                    chi_kij = chi_k(i,j);

                    if (isnan(chi_kij))
                        continue;
                    end

                    if (chi_kij > 0)
                        adr_num = adr_num + 1;
                    end

                    chi_sum = chi_sum + chi_kij;
                    den = den + 1;
                end
            end

            achi(k) = chi_sum / den;
            adr(k) = adr_num / den;
        end
    else
        for k = 1:t
            chi_k = chi(:,:,k);

            adr_num = 0;
            chi_sum = 0;
            den = 0;

            for i = 1:n
                for j = 1:n
                    if (i == j)
                        continue;
                    end

                    chi_kij = chi_k(i,j);

                    if (isnan(chi_kij))
                        continue;
                    end

                    if (chi_kij > 0)
                        adr_num = adr_num + 1;
                    end

                    chi_sum = chi_sum + chi_kij;
                    den = den + 1;
                end
            end

            achi(k) = chi_sum / den;
            adr(k) = adr_num / den;
        end
    end

end

function chi = validate_input(chi)

    [n1,n2,t] = size(chi);

    if ((n1 ~= n2) || (min(n1,n2) < 2))
        error('The value of ''chi'' is invalid. Expected input to be a square 3d matrix with a minimum size of 2x2xt.');
    end

    if (t < 5)
        error('The value of ''chi'' is invalid. Expected input to be a square 3d matrix with at least elements on the third dimension.');
    end

    chiv = chi(:);
    chiv(isnan(chiv)) = [];

    if (any((chiv < 0) | (chiv > 1)))
        error('The value of ''chi'' is invalid. Expected input to contain non-NaN values greater than or equal to 0 and less than or equal to 1.');
    end

end


function [estSig, cost] = ILRMAISS(mixSig, nSrc, sampFreq, nBases, fftSize, shiftSize, windowType, nIter, ilrmaType, refMic, applyNormalize, applyWhitening, drawConv)
% Blind source separation with independent low-rank matrix analysis (ILRMA) based on iterative source steering update rule
%
% Coded by D. Kitamura (d-kitamura@ieee.org)
%
% Copyright 2021 Daichi Kitamura
%
% These programs are distributed only for academic research at
% universities and research institutions.
% It is not allowed to use or modify these programs for commercial or
% industrial purpose without our permission.
% When you use or modify these programs and write research articles,
% cite the following references:
%
% # Original paper (The algorithm was called "Rank-1 MNMF" in this paper)
% D. Kitamura, N. Ono, H. Sawada, H. Kameoka, H. Saruwatari, "Determined
% blind source separation unifying independent vector analysis and
% nonnegative matrix factorization," IEEE/ACM Trans. ASLP, vol. 24,
% no. 9, pp. 1626-1641, September 2016.
%
% # Book chapter (The algorithm was renamed as "ILRMA")
% D. Kitamura, N. Ono, H. Sawada, H. Kameoka, H. Saruwatari, "Determined
% blind source separation with independent low-rank matrix analysis,"
% Audio Source Separation. Signals and Communication Technology.,
% S. Makino, Ed. Springer, Cham, pp. 125-155, March 2018.
%
% # Original AuxIVA-ISS paper
% S. Robin and N. Ono,
% "Fast and stable blind source separation with rank-1 updates,"
% Proc. ICASSP, pp.236–240, 2020.
%
% See also:
% http://d-kitamura.net
% http://d-kitamura.net/demo-ILRMA_en.html
%
% [syntax]
%   [estSig,cost] = ILRMAISS(mixSig,nSrc,sampFreq,nBases,fftSize,shiftSize,windowType,nIter,ilrmaType,refMic,applyNormalize,applyWhitening,drawConv)
%
% [inputs]
%         mixSig: observed mixture (sigLen x nCh)
%           nSrc: number of sources in the mixture (scalar)
%       sampFreq: sampling frequency [Hz] of mixSig (scalar)
%         nBases: number of bases in NMF model (scalar, # of bases for "each" source when ilrmaType=1, and # of bases for "all" sources when ilrmaType=2, default: 4)
%        fftSize: window length [points] in STFT (scalar, default: next higher power of 2 that exceeds 0.256*sampFreq)
%      shiftSize: shift length [points] in STFT (scalar, default: fftSize/2)
%     windowType: window function used in STFT (name of window function, default: 'hamming')
%          nIter: number of iterations in the parameter update in ILRMA (scalar, default: 100)
%      ilrmaType: without or with partitioning function (1: ILRMA without partitioning function (ILRMA type1), 2: ILRMA with partitioning function (ILRMA type2), default: 1)
%         refMic: reference microphone for applying back projection (default: 1)
% applyNormalize: normalize parameters in each iteration to avoid numerical divergence (0: do not apply, 1: average-power-based normalization, 2: back projection (only for ILRMA type1), normalization may collapse monotonic decrease of the cost function, default: 0)
% applyWhitening: apply whitening to the observed multichannel spectrograms or not (true or false, default: true)
%       drawConv: plot cost function values in each iteration or not (true or false, default: false)
%
% [outputs]
%         estSig: estimated signals (sigLen x nCh x nSrc)
%           cost: convergence behavior of cost function in ILRMA (nIter+1 x 1)
%

% Arguments check and set default values
arguments
    mixSig (:,:) double
    nSrc (1,1) double {mustBeInteger(nSrc)}
    sampFreq (1,1) double
    nBases (1,1) double {mustBeInteger(nBases)} = 4
    fftSize (1,1) double {mustBeInteger(fftSize)} = 2^nextpow2(0.256*sampFreq)
    shiftSize (1,1) double {mustBeInteger(shiftSize)} = fftSize/2
    windowType char {mustBeMember(windowType,{'hamming','hann','rectangular','blackman','sine'})} = 'hamming'
    nIter (1,1) double {mustBeInteger(nIter)} = 100
    ilrmaType (1,1) double {mustBeInteger(ilrmaType)} = 1
    refMic (1,1) double {mustBeInteger(refMic)} = 1
    applyNormalize (1,1) double {mustBeInteger(applyNormalize)} = 0
    applyWhitening (1,1) logical = true
    drawConv (1,1) logical = false
end

% Error check
[sigLen, nCh] = size(mixSig); % sigLen: signal length, nCh: number of channels
if sigLen < nCh; error("The size of mixSig might be wrong.\n"); end
if nCh < nSrc || nSrc < 2; error("The number of channels must be equal to or grater than the number of sources in the mixture.\n"); end
if sampFreq <= 0; error("The sampling frequency (sampFreq) must be a positive value.\n"); end
if nBases < 1; error("The number of bases (nBases) must be a positive integer value.\n"); end
if fftSize < 1; error("The FFT length in STFT (fftSize) must be a positive integer value.\n"); end
if shiftSize < 1; error("The shift length in STFT (shiftSize) must be a positive integer value.\n"); end
if nIter < 1; error("The number of iterations (nIter) must be a positive integer value.\n"); end
if ilrmaType ~= 1 && ilrmaType ~= 2; error("The ILRMA type (ilrmaType) must be set to 1 or 2.\n"); end
if refMic < 1 || refMic > nCh; error("The reference microphone must be an integer between 1 and nCh.\n"); end
if applyNormalize ~= 0 && applyNormalize ~= 1 && applyNormalize ~= 2; error("The normalization type (applyNormalize) must be set to 0, 1, or 2.\n"); end
if applyNormalize == 2 && ilrmaType == 2; error("The back-projection-based normalization only supports ILRMA type 1.\n"); end

% Apply multichannel short-time Fourier transform (STFT)
[mixSpecgram, windowInStft] = STFT(mixSig, fftSize, shiftSize, windowType);

% Apply whitening (decorrelate X so that the correlation matrix becomes an identity matrix) based on principal component analysis
if applyWhitening
    inputMixSpecgram = whitening(mixSpecgram, nSrc); % apply whitening, where dimension is reduced from nCh to nSrc when nSrc < nCh
else
    inputMixSpecgram = mixSpecgram(:,:,1:nSrc); % when nSrc < nCh, only mixSpecgram(:,:,1:nSrc) is input to ILRMA so that the number of microphones equals to the number of sources (determined condition)
end

% Apply ILRMA
if ilrmaType == 1
    [estSpecgram, cost] = local_ILRMA1ISS(inputMixSpecgram, nIter, nBases, applyNormalize, drawConv, mixSpecgram(:,:,refMic));
else
    [estSpecgram, cost] = local_ILRMA2ISS(inputMixSpecgram, nIter, nBases, applyNormalize, drawConv);
end

% Apply back projection (fix the scale ambiguity using the reference microphone channel)
scaleFixedSepSpecgram = backProjection(estSpecgram, mixSpecgram(:,:,refMic)); % scale-fixed estimated signal

% Inverse STFT for each source
estSig = ISTFT(scaleFixedSepSpecgram, shiftSize, windowInStft, sigLen);
end

%% Local function for ILRMA type1 (without pertitioning function) based on ISS
function [Y, cost] = local_ILRMA1ISS(X, nIter, L, applyNormalize, drawConv, refMixSpecgram)
% [inputs]
%              X: observed multichannel spectrograms (I x J x M)
%          nIter: number of iterations of the parameter updates
%              L: number of bases in NMF model for each source
% applyNormalize: normalize parameters in each iteration to avoid numerical divergence (0: do not apply, 1: average-power-based normalization, 2: back projection, normalization may collapse monotonic decrease of the cost function)
%       drawConv: plot cost function values in each iteration or not (true or false)
% refMixSpecgram: observed reference spectrogram before apply whitening (I x J)
%
% [outputs]
%              Y: estimated spectrograms of sources (I x J x N)
%           cost: convergence behavior of cost function in ILRMA (nIter+1 x 1)
%
% [scalars]
%              I: number of frequency bins,
%              J: number of time frames
%              M: number of channels (microphones)
%              N: number of sources (equals to M)
%              L: number of bases in NMF model for each source
%
% [matrices]
%              X: observed multichannel spectrograms (I x J x M)
%             pX: permuted observed multichannel spectrograms (M x J x I)
%             pY: permuted separated multisource spectrograms (N x J x I)
%              W: frequency-wise demixing matrix (N x M x I)
%              Y: estimated multisource spectrograms (I x J x N)
%              P: estimated multisource power spectrograms (I x J x N)
%              T: sourcewise basis matrix in NMF (I x L x N)
%              V: sourcewise activation matrix in NMF (L x J x N)
%              R: sourcewise low-rank model spectrogram constructed by T and V (I x J x N)
%              E: identity matrix (N x N)
%              U: model-spectrogram-weighted sample covariance matrix of the mixture (M x M)
%

% Initialization
[I,J,M] = size(X); % I:frequency bins, J: time frames, M: channels
pX = permute(X, [3,2,1]); % permuted X whose dimensions are M x J x I
N = M; % N: number of sources, which equals to M in ILRMA
W = zeros(N,M,I); % frequency-wise demixing matrix
Y = zeros(I,J,N); % estimated spectrograms of sources (Y(i,:,n) =  W(n,:,i)*pX(:,:,i))
for i = 1:I
    W(:,:,i) = eye(N); % initial demixing matrices are set to identity matrices
    Y(i,:,:) = (W(:,:,i)*pX(:,:,i)).'; % initial estimated spectrograms
end
P = max(abs(Y).^2, eps); % power spectrogram of Y
T = max(rand( I, L, N ), eps); % sourcewise basis matrix in NMF
V = max(rand( L, J, N ), eps); % sourcewise activation matrix in NMF
R = zeros(I,J,N); % sourcewise low-rank model spectrogram constructed by T and V (R(:,:,n) = T(:,:,n)*V(:,:,n))
for n = 1:N
    R(:,:,n) = T(:,:,n)*V(:,:,n); % initial source model defined by T and V
end
v = zeros(N,1); % v vector for iterative source steering
cost = zeros(nIter+1, 1);

% Calculate initial cost function value
if drawConv
    cost(1,1) = local_calcCostFunction( P, R, W, I, J );
end

% Optimize parameters in ILRMA (W, T, and V)
fprintf('Iteration:    ');
for iIter = 1:nIter
    fprintf('\b\b\b\b%4d', iIter);

    %%%%% Update parameters %%%%%
    for n = 1:N
        %%%%% Update rule of T %%%%%
        T(:,:,n) = T(:,:,n) .* sqrt((P(:,:,n)./(R(:,:,n).^2))*V(:,:,n).' ./ ( (1./R(:,:,n))*V(:,:,n).' ));
        T(:,:,n) = max(T(:,:,n), eps);
        R(:,:,n) = T(:,:,n)*V(:,:,n);
        %%%%% Update rule of V %%%%%
        V(:,:,n) = V(:,:,n) .* sqrt(T(:,:,n).'*(P(:,:,n)./(R(:,:,n).^2)) ./ ( T(:,:,n).'*(1./R(:,:,n)) ));
        V(:,:,n) = max(V(:,:,n), eps);
        R(:,:,n) = T(:,:,n)*V(:,:,n);
        %%%%% Update rule of Y %%%%%
        YY = Y .* conj(Y(:,:,n)); % I x J x N, using implicit expansion (IxJxN .* IxJx1)
        for i = 1:I
            for nn = 1:N % calculate v vector BEGIN
                d = sum((1./(R(i,:,nn))).*real(YY(i,:,n)), 2) / J; % scalar
                if nn ~= n
                    u = sum((1./(R(i,:,nn))).*YY(i,:,nn), 2) / J; % scalar
                    v(nn,1) = u / d;
                else
                    v(nn,1) = 1 - 1/sqrt(d);
                end
            end % calculate v vector END
            Y(i,:,:) = Y(i,:,:) - permute(v.*Y(i,:,n), [3,2,1]); % update Y, usnig implicit expansion (Nx1 .* 1xJ = NxJ -> 1xJxN)
        end
    end
    P = max(abs(Y).^2, eps); % power spectrogram of Y

    %%%%% Normalization %%%%%
    if applyNormalize == 1 % average-power-based normalization
        lambda = sqrt(sum(sum(P,1),2)/(I*J)); % 1 x 1 x N
        Y = Y./lambda; % I x J x N (use implicit expansion)
        lambdaPow = lambda.^2; % 1 x 1 x N
        P = P./lambdaPow; % I x J x N (use implicit expansion)
        R = R./lambdaPow; % I x J x N (use implicit expansion)
        T = T./lambdaPow; % I x L x N (use implicit expansion)
    elseif applyNormalize == 2 % back projection
        lambda = local_backProjection(Y, refMixSpecgram, I, N); % N x 1 x I
        pLambda = permute(lambda, [3,2,1]); % I x 1 x N
        Y = Y.*pLambda; % I x J x N (use implicit expansion)
        lambdaPow = abs(pLambda).^2; % I x 1 x N
        P = P.*lambdaPow; % I x J x N (use implicit expansion)
        R = R.*lambdaPow; % I x J x N (use implicit expansion)
        T = T.*lambdaPow; % I x L x N (use implicit expansion)
    end

    %%%%% Calculate cost function value %%%%%
    if drawConv
        pY = permute(Y, [3,2,1]); % IxJxN -> NxJxI
        for i = 1:I
            W(:,:,i) = pY(:,:,i)*pX(:,:,i)' / (pX(:,:,i)*pX(:,:,i)'); % derived by "Y(:,:,i) = W(:,:,i)*X(:,:,i)"
        end
        cost(iIter+1,1) = local_calcCostFunction( P, R, W, I, J );
    end
end

% Draw convergence behavior
if drawConv
    figure; plot((0:nIter), cost);
    set(gca, 'FontName', 'Times', 'FontSize', 16);
    xlabel('Number of iterations', 'FontName', 'Arial', 'FontSize', 16);
    ylabel('Value of cost function', 'FontName', 'Arial', 'FontSize', 16);
end
fprintf(' ILRMA1-ISS done.\n');
end

%% Local function for ILRMA type2 (with pertitioning function) based on ISS
function [Y, cost] = local_ILRMA2ISS(X, nIter, K, applyNormalize, drawConv)
% [inputs]
%              X: observed multichannel spectrograms (I x J x M)
%          nIter: number of iterations of the parameter updates
%              K: number of bases in NMF model shared for all the source
% applyNormalize: normalize parameters in each iteration to avoid numerical divergence (0: do not apply, 1: average-power-based normalization, normalization may collapse monotonic decrease of the cost function)
%       drawConv: plot cost function values in each iteration or not (true or false)
%
% [outputs]
%              Y: estimated spectrograms of sources (I x J x N)
%           cost: convergence behavior of cost function in ILRMA (nIter+1 x 1)
%
% [scalars]
%              I: number of frequency bins,
%              J: number of time frames
%              M: number of channels (microphones)
%              N: number of sources (equals to M)
%              K: number of bases in NMF model shared for all the source
%
% [matrices]
%              X: observed multichannel spectrograms (I x J x M)
%             pX: permuted observed multichannel spectrograms (M x J x I)
%             pY: permuted separated multisource spectrograms (N x J x I)
%              W: frequency-wise demixing matrix (N x M x I)
%              Y: estimated multisource spectrograms (I x J x N)
%              P: estimated multisource power spectrograms (I x J x N)
%              T: basis matrix in NMF (I x K)
%              V: activation matrix in NMF (K x J)
%              Z: partitioning function matrix in NMF that clusters K bases into N sources (N x K)
%              R: sourcewise low-rank model spectrogram constructed by T and V (I x J x N)
%              E: identity matrix (N x N)
%              U: model-spectrogram-weighted sample covariance matrix of the mixture (M x M)
%

% Initialization
[I,J,M] = size(X); % I:frequency bins, J: time frames, M: channels
pX = permute(X, [3,2,1]); % permuted X whose dimensions are M x J x I
N = M; % N: number of sources, which equals to M in ILRMA
W = zeros(N,M,I); % frequency-wise demixing matrix
Y = zeros(I,J,N); % estimated spectrograms of sources (Y(i,:,n) =  W(n,:,i)*pX(:,:,i))
for i = 1:I
    W(:,:,i) = eye(N); % initial demixing matrices are set to identity matrices
    Y(i,:,:) = (W(:,:,i)*pX(:,:,i)).'; % initial estimated spectrograms are set to the observed mixture spectrograms
end
P = max(abs(Y).^2, eps); % power spectrogram of Y
T = max(rand( I, K ), eps); % basis matrix in NMF shared for all the sources
V = max(rand( K, J ), eps); % activation matrix in NMF shared for all the sources
Z = max(rand( N, K ), eps); % partitioning function matrix in NMF that clusters K bases into N sources
Z  = Z./sum(Z,1); % ensure sum_n z_{nk} = 1 (use implicit expansion)
tmpT = zeros(size(T)); % temporal variable used in the update rule of T
tmpV = zeros(size(V)); % temporal variable used in the update rule of V
tmpZ = zeros(size(Z)); % temporal variable used in the update rule of Z
R = zeros(I,J,N); % sourcewise low-rank model spectrogram constructed by T, V, and Z (R(:,:,n) = T(:,:,n)*V(:,:,n))
for n = 1:N
    R(:,:,n) = (Z(n,:).*T)*V; % initial source model defined by T, V, and Z (use implicit expansion)
end
v = zeros(N,1); % v vector for iterative source steering
cost = zeros(nIter+1, 1);

% Calculate initial cost function value
if drawConv
    cost(1,1) = local_calcCostFunction( P, R, W, I, J );
end

% Optimize parameters in ILRMA (W, T, and V)
fprintf('Iteration:    ');
for iIter = 1:nIter
    fprintf('\b\b\b\b%4d', iIter);

    %%%%% Update parameters %%%%%
    %%%%% Update rule of Z %%%%%
    for n = 1:N
        tmpZ(n,:) = (( sum((T.'*(P(:,:,n)./(R(:,:,n).^2))).*V, 2) )./( sum((T.'*(1./R(:,:,n))).*V, 2) )).';
    end
    Z = Z .* sqrt(tmpZ);
    Z  = max(Z./sum(Z,1), eps); % ensure sum_n z_{nk} = 1 (use implicit expansion)
    for n = 1:N
        R(:,:,n) = (Z(n,:).*T)*V; % initial source model defined by NMF (use implicit expansion)
    end
    %%%%% Update rule of T %%%%%
    for i = 1:I
        Pi = squeeze(P(i,:,:)); % J x N
        Ri = squeeze(R(i,:,:)); % J x N
        tmpT(i,:) = (( sum((V*(Pi./(Ri.^2))).*(Z.'), 2) )./( sum((V*(1./Ri)).*(Z.'), 2) )).';
    end
    T = max(T.*sqrt(tmpT), eps);
    for n = 1:N
        R(:,:,n) = (Z(n,:).*T)*V; % initial source model defined by NMF (use implicit expansion)
    end
    %%%%% Update rule of V %%%%%
    for j = 1:J
        Pj = squeeze(P(:,j,:)); % I x N
        Rj = squeeze(R(:,j,:)); % I x N
        tmpV(:,j) = ( sum((T.'*(Pj./(Rj.^2))).*(Z.'), 2) )./( sum((T.'*(1./Rj)).*(Z.'), 2) );
    end
    V = max(V.*sqrt(tmpV), eps);
    for n = 1:N
        R(:,:,n) = (Z(n,:).*T)*V; % initial source model defined by NMF (use implicit expansion)
    end
    %%%%% Update rule of Y %%%%%
    for n=1:N
        YY = Y .* conj(Y(:,:,n)); % I x J x N, using implicit expansion (IxJxN .* IxJx1)
        for i = 1:I
            for nn = 1:N % calculate v vector BEGIN
                d = sum((1./(R(i,:,nn))).*real(YY(i,:,n)), 2) / J; % scalar
                if nn ~= n
                    u = sum((1./(R(i,:,nn))).*YY(i,:,nn), 2) / J; % scalar
                    v(nn,1) = u / d;
                else
                    v(nn,1) = 1 - 1/sqrt(d);
                end
            end % calculate v vector END
            Y(i,:,:) = Y(i,:,:) - permute(v.*Y(i,:,n), [3,2,1]); % update Y, usnig implicit expansion (Nx1 .* 1xJ = NxJ -> 1xJxN)
        end
    end
    P = max(abs(Y).^2, eps); % power spectrogram of Y

    %%%%% Normalization %%%%%
    if applyNormalize == 1
        lambda = sqrt(sum(sum(P,1),2)/(I*J)); % 1 x 1 x N
        Y = Y./lambda; % I x J x N (use implicit expansion)
        P = P./lambda.^2; % I x J x N (use implicit expansion)
        R = R./lambda.^2; % I x J x N (use implicit expansion)
        Zlambda = Z./(squeeze(lambda).^2); % N x K
        ZlambdaSum = sum(Zlambda,1); % 1 x K
        T = T.*ZlambdaSum; % I x K (use implicit expansion)
        Z = Zlambda./ZlambdaSum; % N x K (use implicit expansion)
    end

    %%%%% Calculate cost function value %%%%%
    if drawConv
        pY = permute(Y, [3,2,1]); % IxJxN -> NxJxI
        for i = 1:I
            W(:,:,i) = pY(:,:,i)*pX(:,:,i)' / (pX(:,:,i)*pX(:,:,i)'); % derived by "Y(:,:,i) = W(:,:,i)*X(:,:,i)"
        end
        cost(iIter+1,1) = local_calcCostFunction( P, R, W, I, J );
    end
end

% Draw convergence behavior
if drawConv
    figure; plot((0:nIter), cost);
    set(gca, 'FontName', 'Times', 'FontSize', 16);
    xlabel('Number of iterations', 'FontName', 'Arial', 'FontSize', 16);
    ylabel('Value of cost function', 'FontName', 'Arial', 'FontSize', 16);
end
fprintf(' ILRMA2-ISS done.\n');
end

%% Local function for calculating cost function value in ILRMA
function [ cost ] = local_calcCostFunction(P, R, W, I, J)
logDetAbsW = zeros(I,1);
for i = 1:I
    logDetAbsW(i,1) = log(max(abs(det(W(:,:,i))), eps));
end
cost = sum(sum(sum(P./R+log(R),3),2),1) - 2*J*sum(logDetAbsW, 1);
end

%% Local function for applying back projection and returns frequency-wise coefficients
function [ D ] = local_backProjection(Y, X, I, N)
D = zeros(I,N);
for i = 1:I
    Yi = squeeze(Y(i,:,:)).'; % N x J
    D(i,:) = X(i,:,1)*Yi'/(Yi*Yi'); % 1 x N
end
D(isnan(D) | isinf(D)) = 0; % replace NaN and Inf to 0
D = permute(D, [2,3,1]); % N x 1 x I
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% EOF %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


% SCRIPT TEST FOR THE 3 DOF spherical manipulator

% Copyright (C) 2012, by Arturo Gil Aparicio
%
% This file is part of ARTE (A Robotics Toolbox for Education).
%
% ARTE is free software: you can redistribute it and/or modify
% it under the terms of the GNU Lesser General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% ARTE 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 Lesser General Public License for more details.
%
% You should have received a copy of the GNU Leser General Public License
% along with ARTE.  If not, see <http://www.gnu.org/licenses/>.

close all

%general call to inverse dynamic function
%TAU = inversedynamic(param, Q, QD, QDD, GRAV, FEXT)

%three different poses
q1 = [0, pi/2, pi/2];
q2 = [0, pi/2, 0];

%load 3dof spherical robot
%the parameters are loaded from robots/example/3dofspherical/parameters.m
spherical=load_robot('example', '3dofspherical');

spherical.graphical.draw_transparent=1;

fprintf('\nTorques at each joint given position = [0 0 0], and velocity=[0 0 0] and standard gravity acting on Z0')
%Please note that the forces and moments are specified with respect to the
%X3 Y3 Z3 reference system.
%In this mechanism, forces acting on Fx3 or Fz3 do not account for tau, nor
%moments acting on My3 or Mx3
tau = inversedynamic(spherical, q1, [0 0 0], [0 0 0], [0 0 9.81]', [0 0 0 0 0 0]')

%draw the robot at this position
drawrobot3d(spherical,q1)
disp('press any key to continue')
pause

%Now, add a force on Fy, and observe the results on tau
tau = inversedynamic(spherical, q1, [0 0 0], [0 0 0], [0 0 9.81]', [0 1 0 0 0 0]')
%draw the robot at this position
drawrobot3d(spherical,q1)
disp('press any key to continue')
pause

%torques necessary to instantaneously bring the arm to the specified state
fprintf('\nTorques at each joint given position = [0 pi/2 0], and velocity=[1 1 1] and acceleration = [1 1 1] and gravity acting on Z0')
tau = inversedynamic(spherical, q2, [1 1 1], [1 1 1], [0 0 9.81]', [0 0 0 0 0 0]')

%draw the robot
drawrobot3d(spherical,q2)
disp('press any key to continue')
pause

%In this case, all torques are caused by gravity acting on the Center Of
%Mass of each link. Note that we have changed the direction of the gravity
%vector g
fprintf('\nTorques at each joint given position = [0 0 0], and velocity=[0 0 0] and acceleration = [0 0 0] and gravity acting on Y0')
tau = inversedynamic(spherical, q1, [0 0 0], [0 0 0], [0 9.81 0]', [0 0 0 0 0 0]')

%draw the robot
drawrobot3d(spherical,q1)


function g = power(f, b)
%.^   CHEBTECH power.
%   F.^G returns a CHEBTECH F to the scalar power G, a scalar F to the CHEBTECH
%   power G, or a CHEBTECH F to the CHEBTECH power G. F and or G may be complex.
%
%   This function assumes that the curve traced out by F in the complex plane
%   both (1) does not come too close to zero and (2) does not cross over the
%   branch cut in POWER along the negative real axis.  That is, F should not
%   vanish at any point of [-1, 1], and the imaginary part of F should not
%   vanish at any point of (-1, 1) where the real part of F is negative.  If any
%   of these assumptions are violated, garbage may be returned with no warning.
%
%   H = POWER(F, G) is called for the syntax 'F .^ G'.
%
% See also SQRT.

% Copyright 2017 by The University of Oxford and The Chebfun Developers.
% See http://www.chebfun.org/ for Chebfun information.

% If F is complex and the imaginary part of F vanishes exactly at the domain
% boundary, tiny rounding errors in evaluating the imaginary part at the
% boundary points can cause it to appear to change sign there.  If this happens,
% and if the real part of F is negative there, the endpoint function values will
% cross over the branch cut in POWER, creating a jump discontinuity that should
% not be there.  Enabling extrapolation resolves this issue by avoiding these
% function evaluations:

pref = f.techPref();
if ( ~isreal(f) )
    pref.extrapolate = 1;
end

% Simply call the compose function:
if ( isa(f, 'chebtech') && isa(b, 'chebtech') )
    % Both F and G are CHEBTECHs:
	g = compose(f, @power, b, [], [], pref);
elseif ( isa(f, 'chebtech') )
    % F is CHEBTECH and G is constant:
	g = compose(f, @(f) power(f, b), [], [], pref);
else
    % F is constant and G is CHEBTECH:
	g = compose(b, @(b) power(f, b), [], [], pref);
end

end


function [Isub,domSub,femSub] = femSubdivide(dom,fem,Nsub,fig)
%+========================================================================+
%|                                                                        |
%|              OPENFEM - LIBRARY FOR FINITE ELEMENT METHOD               |
%|           openFem is part of the GYPSILAB toolbox for Matlab           |
%|                                                                        |
%| COPYRIGHT : Matthieu Aussal & Francois Alouges (c) 2017-2018.          |
%| PROPERTY  : Centre de Mathematiques Appliquees, Ecole polytechnique,   |
%| route de Saclay, 91128 Palaiseau, France. All rights reserved.         |
%| LICENCE   : This program is free software, distributed in the hope that|
%| it will be useful, but WITHOUT ANY WARRANTY. Natively, you can use,    |
%| redistribute and/or modify it under the terms of the GNU General Public|
%| License, as published by the Free Software Foundation (version 3 or    |
%| later,  http://www.gnu.org/licenses). For private use, dual licencing  |
%| is available, please contact us to activate a "pay for remove" option. |
%| CONTACT   : matthieu.aussal@polytechnique.edu                          |
%|             francois.alouges@polytechnique.edu                         |
%| WEBSITE   : www.cmap.polytechnique.fr/~aussal/gypsilab                 |
%|                                                                        |
%| Please acknowledge the gypsilab toolbox in programs or publications in |
%| which you use it.                                                      |
%|________________________________________________________________________|
%|   '&`   |                                                              |
%|    #    |   FILE       : femSubdivide.m                                |
%|    #    |   VERSION    : 0.40                                          |
%|   _#_   |   AUTHOR(S)  : Matthieu Aussal & François Alouges            |
%|  ( # )  |   CREATION   : 14.03.2017                                    |
%|  / 0 \  |   LAST MODIF : 14.03.2018                                    |
%| ( === ) |   SYNOPSIS   : Subdivide finite element space for domain     |
%|  `---'  |                decomposition method                          |
%|         |                ! ONLY WORKS FOR NON DIRICHLET FEM !          |
%+========================================================================+

% Number of subdomains is power of 2
if (Nsub ~= 2^floor(log2(Nsub)))
    error('femSubdivide.m : number of subdomain is not a power of 2.')
end

% Mesh and dof associated to finite element space
mesh          = fem.msh;
[dof,elt2dof] = fem.dof;                                   % to be improved

% Dof subdivision with binary tree
Nleaf  = ceil(length(fem)/Nsub);
bitree = tree(msh(dof),'binary',Nleaf);
Isub   = bitree{end}.ind;

% Security
if (length(Isub) ~= Nsub)
    error('femSubdivide.m : unavailable case')
end
if (norm(sort(cell2mat(Isub)) - (1:length(fem))','inf') > 1e-12)
    error('femSubdivide.m : unavailable case')
end

% Output initialization
domSub = cell(Nsub,1);
femSub = cell(Nsub,1);

% For each subindices
for i = 1:Nsub
    % Indices of valid dof and extended mesh
    I       = ismember(elt2dof,Isub{i});
    meshSub = mesh.sub(sum(I,2)>0);

    % Quadrature
    domSub{i}     = dom;
    domSub{i}.msh = meshSub;

    % Finite element
    femSub{i}     = fem;
    femSub{i}.msh = meshSub;

    % Dirichlet condition for boundary dof
    dir = setdiff( msh(femSub{i}.dof) , msh(dof(Isub{i},:)) );
    if (size(dir,1)>0)
        femSub{i} = dirichlet(femSub{i},dir);
    end

    % Security
    if (length(femSub{i}) ~= length(Isub{i}))
        error('femSubdivide.m : unavailable case');
    end

    % Graphical representation
    if fig
        tmp = mesh.sub(sum(I,2)==size(elt2dof,2));
        figure(fig)
        hold on
        plot(tmp,'b')
        plot(tmp.bnd,'r')
        plot(setdiff(meshSub,tmp),'w')
        axis equal
    end
end
end


%==============================================================================
% This code is part of the Matlab-based toolbox
% FAIR - Flexible Algorithms for Image Registration.
% For details see
% - https://github.com/C4IR and
% - http://www.siam.org/books/fa06/
%==============================================================================
% function [wOpt,his] = MLPIR(ML,varargin)
%
% Multi-Level Parametric Image Registration
%
% minimizes J^h(w) = D^h(T(Y(w)),R) + S^h(w) for h=coarse:fine
% see PIRobjFctn for a default objective and GaussNewton for a default optimizer
%
% Input:
%   ML        struct of coarse to fine representations of the data,
%           see getMultilevel.m
%  varargin optinonal parameters, see default parameter
%
% Output:
%  wOpt     optimal parameter for the parametric part
%  MLhis    iteration history
%
% for level=minLevel:maxLevel,
%   get data(level)
%   if level>minLevel, w0 = wOpt; end;
%   get wOpt running PIR  using w0 as starting guess
% end%For
%==============================================================================

function [wOpt,his] = MLPIR(ML,varargin)

if nargin == 0,
  help(mfilename);
  E6_HNSP_MLPIR_SSD_affine2D
  return;
end;

% setup default parameter for parametric pre-registration
PIRobj      = @PIRobjFctn;
PIRpara     = optPara('PIR-GN');
w0          = trafo('w0');  % starting guess
wStop       = w0;           % starting value (independent of level) used for stopping only

% configure regularizer
wRef        = w0;           % regularization: (w-wRef)'*M*(w-wRef)
M           = [];           %
beta        = 0;            % regularization: H -> H + beta*I
getGrid     = @getCellCenteredGrid;

% setup additional default parameter
pause       = 0;            % flag for pauses
plotIter    = 0;            % flag for output of iteration history each level
plotMLiter  = 0;            % flag for output of summarized iteration history
dimstr      = @(m) sprintf('m = [%s]',sprintf(' %d',m));

[ML,minLevel,maxLevel] = getMultilevel(ML);

for k=1:2:length(varargin),   % overwrites default parameter
  eval([varargin{k},'=varargin{',int2str(k+1),'};']);
end;

% intialize
omega = ML{end}.omega;
wOpt  = w0;
his   = [];

FAIRmessage(sprintf('%s: MultiLevel Image Registration',mfilename),'#')
fprintf('>> %-20s : %s\n','omega',dimstr(omega));
fprintf('>> %-20s : %s\n','IMAGE MODEL',imgModel);
fprintf('>> %-20s : %s\n','DISTANCE',distance);
fprintf('>> %-20s : %s\n','TRAFO',trafo);
fprintf('>> %-20s :\n','PARARMETRIC-REGISTRATION');
fprintf('>> %-20s : objective=<%s>, method = <%s>\n',...
  'PIR',func2str(PIRobj),func2str(PIRpara.scheme));
fprintf('>> %-20s : minLevel=%d, maxLevel=%d\n','LEVEL',minLevel,maxLevel);
FAIRmessage('#')

FAIRmessage(sprintf('%s, minLevel=%d:maxLevel=%d',...
  'MultiLevel Parametric Image Registration',minLevel,maxLevel));

tic;
% -- for loop over all levels ---------------------------------------------
for level=minLevel:maxLevel;

  FAIRmessage(sprintf('%s: level %d from %d to %d, %s',...
    mfilename,level,minLevel,maxLevel,dimstr(ML{level}.m)));

  % get data for current level, compute interpolation coefficients
  m     = ML{level}.m;
  [T,R] = imgModel('coefficients',ML{level}.T,ML{level}.R,omega);

  % update transformation
  trafo('set','omega',omega,'m',m);

  % initialize plots
  PIRpara.Plots('reset','mode','PIR-multi level','fig',level);
  PIRpara.Plots('init',struct('Tc',T,'Rc',R,'omega',omega,'m',m));

  % ----- call PIR ------------------------------------
  xc   = getGrid(omega,m);
  Rc   = imgModel(R,omega,center(xc,m));
  fctn = @(wc) PIRobj(T,Rc,omega,m,beta,M,wRef,xc,wc);
  if level == minLevel,
    fctn([]);   % report status
  else
    w0 = wOpt;  % update starting guess
  end;
  PIRpara.yStop = wStop;
  f = fieldnames(PIRpara);
  v = struct2cell(PIRpara);
  temp = reshape({f{:};v{:}},1,[]);
  [wOpt,hisPIR] = PIRpara.scheme(fctn,w0,temp{:});
  % ----- end PIR --------------------------------------

  if plotIter,
    plotIterationHistory(hisPIR,'J',[1,2,5],'fig',20+level);
  end;

  % update iteration history
  if level == minLevel,
    his.str = hisPIR.str;
    his.his = hisPIR.his;
  else
    his.his = [his.his;hisPIR.his];
  end;
  doPause(pause)

end;%for level
% -- for loop over all levels ---------------------------------------------
his.time = toc;
if plotMLiter,
  plotMLIterationHistory(his,'fig',30);
end;
if isempty(wStop), wStop = w0; end
his.reduction = hisPIR.his(end,2) / (hisPIR.his(1,2)+(hisPIR.his(1,2)==0));
J = find(his.his(:,1)==-1);
his.iter(minLevel:maxLevel) = his.his([J(2:end)-1;size(his.his,1)],1)';

FAIRmessage([mfilename,' : done !']);

%------------------------------------------------------------------------------

function doPause(p)
if strcmp(p,'on'),
  pause;
elseif p>0,
  pause(p);
end;
%==============================================================================


function [data] = ft_regressconfound(cfg, datain)

% FT_REGRESSCONFOUND estimates the regression weight of a set of confounds
% using a General Linear Model (GLM) and removes the estimated contribution
% from the single-trial data.
%
% Use as
%   timelock = ft_regressconfound(cfg, timelock)
% or as
%   freq     = ft_regressconfound(cfg, freq)
% or as
%   source   = ft_regressconfound(cfg, source)
%
% where timelock, freq, or, source come from FT_TIMELOCKANALYSIS,
% FT_FREQANALYSIS, or FT_SOURCEANALYSIS respectively, with keeptrials = 'yes'
%
% The cfg argument is a structure that should contain
%   cfg.confound    = matrix, [Ntrials X Nconfounds], may not contain NaNs
%
% The following configuration options are supported:
%   cfg.reject      = vector, [1 X Nconfounds], listing the confounds that
%                     are to be rejected (default = 'all')
%   cfg.normalize   = string, 'yes' or 'no', normalization to
%                     make the confounds orthogonal (default = 'yes')
%   cfg.output      = 'residual' (default), 'beta', or 'model'.
%                     If 'residual' is specified, the output is a data
%                     structure containing the residuals after regressing
%                     out the in cfg.reject listed confounds. If 'beta' or 'model'
%                     is specified, the output is a data structure containing
%                     the regression weights or the model, respectively.
%
% This method is described by Stolk et al., Online and offline tools for head
% movement compensation in MEG (Neuroimage, 2013)
%
% To facilitate data-handling and distributed computing you can use
%   cfg.inputfile   =  ...
%   cfg.outputfile  =  ...
% If you specify one of these (or both) the input data will be read from a *.mat
% file on disk and/or the output data will be written to a *.mat file. These mat
% files should contain only a single variable, corresponding with the
% input/output structure.
%
% See also FT_REJECTCOMPONENT, FT_REJECTARTIFACT

% Undocumented local options:
%   cfg.ftest       = string array, {N X Nconfounds}, to F-test whether
%                     the full model explains more variance than reduced models
%                     (e.g. {'1 2'; '3 4'; '5'} where iteratively the added value of
%                     regressors 1 and 2, and then 3 and 4, etc., are tested)
%   cfg.statistics  = string, 'yes' or 'no', whether to add the statistics
%                     on the regression weights to the output (default = 'no',
%                     applies only when cfg.output = 'beta')

% Copyright (C) 2011-2017, Arjen Stolk, Robert Oostenveld, Lennart Verhagen
%
% This file is part of FieldTrip, see http://www.fieldtriptoolbox.org
% for the documentation and details.
%
%    FieldTrip 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 3 of the License, or
%    (at your option) any later version.
%
%    FieldTrip 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 FieldTrip. If not, see <http://www.gnu.org/licenses/>.
%
% $Id$

% these are used by the ft_preamble/ft_postamble function and scripts
ft_revision = '$Id$';
ft_nargin   = nargin;
ft_nargout  = nargout;

% do the general setup of the function
ft_defaults
ft_preamble init
ft_preamble debug
ft_preamble loadvar datain
ft_preamble provenance datain
ft_preamble trackconfig

% the ft_abort variable is set to true or false in ft_preamble_init
if ft_abort
  return
end

% check if the input data is valid for this function
datain = ft_checkdata(datain, 'datatype', {'timelock', 'freq', 'source'}, 'feedback', 'yes');

if isfield(cfg, 'beta') || isfield(cfg, 'model')
 ft_error('The options cfg.beta and cfg.model have been removed as of Aug 2017, please use cfg.output instead');
end

% ensure that the required options are present
cfg = ft_checkconfig(cfg, 'required', {'confound'}, 'renamed', {'Ftest','ftest'}, 'forbidden', {'beta','model'});

% specify the defaults
cfg.confound   = ft_getopt(cfg, 'confound');
cfg.reject     = ft_getopt(cfg, 'reject', 'all');
cfg.normalize  = ft_getopt(cfg, 'normalize', 'yes');
cfg.output     = ft_getopt(cfg, 'output', 'residual');
cfg.statistics = ft_getopt(cfg, 'statistics', 'no');
cfg.ftest      = ft_getopt(cfg, 'ftest');
cfg.parameter  = ft_getopt(cfg, 'parameter'); % the default is handled further down

regr = cfg.confound;
if any(isnan(regr(:)))
  ft_error('the confounds may not contain NaNs');
end
nconf     = size(regr,2);
conflist  = 1:nconf;
if strcmp(cfg.reject, 'all')
  cfg.reject = conflist(1:end); % to be removed
else
  cfg.reject = intersect(conflist, cfg.reject); % to be removed
end

fprintf('removing confound %s \n', num2str(cfg.reject));
kprs = setdiff(conflist, cfg.reject); % to be kept
fprintf('keeping confound %s \n', num2str(kprs));

% confound normalization for orthogonality
if strcmp(cfg.normalize, 'yes')
  fprintf('normalizing the confounds, except the constant \n');
  for c = 1:nconf
    AVG = mean(regr(:,c));
    STD = std(regr(:,c),0,1);
    if abs(STD/AVG)<10*eps
      fprintf('confound %s is a constant \n', num2str(c));
    else
      regr(:,c) = (regr(:,c) - AVG) / STD;
    end
    clear AVG STD;
  end
else
  fprintf('skipping normalization procedure \n');
end

switch ft_datatype(datain)
  case 'freq'
    cfg.parameter = ft_getopt(cfg, 'parameter', 'powspctrm');
  case 'timelock'
    cfg.parameter = ft_getopt(cfg, 'parameter', 'trial');
  case 'source'
    cfg.parameter = ft_getopt(cfg, 'parameter', 'pow');
end

dimord = getdimord(datain, cfg.parameter);
dimsiz = getdimsiz(datain, cfg.parameter);
dimtok = tokenize(dimord, '_');
rptdim = find(strcmp(dimtok, 'rpt'));
datdim = setdiff(1:length(dimtok), rptdim);

nrpt = dimsiz(rptdim);

dat = datain.(cfg.parameter);
if strcmp(dimtok{1}, '{pos}')
  indx = find(datain.inside);
  npos = length(indx);
  tmp = nan([nrpt npos datdim(2:end)]); % only positions inside the brain
  for i=indx'
    tmp(:,i,:,:,:) = dat{i}(:,:,:,:);
  end
  dat = tmp;
  haspermuted = false;
else
  dat = permute(dat, [rptdim datdim]);
  haspermuted = true;
end

dat = reshape(dat, nrpt, []);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% GLM MODEL
%   Y = X * B + err, where Y is data, X is the model, and B are beta's
% which means
%   Best = X\Y ('matrix division', which is similar to B = inv(X)*Y)
% or when presented differently
%   Yest = X * Best
%   Yest = X * X\Y
%   Yclean = Y - Yest (the true 'clean' data is the recorded data 'Y' -
%   the data containing confounds 'Yest')
%   Yclean = Y - X * X\Y
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% estimate and remove the confounds
fprintf('estimating the regression weights and removing the confounds \n');
if isempty(find(isnan(dat))) % if there are no NaNs, process all at once
  beta = regr\dat;                                                        % B = X\Y
else % otherwise process per colum set as defined by the nan distribution
  [u,i,j] = unique(~isnan(dat)','rows','first'); % find unique rows
  uniquecolumns = u'; % unique column types
  Nuniques = numel(i); % number of unique types
  beta_temp = NaN(Nuniques, nconf, size(dat,2)); % declare empty variable
  for n = 1:Nuniques % for each unique type
    rowidx = find(uniquecolumns(:,n)==1); % row indices for unique type
    colidx = find(j==n); % column indices for unique type
    if any(uniquecolumns(:,n)) % if vector contains a nonzero number
      beta_temp(n,:,colidx) = regr(rowidx,:)\dat(rowidx,colidx);         % B = X\Y
    end
  end
  beta = reshape(nansum(beta_temp,1),[nconf size(dat,2)]); % sum the betas
  clear beta_temp
end

model = regr(:, cfg.reject) * beta(cfg.reject, :);                        % model = confounds * weights = X * X\Y
Yc = dat - model;                                                         % Yclean = Y - X * X\Y

% reduced models analyses
if ~isempty(cfg.ftest)
  dfe        = nrpt - nconf;                                              % degrees of freedom
  err        = dat - regr * beta;                                         % err = Y - X * B
  tmse       = sum((err).^2)/dfe;                                         % mean squared error
  for iter = 1:numel(cfg.ftest)
    % regressors to test if they explain additional variance
    r          = str2num(cfg.ftest{iter});
    fprintf('F-testing explained additional variance of regressors %s \n', num2str(r));
    % regressors in reduced design (that is the original design)
    ri         = ~ismember(1:size(regr,2),r);
    rX         = regr(:,ri);               % reduced design
    rnr        = size(rX,2);               % number of regressors in reduced design
    % estimate reduced model betas
    rXcov      = pinv(rX'*rX);             % inverse design covariance matrix
    rb         = rXcov*rX'*dat;          	 % beta estimates using pinv
    % calculate mean squared error of reduced model
    rdfe       = size(dat,1) - size(rX,2); % degrees of freedom of the error
    rerr       = dat-rX*rb;                % residual error
    rmse       = sum(rerr'.^2,2)./rdfe;	   % mean squared error
    % F-test
    F(iter,:)  = ((rmse'-tmse)./(nconf-rnr)) ./ (tmse./(dfe-2));
    % Rik Henson defined F-test
    % F = ( ( rerr'*rerr - err'*err ) / ( nconf-rnr ) ) / ( err'*err/ ( nrpt-nconf ) );
    % convert F-value to p-value
    idx_pos    = F(iter,:) >= 0;
    idx_neg    = ~idx_pos;
    p(iter,:)     = nan(1,size(F(iter,:),2));
    p(iter,idx_pos) = (1-fcdf(F(iter,idx_pos),rnr,rdfe));
    p(iter,idx_neg) = fcdf(-F(iter,idx_neg),rnr,rdfe);
    clear rerr rmse
    % FIXME: drop in replace tcdf from the statfun/private dir
  end
  clear dfe err tmse
end

% organize the output
dataout = keepfields(datain, {'label', 'time', 'freq', 'pos', 'dim', 'transform', 'inside', 'outside', 'trialinfo', 'sampleinfo', 'dimord'});
switch cfg.output
  case 'residual'
    dataout.(cfg.parameter) = reshape(Yc, [nrpt dimsiz(datdim)]); % either powspctrm, trial, or pow
    if haspermuted
      dataout.(cfg.parameter) = ipermute(dataout.(cfg.parameter), [rptdim datdim]);
    end
    clear Yc
  case 'beta'
    dataout.beta = reshape(beta, [nconf, dimsiz(datdim)]);
    if haspermuted
      dataout.beta = ipermute(dataout.beta, [rptdim datdim]);
    end
    if strcmp(cfg.statistics, 'yes') % beta statistics
      fprintf('performing statistics on the regression weights \n');
      dfe        = nrpt - nconf;                                              % degrees of freedom
      err        = dat - regr * beta;                                         % err = Y - X * B
      mse        = sum((err).^2)/dfe;                                         % mean squared error
      covar      = diag(regr'*regr)';                                         % regressor covariance
      bvar       = repmat(mse',1,size(covar,2))./repmat(covar,size(mse,2),1); % beta variance
      tval       = (beta'./sqrt(bvar))';                                      % betas -> t-values
      prob       = (1-tcdf(tval,dfe))*2;                                      % p-values
      clear err dfe mse bvar
      % FIXME: drop in replace tcdf from the statfun/private dir
      dataout.stat = reshape(tval, [nconf dimsiz(datdim)]);
      dataout.prob = reshape(prob, [nconf dimsiz(datdim)]);
      if haspermuted
        dataout.stat = ipermute(dataout.stat, [rptdim datdim]);
        dataout.prob = ipermute(dataout.prob, [rptdim datdim]);
      end
      clear tval prob
    end
  case 'model'
    dataout.model = keepfields(datain, {'label', 'time', 'freq', 'pos', 'dim', 'transform', 'inside', 'outside', 'trialinfo', 'sampleinfo', 'dimord'});
    dataout.model.(cfg.parameter) = reshape(model, [nrpt, dimsiz(datdim)]);
    if haspermuted
      dataout.model.(cfg.parameter) = ipermute(dataout.model.(cfg.parameter), [rptdim datdim]);
    end
  otherwise
    error('output ''%s'' is not supported', cfg.output);
end

% reduced models analyses
if ~isempty(cfg.ftest)
  dataout.fvar   = reshape(F, [numel(cfg.ftest) dimsiz(datdim)]);
  dataout.pvar   = reshape(p, [numel(cfg.ftest) dimsiz(datdim)]);
  clear F p
end

% discard the gradiometer information because the weightings have been changed
if isfield(dataout, 'grad')
  ft_warning('discarding gradiometer information because the weightings have been changed');
  dataout = rmfield(dataout, 'grad');
end

% discard the electrode information because the weightings have been changed
if isfield(dataout, 'elec')
  ft_warning('discarding electrode information because the weightings have been changed');
  dataout = rmfield(dataout, 'elec');
end

% do the general cleanup and bookkeeping at the end of the function
ft_postamble debug
ft_postamble trackconfig
ft_postamble previous datain

% rename the output variable to accomodate the savevar postamble
data = dataout;

ft_postamble provenance data
ft_postamble history    data
ft_postamble savevar    data


function [ar,e,dc]=lpccovar(s,p,t,w)
%LPCCOVAR performs covariance LPC analysis [AR,E,DC]=(S,P,T)
%
%  Inputs:  S(NS)    is the input signal
%           P        is the order (default: 12)
%           T(NF,:)  specifies the frames size details: each row specifies one frame
%                    T can be a cell array if rows have unequal numbers of values
%                       T(:,1) gives the start of the analysis interval: must be >P
%                       T(:,2) gives the end of the anaylsis interval [default: t(:+1,1)-1]
%                       subsequent pairs can be used to specify multiple disjoint segments
%                    If T is omitted, T(1,1)=P+1, T(1,2)=NS;
%                    The elements of t need not be integers.
%           W(NS)    The error at each sample is weighted by W^2 (default: 1)
%
% Outputs:  AR(NF,P+1)  are the AR coefficients with AR(:,1) = 1
%           E(NF,2)     is the energy in the residual and in the original window.
%                       sqrt(E) is often called the 'gain' of the LPC filter.
%           DC          is the DC component of the signal S. If this output is included,
%                       the LPC equations are modified to include a DC offset.

% Notes:
%
% (1) For speech processing P should be at least 2*F*L/C where F is the sampling
%     frequency, L the vocal tract length and C the speed of sound. For a typical
%     male (l=17 cm) this gives f/1000.
%
% (2) Each analysis frame should contain at least 2P samples. If note (1) is followed
%     this implies at least 2 ms of speech signal per frame.
%
% (3) It can be advantageous to restrict the analysis regions to time intervals
%     when the glottis is closed (closed-phase analysis). This can be achieved by
%     setting the T input parameter appropriately. If the closed-phase is shorter than
%     2 ms then two or more successive closed-phases should be used by defining 4 or more
%     elements in the corresponding row of T.
%
% (4) A previous version of this routine allowed T() to have a single row which would
%     be replicated for the entire file length. This has be removed because it gave rise
%     to an ambiguity.

%  Bugs: sould really detect a singular matrix and reduce the order accordingly

%	   Copyright (C) Mike Brookes 1995
%      Version: $Id: lpccovar.m 713 2011-10-16 14:45:43Z dmb $
%
%   VOICEBOX is a MATLAB toolbox for speech processing.
%   Home page: http://www.ee.ic.ac.uk/hp/staff/dmb/voicebox/voicebox.html
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   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 can obtain a copy of the GNU General Public License from
%   http://www.gnu.org/copyleft/gpl.html or by writing to
%   Free Software Foundation, Inc.,675 Mass Ave, Cambridge, MA 02139, USA.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

s = s(:); % make it a column vector
if nargin < 2 p=12; end;
if nargin < 3 t=[p+1 length(s)]; end;
wq = nargin>3;
[nf,ng]=size(t);
if iscell(t)
   t{nf+1}=length(s)+1;
else
   if rem(ng,2)
      t(:,end+1)=[t(2:nf,1)-1; length(s)];
   end
end
ar=zeros(nf,p+1);
ar(:,1)=1;
e=zeros(nf,2);
dc=zeros(nf,1);
d0=nargout >2;

rs=(1:p);

for jf=1:nf
   if iscell(t)
      tj=t{jf};
      if rem(length(tj),2)
         tj(end+1)=t{jf+1}(1)-1;
      end
   else
      tj=t(jf,:);
   end

   ta = ceil(tj(1));
   tb = floor(tj(2));
   cs = (ta:tb).';
   for js=3:2:length(tj)
      ta = ceil(tj(js));
      tb = floor(tj(js+1));
      cs = [cs; (ta:tb).'];
   end
   %disp(cs([logical(1); (cs(2:end-1)~=cs(1:end-2)+1)|(cs(2:end-1)~=cs(3:end)-1); logical(1)])');
   nc = length(cs);
   pp=min(p,nc-d0);
   dm=zeros(nc,pp);	% predefine shape
   dm(:) = s(cs(:,ones(1,pp))-rs(ones(nc,1),1:pp));
   if nargout>2
      if wq
         dm = [ones(nc,1) dm].*w(cs(:,ones(1,1+pp)));
         sc=(s(cs).*w(cs));
         aa = (dm\sc).';
      else
         dm = [ones(nc,1) dm];
         sc=s(cs);
         aa = (dm\sc).';
      end
      ar(jf,2:pp+1) = -aa(2:pp+1);
      e(jf,1)=sc.'*(sc - dm*aa.');
      e(jf,2)=sc.'*sc;
      dc(jf)=aa(1)/sum(ar(jf,:));
   else
      if wq
         dm = dm.*w(cs(:,ones(1,pp)));
         sc=(s(cs).*w(cs));
         aa = (dm\sc).';
      else
         sc=s(cs);
         aa = (dm\sc).';
      end;
      ar(jf,2:pp+1) = -aa;
      if nargout>1
         e(jf,1)=sc.'*(sc - dm*aa.');
         e(jf,2)=sc.'*sc;
      end
   end
end
if ~nargout
   lpcar2db(ar,127);
end


function [xUpdate,SUpdate,innov,Szz,W]=sqrtCubKalUpdateWithPred(z,SR,zPred,otherInfo)
%%SQRTCUBKALUPDATEWITHPRED Given the output of the measurement prediction
%           step from sqrtCubKalMeasPred and a measurement, complete the
%           measurement update step of the square root cubature Kalman
%           filter. Separating the measurement prediction step from the
%           rest of the update step can make the creation of multiple
%           measurement association hypotheses from a single target
%           prediction more efficient. The full measurement update function
%           is sqrtCubKalUpdate.
%
%INPUTS: z The zDimX1 measurement vector.
%       SR The zDimXzDim lower-triangular square root of the measurement
%          covariance matrix in the native coordinate system of the
%          measurement.
%    zPred The zDimXnumComp measurement predictions from the filter.
% otherInfo The intermediate results returned in the otherInfo output of
%          the sqrtCubKalMeasPred function.
%
%OUTPUTS: xUpdate The xDimXnumComp updated state vectors.
%         SUpdate The updated xDimXxDimXnumComp lower-triangular square-
%                 root state covariance matrices.
%      innov, Szz The zDimXnumComp innovations and the zDimXzDimXnumComp
%                 square-root innovation covariance matrices are returned
%                 in case one wishes to analyze the consistency of the
%                 estimator or use those values in gating or likelihood
%                 evaluation.
%               W The xDimXzDimXnumComp gains used in the update.
%
%See the comments to the function sqrtCubKalMeasPred for an example of
%usage of this function. See the comments to sqrtCubKalUpdate for more
%information on the algorithm.
%
%June 2018 David F. Crouse, Naval Research Laboratory, Washington D.C.
%(UNCLASSIFIED) DISTRIBUTION STATEMENT A. Approved for public release.

    xPred=otherInfo.xPred;
    innovTrans=otherInfo.innovTrans;
    stateDiffTrans=otherInfo.stateDiffTrans;
    xPredCenPoints=otherInfo.xPredCenPoints;
    stateTrans=otherInfo.stateTrans;
    zPredCenPoints=otherInfo.zPredCenPoints;
    Pxz=otherInfo.Pxz;

    xDim=size(xPred,1);
    numComp=size(xPred,2);
    zDim=size(z,1);

    xUpdate=zeros(xDim,numComp);
    SUpdate=zeros(xDim,xDim,numComp);
    innov=zeros(zDim,numComp);
    Szz=zeros(zDim,zDim,numComp);
    W=zeros(xDim,zDim,numComp);

    for k=1:numComp
        Szz(:,:,k)=tria([zPredCenPoints(:,:,k),SR]);

        %The filter gain
        W(:,:,k)=(Pxz(:,:,k)/Szz(:,:,k)')/Szz(:,:,k);

        %The innovation, transformed as necessary to keep values in a
        %desired range.
        innov(:,k)=innovTrans(z,zPred(:,:,k));

        %Updated state estimate
        xUpdate(:,k)=stateTrans(xPred(:,k)+W(:,:,k)*innov(:,k));

        %Updated state root covariance
        SUpdate(:,:,k)=tria([stateDiffTrans(xPredCenPoints(:,:,k)-W(:,:,k)*zPredCenPoints(:,:,k)),W(:,:,k)*SR]);
    end
end

%LICENSE:
%
%The source code is in the public domain and not licensed or under
%copyright. The information and software may be used freely by the public.
%As required by 17 U.S.C. 403, third parties producing copyrighted works
%consisting predominantly of the material produced by U.S. government
%agencies must provide notice with such work(s) identifying the U.S.
%Government material incorporated and stating that such material is not
%subject to copyright protection.
%
%Derived works shall not identify themselves in a manner that implies an
%endorsement by or an affiliation with the Naval Research Laboratory.
%
%RECIPIENT BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE
%SOFTWARE AND ANY RELATED MATERIALS, AND AGREES TO INDEMNIFY THE NAVAL
%RESEARCH LABORATORY FOR ALL THIRD-PARTY CLAIMS RESULTING FROM THE ACTIONS
%OF RECIPIENT IN THE USE OF THE SOFTWARE.


classdef DWU < ALGORITHM
% <multi> <real/integer/label/binary/permutation>
% Dominance-weighted uniformity multi-objective evolutionary algorithm

%------------------------------- Reference --------------------------------
% G. Moreira and L. Paquete, Guiding under uniformity measure in the
% decision space, Proceedings of the IEEE Latin American Conference on
% Computational Intelligence, 2019.
%------------------------------- Copyright --------------------------------
% Copyright (c) 2023 BIMK Group. You are free to use the PlatEMO for
% research purposes. All publications which use this platform or any code
% in the platform should acknowledge the use of "PlatEMO" and reference "Ye
% Tian, Ran Cheng, Xingyi Zhang, and Yaochu Jin, PlatEMO: A MATLAB platform
% for evolutionary multi-objective optimization [educational forum], IEEE
% Computational Intelligence Magazine, 2017, 12(4): 73-87".
%--------------------------------------------------------------------------

% This function is written by Gladston Moreira

    methods
        function main(Algorithm,Problem)
            %% Generate random population
            Population = Problem.Initialization();
            [Population,FrontNo] = EnvironmentalSelection(Population,Problem.N);

            %% Optimization
            while Algorithm.NotTerminated(Population)
                MatingPool = TournamentSelection(2,Problem.N,FrontNo);
                Offspring  = OperatorGA(Problem,Population(MatingPool));
                [Population,FrontNo,] = EnvironmentalSelection([Population,Offspring],Problem.N);
            end
        end
    end
end


function node_order = tet_mesh_node_order ( tetra_order, tetra_num, ...
  tetra_node, node_num )

%*****************************************************************************80
%
%% TET_MESH_NODE_ORDER: determines the order of nodes.
%
%  Discussion:
%
%    The order of a node is the number of tetrahedrons that use that node
%    as a vertex.
%
%  Licensing:
%
%    This code is distributed under the GNU LGPL license.
%
%  Modified:
%
%    27 October 2005
%
%  Author:
%
%    John Burkardt
%
%  Parameters:
%
%    Input, integer TETRA_ORDER, the order of the tetrahedrons.
%
%    Input, integer TETRA_NUM, the number of tetrahedrons.
%
%    Input, integer TETRA_NODE(TETRA_ORDER,TETRA_NUM), the nodes
%    that make up the tetrahedrons.
%
%    Input, integer NODE_NUM, the number of nodes.
%
%    Output, integer NODE_ORDER(NODE_NUM), the order of each node.
%
  node_order(1:node_num) = 0;

  for tetra = 1 : tetra_num
    for i = 1 : tetra_order
      node = tetra_node(i,tetra);
      if ( node < 1 | node_num < node )
        fprintf ( 1, '\n' );
        fprintf ( 1, 'TET_MESH_NODE_ORDER - Fatal error!\n' );
        fprintf ( 1, '  Illegal entry in TETRA_NODE.\n' );
        error ( 'TET_MESH_NODE_ORDER - Fatal error!' );
      else
        node_order(node) = node_order(node) + 1;
      end
    end
  end

  return
end