% Identification of ARMAX models through PEM

function ARMAX_main()
clear all
close all
clc
 
  

    N=8000; % Length of training dataset
    Nval=2000; % Length of validation dataset

% Initialize random generator
    randn('seed',1000);
    rand('seed',900);
    
% True ARMAX system parameters (used only to generate artificial data). The
% system has 1 input and 1 output
    a1o=0.9; a2o=0.3; 
    b1o=1; b2o=-0.5;  
    c1o=0.4; c2o=0.8;  
        
    z=tf('z',1);
    Gou=[b1o*z^(-1)+b2o*z^(-2)]/(1+a1o*z^(-1)+a2o*z^(-2));   %True  tf from input  to output
    Goe=[1+c1o*z^(-1)+c2o*z^(-2)]/(1+a1o*z^(-1)+a2o*z^(-2)); %True DT from  error to output
    
% Generate input u
    u=rand(N+Nval,1);

% Generate white noise e
    sigmae=0.3; %Standard deviation of the noise
    e=randn(N+Nval,1)*sigmae;
    
% Generate  output     
    y=lsim([Gou Goe],[u e]); 

% Generate noise-free output (only used to assess model quality)
    yo=lsim([Gou],[u]); 
    
    
% Compute Signal to Noise Ratio (SNR) on each output channel to see the amplitude of the error w.r.t the output        
 
        SNR=10*log10(norm(y(:,1))^2/norm(e(:,1))^2);
        fprintf('SNR on output channel %2.0f =  %3.2f db \n',1, SNR);
 
    
% Identification
    na=3; nb=[4]; nc=[2]; %Chosen order of the models (an overparameterized model is considered) 
    
    
    %Initial conditions for theta (be careful: even at the initial conditions, the filter 1/(C(q)) should
    %be stable!)
    theta0=zeros(na+nb+nc,1);
    maxiter=200; %Maximum number of iterations in the numerical optimization
    tol=10^(-8); %Stopping criterion
    opt_method='gradient'; %Optimization algorithm: it can be either gradient or Gauss-Newton method
    maxgridlinesearch=100; % Number of grid points used in the line search
    delta=0.001; %Regularization parameter for Gauss-Newton method
    
    [modelestLS Phi thetaLS]=estimateLS(u(1:N),y(1:N),na,nb); %Compute LS estimate (just for comparison)
    fprintf('LS estimate: Identified parameters: \n');
    as=sprintf('%5.2f ', modelestLS.a');
    bs=sprintf('%5.2f ', modelestLS.b{1}');
    fprintf('a: %s \n',as);
    fprintf('b: %s \n',bs);
 
    fprintf('Press a bottom to continue \n');
    pause();
    
    fprintf('Identification of ARMAX models through PEM started\n');
    modelest=id_ARMAX(u(1:N),y(1:N),na,nb,nc,opt_method,theta0,maxiter,tol,maxgridlinesearch,delta);   %Identification of ARMAX models through PEM
    fprintf('Identification of ARMAX models through PEM  ended\n');
    
    fprintf('PEM estimate: Identified parameters: \n');
    as=sprintf('%5.2f ', modelest.a');
    bs=sprintf('%5.2f ', modelest.b{1}');
    cs=sprintf('%5.2f ', modelest.c');
    fprintf('a: %s \n',as);
    fprintf('b: %s \n',bs);
    fprintf('c: %s \n',cs);

    % simulate model output
    ysim=simulateModel(u(N+1:N+Nval),1,modelest);
    
  
    figure, plot(ysim,'r'), hold on, plot(y(N+1:Nval+N),'b'), plot(yo(N+1:Nval+N),'k');
    xlabel('samples')
    ylabel('Output y')
    legend('simulated','measured','noise-free')
 
  
    
end




function modelest=id_ARMAX(u,y,na,nb,nc,opt_method,theta0,maxiter,tol,maxgrid,delta)
% modelest=id_ARMAX(u,y,na,nb,nc,opt_method,theta0,maxiter,tol,maxgrid,delta) 
% Estimate of ARMAX models through PEM
% Inputs: 
%   u: input data
%   y: output data
%   na,nb,nc: order of ARMAX polynomials A(q), B(q) and C(q)
%   opt_method: optimization method (either 'gradient' or 'Gauss-Newton')
%   theta0: initial conditions for optimization
%   maxiter maximum numberr of iterations (optional, default:100)
%   tolerance: exit criterion (optional, default: 10^-6)
%   maxgrid: maximum number of gridding points for exact line search (optional, default:100)
%   delta= regularization parameter for the Hessian (optional, default 0 )
% Outputs:
%   modelest: structure containing the estimated model parameters 


if exist('maxiter')==0
   maxiter=100; 
end
if exist('tol')==0
   tol=10^(-6); 
end
if exist('maxgrid')==0
   maxgrid=50; 
end
if exist('delta')==0
   delta=0; 
end
if (strcmp(opt_method,'Gauss-Newton')+strcmp(opt_method,'gradient'))==0
   error('Not valid optimization method'); 
end

%Initial conditios 
    aest=theta0(1:na);
    best=theta0(na+1:nb+na);
    cest=theta0(na+nb+1:nb+na+nc);
   % thetaest=theta0;    
    z=tf('z',1);
    
    for inditer=1:maxiter %Numerical algorithm for optimization
        thetaest=[aest;best;cest];

        %Define filter 1/C;
        numFilterC=[1;zeros(length(cest),1)]';
        denFilterC=[1; cest]';
        FilterC=tf(numFilterC, denFilterC,1); % 1/C(z)
        
        % Define filter A/C;
        FilterAC=tf([1; aest]', denFilterC,1); % A(z)/C(z)
        if na-nc>0
            FilterAC=FilterAC/z^(na-nc);
        end
        
        % Define filter B/C;
        FilterBC=tf([0; best]', denFilterC,1); % B(z)/C(z)
        if nb-nc>0
            FilterBC=FilterBC/z^(nb-nc);
        end
        
        perr=lsim([FilterAC -FilterBC],[y u]); %Prediction error  
        
        yf=lsim(FilterC,y); % Filtered output
        uf=lsim(FilterC,u); % Filtered input
        pef=lsim(FilterC,perr); % Filtered prediction error
        
        % Compute gradient: d pe(k)/ d theta=[yf(k-1)...yf(k-na) -uf(k-1) ... -uf(k-nb) -pef(k-1) .... -pef(k-nc)];
            
        delaymax=max(max(na,nb),nc);        
        pesgrad=zeros(length(theta0),1);
        
        for ind=delaymax+1:size(y,1)            
            grad(:,ind)=[yf(ind-1:-1:ind-na); -uf(ind-1:-1:ind-nb); -pef(ind-1:-1:ind-nc)];
            pesgrad=pesgrad+perr(ind)*grad(:,ind); %prediction error e(k)* d pe(k)/ d theta            
        end
        
        VNprime=2/(size(y,1)-delaymax)*pesgrad; %Gradient
        
        if strcmp(opt_method,'Gauss-Newton') %Compute Hessian
            Htot=delta*eye(length(theta0),length(theta0));
            for ind=delaymax+1:size(y,1) 
                Htot=Htot+grad(:,ind)*grad(:,ind)';  % used to define the Hessian of VN
            end
            VNHessian=2/(size(y,1)-delaymax)*Htot; %Hessian
            deltatheta=-VNHessian\VNprime; %Search direction
            if -VNprime'*deltatheta<tol  %Stopping criterion for Gauss-Newton method
                break
            end
            
        else
            if norm(VNprime,2)<tol %Stopping criterion for Gradient method
                break
            end
            deltatheta=-VNprime; %Search direction for Gauss-Newton method
        end
          
        
        [t VNopt(inditer)]=linesearch(thetaest,deltatheta,u,y,na,nb,nc,maxgrid); %Line search
        
        thetaest=thetaest+t*deltatheta;
        fprintf('PEM  iteration: %3.0f \n',inditer);
        fprintf('%12s\t  \n','Estimate');
        fprintf('%+9.3f\t  \n',thetaest); 
        
        
        aest=thetaest(1:na);
        best=thetaest(na+1:nb+na);
        cest=thetaest(na+nb+1:nb+na+nc);

 
        
    end
     
modelest.a=aest;
modelest.b{1}=best;
modelest.c=cest;

fprintf('Maximum number of iterations reached \n')

end
 





function [topt VNopt]=linesearch(thetaest,deltatheta,u,y,na,nb,nc,maxgrid)
% [topt VNopt]=linesearch(thetaest,deltatheta,u,y,na,nb,nc,maxgrid) 
% Exact line search for identification of ARMAX through PEM
% Inputs: 
%   thetaest: current estimate of theta
%   deltatheta: search direction
%   u: input data
%   y: output data
%   na,nb,nc: order of ARMAX polynomials A(q), B(q) and C(q)
%   maxgrid: maximum number of gridding points for exact line search  
% Outputs:
%   topt: step size
%   VNopt: Value of the cost function after line search 

vect=linspace(0.0001,1,maxgrid); 
z=tf('z');
delaymax=max(max(na,nb),nc);

    for ind=1:length(vect)
        
        t=vect(ind);
        thetaestp1=thetaest+t*deltatheta;
        
        aest=thetaestp1(1:na);
        best=thetaestp1(na+1:nb+na);
        cest=thetaestp1(na+nb+1:nb+na+nc);
        
 
        %Define filter 1/C;
        numFilterC=[1;zeros(length(cest),1)]';
        denFilterC=[1; cest]';
        FilterC=tf(numFilterC, denFilterC,1); % 1/C(z)
        
        % Define filter A/C;
        FilterAC=tf([1; aest]', denFilterC,1); % A(z)/C(z)
        if na-nc>0
            FilterAC=FilterAC/z^(na-nc);
        end
        
        % Define filter B/C;
        FilterBC=tf([0; best]', denFilterC,1); % B(z)/C(z)
        if nb-nc>0
            FilterBC=FilterBC/z^(nb-nc);
        end
        
        pe=lsim([FilterAC -FilterBC],[y u]); %Prediction error                       
        VN(ind)=norm(pe(delaymax+1:end),2);
        
    end
    
    [VNopt indopt]=min(VN);
    topt=vect(indopt);
         
end




function  [modelestLS Phi thetaLS]=estimateLS(u,y,na,nb) 
% [modelestLS Phi]=estimateLS(u,y,na,nb,delaymax)  LS estimate
% Least-Squares algorithm for MIMO systems with common denominator A(q^(-1)theta)
% Inputs:
    %u: Training input signal 
    %u: Training output signal
    %na: Order of the denominator of the model transfer function (scalar)
    %nb: Order of the numerators of the model transfer functions (matrix) 
% Outputs:
%    modelest: Structure containing the estimated model parameters.
%              modelest.a= estimated parameters of the tf denominator (vector)
%              modelest.b= estimated parameters of the tf numerators (cell)
%    Phi: Regressor matrix     
%    thetaLS: Least-squares parameter estimates in a vector form


delaymax=max(na,max(max(nb))); % Used to handle the initial conditions  
  
[Phi ytot]=BuildRegr(u,y,na,nb); %Construct concated regressor and concatenated output
        
thetaLS=Phi\ytot;
         
modelestLS.a= thetaLS(1:na);

nbvec=reshape(nb',size(nb,1)*size(nb,2),1);         
for indu=1:length(nbvec)        
    modelestLS.b{indu}=thetaLS(na+1+sum(nbvec(1:indu-1)):na+sum(nbvec(1:indu))); 
end %for indu=1:length(nbvec)  
 
end %end function



function [Phi ytot]=BuildRegr(u,y,na,nb)
% [Phi ytot]=BuildRegr(u,y,na,nb)
% Build regressor matrix Phi and extended output 
% Inputs:
    %u: Training input signal 
    %u: Training output signal
    %na: Order of the denominator of the model transfer function (scalar)
    %nb: Order of the numerators of the model transfer functions (matrix) 
% Outputs:
%    Phi: Regressor matrix    
%    y: extended output vector


delaymax=max(na,max(max(nb))); % Used to handle the initial conditions 
ny=size(y,2);
nu=size(u,2);
Phi=zeros(ny*(size(u,1)-delaymax),na+sum(sum(nb)));
ytot=zeros(ny*(size(u,1)-delaymax),1);
N=size(u,1)-delaymax;


for indk=delaymax+1:size(u,1)
     
    for ind=1:ny   
        Phi(N*(ind-1)+indk-delaymax,1:na)=[-y(indk-1:-1:indk-na,ind)'];
         uvec=[];
         for indu=1:nu
            uvec=[uvec; u(indk-1:-1:indk-nb(ind,indu),indu)];
         end         
         Phi(N*(ind-1)+indk-delaymax,na+1+sum(sum(nb(1:ind-1,:))):na+sum(sum(nb(1:ind,:))))=uvec';
    end
end

     for ind=1:ny 
         ytot((ind-1)*N+1:ind*N)=y(delaymax+1:size(u,1),ind);
     end
end % end function





function  ysim=simulateModel(u,ny,modelest)
% ysim=simulateModel(u,y,modelest)
% Simulate model
% Input: 
%   u: input signal
%   ny: number of outputs
%   modelest: estimated parameters
% Output:
%   ysim: Simulated output

Bvec=reshape(modelest.b,size(u,2),ny)';
 
for indy=1:ny
    for indu=1:size(u,2)
        Gest(indy,indu)=tf([0;Bvec{indy,indu}]',[1;modelest.a;zeros(length(Bvec{indy,indu})-length(modelest.a),1)]',1);
    end
end

ysim=lsim(Gest,u);

end



 