function [X,r,s,T] = mrgras(X0,u,v,G,Q,W,eps)
% PURPOSE: estimate a new multi-regional (or any partitioned) matrix X as 
% close as possible to a given matrix X0 subject to the row sums, column 
% sums and non-overlapping aggregation constraints, using MR-GRAS approach
% -------------------------------------------------------------------------
% USAGE: 
% Write X = mrgras(X0,u,v,G,Q,W) OR [X,r,s,T] = mrgras(X0,u,v,G,Q,W) with  
% or without eps as the seventh optional input argument, where
%
% INPUT:
% -> X0 = benchmark (original) matrix, not necessarily square
% -> u = column vector of row totals (new row sums) 
% -> v = column vector of column totals (new column sums)
% -> W = non-overlapping aggregation constraints matrix
% -> G & Q = the row and column aggregator matrices such that G*X0*Q = W;
%            non-overlapping aggergation necessarily implies that the 
%            column sums of G  and the row sums of Q must be all unity
% -> eps = convergence tolerance level; if empty, the default threshold 
%          level is 0.1e-5 (=0.000001)
% -> In case of missing w_IJ, set the corresponding missing number to 
%    w_IJ = 1010101 (assuming that 1010101 is not among the w_IJ's)
%
% OUTPUT (input-output analysis-related interpretation):
% -> X = updated/balanced/adjusted/projected matrix
% -> r = substitution effects (row multipliers)
% -> s = fabrication effects (column multipliers)
% -> T = technology or regional effects (aggregation multipliers)
% -------------------------------------------------------------------------
% REFERENCES: 
% Temursho, U., Oosterhaven, J. and M.A. Cardenete (2019), A multi-regional 
% generalized RAS updating technique, IOpedia Research Paper No. 2,  
% September 2019, www.IOpedia.eu
% -------------------------------------------------------------------------
% NOTE FROM THE AUTHOR: Using this program and publishing the results in
% the form of a report, working/discussion paper, journal article, etc.
% requires citation of the above reference paper. 
% -------------------------------------------------------------------------
% Written by:   Umed Temursho, May 2019 
%               E-mail: utemursho@gmail.com

[m,n] = size(X0);
[h,k] = size(W);
N = zeros(m,n);
N(X0<0) = -X0(X0<0); 
N = sparse(N);      % could save memory with large-scale matrices
P = X0+N;
P = sparse(P);
%
r0 = ones(m,1);      % initial guess for r in step 0 
T0 = ones(h,k);      % initial guess for T in step 0
Te = G'*T0*Q';       % T expanded to fit the dimention of X0
p_rt = (P.*Te)'*r0;
n_rt = (N.*invM(Te))'*invd(r0)*ones(m,1);
s1 = invd(2*p_rt)*(v+sqrt(v.^2+4*p_rt.*n_rt));    % first step s
ss = -invd(v)*n_rt;
s1(p_rt==0) = ss(p_rt==0); 
%
p_st = (P.*Te)*s1;
n_st = (N.*invM(Te))*invd(s1)*ones(n,1);
r1 = invd(2*p_st)*(u+sqrt(u.^2+4*p_st.*n_st));     % first step r
rr = -invd(u)*n_st;
r1(p_st==0) = rr(p_st==0);
%
P_rs = G*(diag(r1)*P*diag(s1))*Q;
N_rs = G*(invd(r1)*N*invd(s1))*Q;
T1 = invM(2*P_rs).*(W+sqrt(W.^2+4*P_rs.*N_rs));    % first step T
TT = -invM(W).*N_rs;
T1(P_rs==0) = TT(P_rs==0);
T1(W==1010101) = 1;    % set t_IJ=1 for missing aggregation total w_IJ 
Te = G'*T1*Q'; 
% 
p_rt = (P.*Te)'*r1;
n_rt = (N.*invM(Te))'*invd(r1)*ones(m,1);
s2 = invd(2*p_rt)*(v+sqrt(v.^2+4*p_rt.*n_rt));    % second step s  
ss = -invd(v)*n_rt;
s2(p_rt==0) = ss(p_rt==0); 
%
p_st = (P.*Te)*s2;
n_st = (N.*invM(Te))*invd(s2)*ones(n,1);
r2 = invd(2*p_st)*(u+sqrt(u.^2+4*p_st.*n_st));     % second step r
rr = -invd(u)*n_st;
r2(p_st==0) = rr(p_st==0);
%
P_rs = G*(diag(r2)*P*diag(s2))*Q;
N_rs = G*(invd(r2)*N*invd(s2))*Q;
T2 = invM(2*P_rs).*(W+sqrt(W.^2+4*P_rs.*N_rs));    % second step T
TT = -invM(W).*N_rs;
T2(P_rs==0) = TT(P_rs==0);
T2(W==1010101) = 1;    % set t_IJ=1 for missing w_IJ 
%
tmax = max(max(abs(T2-T1)));
dif = [s2-s1;r2-r1;tmax];
iter = 1              % first iteration
if  nargin < 7 || isempty(eps)
    eps = 0.1e-5;     % default tolerance level
end        
M = max(abs(dif));
while (M > eps)
    s1 = s2;
    r1 = r2;
    T1 = T2;
    Te = G'*T1*Q';
    %
    p_rt = (P.*Te)'*r1;
    n_rt = (N.*invM(Te))'*invd(r1)*ones(m,1);
    s2 = invd(2*p_rt)*(v+sqrt(v.^2+4*p_rt.*n_rt));      % next step s  
    ss = -invd(v)*n_rt;
    s2(p_rt==0) = ss(p_rt==0);  
    %
    p_st = (P.*Te)*s2;
    n_st = (N.*invM(Te))*invd(s2)*ones(n,1);
    r2 = invd(2*p_st)*(u+sqrt(u.^2+4*p_st.*n_st));      % next step r
    rr = -invd(u)*n_st;
    r2(p_st==0) = rr(p_st==0);
    %
    P_rs = G*(diag(r2)*P*diag(s2))*Q;
    N_rs = G*(invd(r2)*N*invd(s2))*Q;
    T2 = invM(2*P_rs).*(W+sqrt(W.^2+4*P_rs.*N_rs));     % next step T
    TT = -invM(W).*N_rs;
    T2(P_rs==0) = TT(P_rs==0);
    T2(W==1010101) = 1;                 % set t_IJ=1 for missing w_IJ
    %
    tmax = max(max(abs(T2-T1)));
    dif = [s2-s1;r2-r1;tmax];
    iter = iter+1
    M = max(abs(dif));
end
s = s2;     % final step s
r = r2;     % final step r
T = T2;     % final step T
Te = G'*T*Q'; 
%
X = Te.*(diag(r)*P*diag(s))-invM(Te).*(invd(r)*N*invd(s));  % updated matrix
end

function invd = invd(x)     % auxiliary function used above
invd = 1./x;
invd(x==0) = 1;
invd = diag(invd);
end

function invM = invM(X)     % auxiliary function used above
invM = 1./X;
invM(X==0) = 1;
end
%--------------------------------------------------------------------------
% END OF THE CODE
%--------------------------------------------------------------------------