diff --git a/scripts/Koopman-ID/.DS_Store b/scripts/Koopman-ID/.DS_Store
new file mode 100644
index 0000000..da67d24
Binary files /dev/null and b/scripts/Koopman-ID/.DS_Store differ
diff --git a/scripts/Koopman-ID/DMD/CompanionMatrix_DMD.m b/scripts/Koopman-ID/DMD/CompanionMatrix_DMD.m
new file mode 100755
index 0000000..9de28e4
--- /dev/null
+++ b/scripts/Koopman-ID/DMD/CompanionMatrix_DMD.m
@@ -0,0 +1,54 @@
+function [ KModes,KEv,Norms ] = CompanionMatrix_DMD( Data )
+% Dynamic Mode Decomposition as presented by 
+% "Spectral analysis of nonlinear flows" by Rowley et al., Journal of FLuid
+% Mechanics, 2009
+
+
+% inputs : 
+% Data - the data matrix: each column is a a set of measurements done at
+% each instant - the sampling frequency is assumed to be constant
+
+% outputs: 
+% 1 - KModes- Koopman or Dynamic Modes: structures that scale exponentially
+% with time - the exponent is given by Koopman or Dynamic Eigenvalues and
+% could be imaginary as well
+
+% 2- KEv - Koopman or Dynamic Eigenvalues: the exponents and frequencies
+% that make up the time-varaiation of data in time
+
+% 3- Norms - Euclidean (vector) norm of each mode, used to sort the data
+
+
+
+
+X=Data(:,1:end-1);
+
+c = pinv(X)*Data(:,end);
+
+m = size(Data,2)-1;
+C = spdiags(ones(m,1),-1,m,m);
+C(:,end)=c;                  %% companion matrix
+
+[P,D]=eig(full(C));                           %% Ritz evalue and evectors
+
+KEunsrtd = diag(D);                     %% Koopman Evalues
+KMunsrtd = X*P;                         %% Koopman Modes
+
+Cnorm = sqrt(sum(abs(KMunsrtd).^2,1));  %% Euclidean norm of Koopman Modes
+[Norms, ind]=sort(Cnorm,'descend');      %% sorting based on the norm
+
+KEv= KEunsrtd(ind);           %% Koopman Eigenvalues
+
+KModes = KMunsrtd(:,ind);                %% Koopman Modes
+
+
+
+end
+
+
+%=========================================================================%
+% Hassan Arbabi - 08-17-2015
+% Mezic research group
+% UC Santa Barbara
+% arbabiha@gmail.com
+%=========================================================================%
\ No newline at end of file
diff --git a/scripts/Koopman-ID/DMD/Exact_DMD.m b/scripts/Koopman-ID/DMD/Exact_DMD.m
new file mode 100755
index 0000000..3c5c7c4
--- /dev/null
+++ b/scripts/Koopman-ID/DMD/Exact_DMD.m
@@ -0,0 +1,71 @@
+function [ ProjectedModes,DEv,ExactModes,Norm ] = Exact_DMD( X,Y,varargin )
+% Dynamic Mode Decomposition as presented by 
+% "On Dynamic Mode Decomposition: theory and applications" by Tu et al.,
+% arXiv, 2013
+
+
+
+% inputs : 
+% Data Sets X and Y- should have the same size 
+% each column of X is a set of measurements done at an instant 
+% each column of Y is the image of the corresponding columns from X
+
+% or 
+
+% ( X,Y,Tol ) - with Tol (optional) being the threshold for filtering thru SVD - the
+% default value is 1e-10
+
+
+
+% outputs: 
+% 1 - Projected Dynamic Modes
+% 2 - Exact Dynamic Modes
+% 3 - Dynamic Eigenvalues
+% 4 - Norms - Euclidean (vector) norm of each mode, used to sort the data
+
+% setting SVD hard threshold
+if isempty(varargin)
+    Tol=1e-10;
+else
+    Tol=varargin{1};
+end
+
+
+disp(['Tolerance used for filtering in DMD:',num2str(Tol)])
+
+[U,S,V]=svd(X,'econ');
+
+k = find(diag(S)>Tol,1,'last');
+disp(['DMD subspace dimension:',num2str(k)])
+
+U = U(:,1:k); V=V(:,1:k); S=S(1:k,1:k);
+
+Atilde = ((U'*Y) *V )* diag((1./diag(S)));
+
+
+[W,DEv]=eig(Atilde);
+DEv = diag(DEv);
+
+ProjectedModes = U*W;
+
+ExactModes = bsxfun(@times,1./DEv.',((Y*V)*S^(-1))*W);
+
+
+
+if nargout>3
+   b = pinv(ProjectedModes)*X(:,end);     % terminal coordinates in the Koopman subspace
+    [Norm,Index]=sort(abs(b),'descend');
+    DEv = DEv(Index);
+    ProjectedModes = ProjectedModes(:,Index);
+    disp('modes sorted based on energy contribution to the last snapshot')
+end
+NORM=Norm
+end
+
+
+%=========================================================================%
+% Hassan Arbabi - 08-17-2015
+% Mezic research group
+% UC Santa Barbara
+% arbabiha@gmail.com
+%=========================================================================%
\ No newline at end of file
diff --git a/scripts/Koopman-ID/DMD/Hankel_DMD.m b/scripts/Koopman-ID/DMD/Hankel_DMD.m
new file mode 100755
index 0000000..710f54e
--- /dev/null
+++ b/scripts/Koopman-ID/DMD/Hankel_DMD.m
@@ -0,0 +1,65 @@
+function [ HModes, HEvalues, Norms ] = Hankel_DMD( Data,n,m,varargin )
+%HANKEL-Dynamic Mode Decompsition as presented in
+% "Ergodic theory, dynamic mode decomposition and computation of Koopman
+% spectral properties" by H. Arbabi and I. Mezic, arXiv:1611.06664
+
+% inputs : 
+% Data - the data matrix: each row is a time series data on an observable
+% m,n: number of columns and rows, respectively, in the data Hankel
+% matrices: size(Data,2)>m+n and preferably m>>n 
+
+% or 
+
+% ( Data,n,m,tol )- with Tol (optional) being the threshold for filtering thru SVD - the
+% default value is 1e-10
+
+
+
+
+% outputs: 
+% 1 - HModes- Dynamic modes which approximate the Koopman eigenfunctions
+
+% 2- HEvalues - Koopman or Dynamic Eigenvalues: the exponents and frequencies
+% that make up the time-varaiation of data in time
+
+% 3- Norms - The L2-norm of Koopman eigenfunction contribution to the
+% observable
+
+
+% setting the SVD hard threshold for DMD
+if isempty(varargin)
+    Tol=1e-10;
+else
+    Tol=varargin{1};
+end
+
+index1 = 1:n;
+index2 = n:n+m-1;
+
+X = []; Y=[];
+
+for ir = 1:size(Data,1)
+   
+    % Hankel blocks ()
+    c = Data(ir,index1).'; r = Data(ir,index2);
+    H = hankel(c,r).';
+    c = Data(ir,index1+1).'; r = Data(ir,index2+1);
+    UH= hankel(c,r).';
+    
+    % scaling of Hankel blocks
+    if ir>1
+        alpha = norm(H(:,1))/norm(X(:,1));
+        H = alpha*H;
+        UH= alpha* UH;
+    end
+    
+    % the data matrices fed to exact DMD
+    X=[X,H]; Y=[Y,UH];
+    
+end
+    % the DMD
+   [HModes,HEvalues,~,Norms ] = DMD.Exact_DMD( X,Y,Tol );
+
+    
+end
+
diff --git a/scripts/Koopman-ID/DMD/HarmonicAverage.m b/scripts/Koopman-ID/DMD/HarmonicAverage.m
new file mode 100755
index 0000000..5c059d9
--- /dev/null
+++ b/scripts/Koopman-ID/DMD/HarmonicAverage.m
@@ -0,0 +1,22 @@
+function [ Average ] = HarmonicAverage( Data,w,t,Filter)
+% computes the harmonic average using the Filter
+
+HarmonicWeight = 2*exp(-1i*w'*t);            
+
+% using the specified filter
+if exist('Filter','var')
+   switch Filter
+       case 'Hamming'
+           HarmonicWeight=bsxfun(@times,HarmonicWeight,hamming(length(t))');
+       case 'Hann'
+           HarmonicWeight=bsxfun(@times,HarmonicWeight,hann(length(t))');
+       case 'exponential'       % the exponential weighting defined by Das & Yorke 2015
+           HarmonicWeight=bsxfun(@times,HarmonicWeight,ExpWeight(length(t))*length(t));
+       otherwise
+           disp('filter not defined ... no filter used')
+   end 
+end
+
+Average = HarmonicWeight*Data'./length(t);     
+
+end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/DMD/SVDenhanced_DMD.m b/scripts/Koopman-ID/DMD/SVDenhanced_DMD.m
new file mode 100755
index 0000000..13ac713
--- /dev/null
+++ b/scripts/Koopman-ID/DMD/SVDenhanced_DMD.m
@@ -0,0 +1,74 @@
+function [ DModes,DEv,Norm ] = SVDenhanced_DMD( Data,varargin )
+% Dynamic Mode Decomposition as presented by 
+% "Dynamic Mode Decomposition of numerical and experimental data" by Peter
+% J. Schmid, Journal of Fluid Mechanics, 2010
+
+% inputs : 
+% Data - the data matrix: each column is a a set of measurements done at
+% each instant - the sampling frequency is assumed to be constant
+
+% or 
+
+% ( Data,Tol ) - with Tol (optional) being the threshold for filtering thru SVD - the
+% default value is 1e-10
+
+% outputs: 
+% 1 - DModes- Koopman or Dynamic Modes: structures that scale exponentially
+% with time - the exponent is given by Koopman or Dynamic Eigenvalues and
+% could be imaginary as well
+
+% 2- DEv - Koopman or Dynamic Eigenvalues: the exponents and frequencies
+% that make up the time-varaiation of data in time
+
+% 3- Norm - The norm here is set to the energy contribution of each mode to
+% the last snapshot of data and then used to sort the modes
+
+
+
+
+tic
+if isempty(varargin)
+    Tol=1e-10;
+else
+    Tol=varargin{1};
+end
+
+disp(['Tolerance used for filtering in DMD:',num2str(Tol)])
+
+X=Data(:,1:end-1);
+Y=Data(:,2:end);
+
+
+[U,S,V]=svd(X,0);
+r = find(diag(S)>Tol,1,'last');
+disp(['DMD subspace dimension:',num2str(r)])
+
+U = U(:,1:r);
+S = S(1:r,1:r);
+V = V(:,1:r);
+
+Atilde = (U'*(Y*V)) / S; 
+
+[w,lambda]=eig(Atilde);
+
+DModes = U*w;
+DEv = diag(lambda);
+
+if nargout>2
+    b = pinv(DModes)*Data(:,end);     % terminal coordinates in the Koopman subspace
+    [Norm,Index]=sort(abs(b),'descend');
+    DEv = DEv(Index);
+    DModes = DModes(:,Index);
+    disp('modes sorted based on energy contribution to the last snapshot')
+end
+
+disp(['time to compute DMD:',num2str(toc)])
+
+end
+
+%=========================================================================%
+% Hassan Arbabi - 08-17-2015
+% Mezic research group
+% UC Santa Barbara
+% arbabiha@gmail.com
+%=========================================================================%
diff --git a/scripts/Koopman-ID/KEEDMD.m b/scripts/Koopman-ID/KEEDMD.m
new file mode 100644
index 0000000..df8fe2e
--- /dev/null
+++ b/scripts/Koopman-ID/KEEDMD.m
@@ -0,0 +1,82 @@
+function [A_koop, B_koop, C_koop, liftFun] = KEEDMD(n,m,Ntraj, Ntime, N_basis,...
+        pow_eig_pairs, Xstr, Ustr, qf, A_nom, B_nom, K_nom, deltaT)
+    %Identify lifted state space model using Extended Dynamic Mode
+    %Decomposition
+    
+    %Inputs:
+    %   n               - Number of states
+    %   m               - Number of control inputs
+    %   N_basis         - Number of basis functions to use
+    %   basis_function  - Type of basis functions (only rbfs implemented)
+    %   rbf-type        - Type of rbf function 
+    %   center_type     - What centers to use for rbfs ('data'/'random') 
+    %   eps             - Width of rbf if rbf type has width parameter
+    %   plot_basis      - (true/false) determines to plot basis functions
+    %   xlim            - Plot limits x-axis
+    %   ylim            - Plot limits y-axis
+    %   Xacc            - Trajectory data
+    %   Yacc            - Trajectory data at next time step
+    %   Uacc            - Control input for trajectory data
+
+    %Outputs:
+    %   A_edmd          - Passive dynamics matrix in lifted space
+    %   B_edmd          - Actuation matrix in lifted space
+    %   C_edmd          - Projection matrix from lifted space to outputs
+
+    % ************************** Prepare data *****************************
+    Xstr_shift = zeros(size(Xstr)); %Shift dynamics such that origin is fixed point
+    X = [];
+    X_dot = [];
+    U = [];
+    for i = 1 : Ntraj
+       Xstr_shift(:,i,:) = Xstr(:,i,:)-qf(:,i);
+       X = [X reshape(Xstr_shift(:,i,:),size(Xstr_shift,1),size(Xstr_shift,3))];
+       X_dot = [X_dot num_diff(reshape(Xstr_shift(:,i,:),size(Xstr_shift,1),size(Xstr_shift,3)),deltaT)];
+       U = [U reshape(Ustr(:,i,:),size(Ustr,1),size(Ustr,3))];
+    end
+    
+    % ******************** Construct eigenfunctions ***********************
+    [A_eigfuncs, liftFun] = construct_eigfuncs(n, N_basis,pow_eig_pairs, ...
+                                            A_nom, B_nom, K_nom, X, X_dot);
+    
+    %Perform lifting
+    liftFun = @(xx) [xx; liftFun(xx)];
+    Nlift = length(liftFun(zeros(n,1)));
+    Xlift = [];
+    Xlift_dot = [];
+    for i = 1 : Ntraj
+       Xlift_temp = liftFun(reshape(Xstr_shift(:,i,:),size(Xstr_shift,1),size(Xstr_shift,3)));
+       Xlift = [Xlift Xlift_temp];
+       Xlift_dot = [Xlift_dot num_diff(Xlift_temp,deltaT)];
+    end
+
+    % ********************** Build predictor *********************************
+
+    %Set up regression for A and B:
+    X_vel = [Xlift; U];
+    Y_vel = Xlift_dot(n/2+1:n,:);
+    A_vel_x = lasso(X_vel',Y_vel(1,:)','Lambda',1e-3, 'Alpha', 0.5);
+    A_vel_th = lasso(X_vel',Y_vel(2,:)','Lambda',1e-3, 'Alpha', 0.5);
+    
+    %Perform regression and enforce known structure:
+    A_koop = zeros(Nlift);
+    A_koop(1:n/2,:) = [zeros(n/2) eye(n/2) zeros(n/2,size(A_koop,2)-n)];
+    A_koop(n/2+1:n,:) = [A_vel_x(1:Nlift,:)'; A_vel_th(1:Nlift,:)']; 
+    A_koop(n+1:end,n+1:end) = A_eigfuncs;
+
+    Y_kin = X_dot(1:n/2,:) - A_koop(1:n/2,:)*Xlift;
+    X_kin = U;
+    B_kin = X_kin'\Y_kin';
+    
+    Y_eig = Xlift_dot(n+1:end,:) - A_eigfuncs*Xlift(n+1:end,:);
+    X_eig = U-K_nom*X;
+    B_eig = X_eig'\Y_eig';
+    
+    B_koop = [B_kin'; A_vel_x(Nlift+1:end,:); A_vel_th(Nlift+1:end,:); B_eig'];  
+    
+    C_koop = zeros(n,size(A_koop,1));
+    C_koop(1:n,1:n) = eye(n);
+    
+    %Subtract effect of nominal controller from A:
+    A_koop(n+1:end,1:n) = A_koop(n+1:end,1:n)-B_eig'*K_nom;
+end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/Resources/.DS_Store b/scripts/Koopman-ID/Resources/.DS_Store
new file mode 100644
index 0000000..22ca612
Binary files /dev/null and b/scripts/Koopman-ID/Resources/.DS_Store differ
diff --git a/scripts/Koopman-ID/Resources/qpOASES-3.1.0/interfaces/matlab/qpOASES.mexmaci64 b/scripts/Koopman-ID/Resources/qpOASES-3.1.0/interfaces/matlab/qpOASES.mexmaci64
old mode 100644
new mode 100755
index 5deca20..15e50b5
Binary files a/scripts/Koopman-ID/Resources/qpOASES-3.1.0/interfaces/matlab/qpOASES.mexmaci64 and b/scripts/Koopman-ID/Resources/qpOASES-3.1.0/interfaces/matlab/qpOASES.mexmaci64 differ
diff --git a/scripts/Koopman-ID/Resources/qpOASES-3.1.0/interfaces/matlab/qpOASES_sequence.mexmaci64 b/scripts/Koopman-ID/Resources/qpOASES-3.1.0/interfaces/matlab/qpOASES_sequence.mexmaci64
old mode 100644
new mode 100755
index bcd7d77..572f213
Binary files a/scripts/Koopman-ID/Resources/qpOASES-3.1.0/interfaces/matlab/qpOASES_sequence.mexmaci64 and b/scripts/Koopman-ID/Resources/qpOASES-3.1.0/interfaces/matlab/qpOASES_sequence.mexmaci64 differ
diff --git a/scripts/Koopman-ID/XY_DDMD_R4.m b/scripts/Koopman-ID/XY_DDMD_R4.m
index 986c3a2..b397ca8 100644
--- a/scripts/Koopman-ID/XY_DDMD_R4.m
+++ b/scripts/Koopman-ID/XY_DDMD_R4.m
@@ -1,391 +1,780 @@
-function [ Z, Lambda, rez, RQ_ref, RSI, Z_ref, rez_ref, U, AxU_k ] = ...

-           XY_DDMD_R4( X, Y, to_scale, to_center, k_mode, tol, nobal, ...

-           n_ref, ref_select, target, overwrite, file_save )

-%       

-% Refined Rayleigh-Ritz procedure for Data Driven Dynamic Mode Decomposition

-% with data driven Residual bounds(DD-DMD-RRRR).

-% This is a sample code that provides algorthmic details of the method

-% described in

-% [1] Z. Drmac, I. Mezic, R. Mohr: Data driven modal decompositions: analysis

-%        and enhancements, SIAM Journal on Scientific Computing.

-%	     (arXiv preprint, https://arxiv.org/abs/1708.02685v1)

-% The code is prepared as a suplement to [1], only to facilitate easier

-% understanding of algorithmic details and experimenting.

-% Released under the 3-Clause BSD License, see below.

-% =====.........................................................................

-% Input

-% =====

-%

-% X (real/complex 2D-array) :: Data matrix. Assumed with more rows than columns. 

-%

-% Y (real/complex 2D-array) ::  Data matrix, Y = A*X for some otherwise

-%                               inaccessible matrix A. Both X and Y should

-%                               be tall matrices

-%                               Separate code (another function, more memory 

-%                               efficient) is used if Y is just shifted X

-%

-% to_scale (string)         ::  Specifies whether snapshot normalization is

-%                               needed. Equilibrating the snapshots in ell_2

-%                               is a good idea in this setting. It does not

-%                               apply if a forgetting factor is incorporated.

-%                           --> If to_scale == 'Scale'

-%                               then X = X*D, Y = Y*D where D is diagonal

-%                               matrix such that the new X has unit columns.

-%

-% to_center (string)        ::  Optional data centering. For experiments only!

-%                           --> If to_center == 'Center'

-%                               then X = X-(1/m)*X*ones(m,1)*ones(1,m) and

-%                               Y = Y-(1/m)*Y*ones(m,1)*ones(1,m). Note that

-%                               after centering, again, Y = A*X.

-%                           --> Set to_center = 'Ignore', except for R&D

-%                               experimenting.

-% 

-% k_mode (integer)          ::  Truncation mode for setting numerical rank k of

-%                               X, based on the singular values sigma(i).

-%                           --> If k_mode == -1, k is selected as

-%                               k = max( i : sigma(i) >= sigma(1)*tol. )

-%                           --> If k_mode == -2, k is selected as

-%                               k = max( i>1 : sigma(i) > sigma(i-1)*tol

-%                           --> If k > 0 and k <= min(m,n)

-%                               the value of k_mode is then understood

-%                               as the caller's request to use k=k_mode

-%                               dominant left singular vectors of X.

-%

-% tol (real, >0 )           ::  Tolerance threshold for truncating the 

-%                               singular values in defining numerical

-%                               rank of scaled data. (Only for the values

-%                               k_mode = -1, k_mode = -2.)

-%                              [!] the recommended value is tol=n*eps,

-%                               where eps is the round-off unit.

-%

-% nobal (string)            ::  Specifies whether to switch of the balancing

-%                               in the function eig() that is used to

-%                               compute the eigenvalues and eigenvectors

-%                               of the Rayleigh quotient.

-%                           --> If nobal == 'NoBalance', the balancing is

-%                               not used.

-%                           --> Set nobal = 'Ignore' if not sure what is

-%                               this about.

-%

-% n_ref (integer)           ::  Specifies how many of the Ritz pairs will be

-%                               refined. If n_ref exceds k, then n_ref is

-%                               set to k (=the total number of Ritz pairs).

-%                              [!!] In case of large dimensions, refining

-%                              too many Ritz pairs is time consuming.

-%

-% ref_select (integer)      ::  Specifies how to select n_ref Ritz pairs

-%                               that will be refined.

-%                           --> If ref_select == 1, 

-%                               refine n_ref pairs with smallest residuals

-%                           --> If ref_select == 2, 

-%                               refine n_ref pairs closest to a given target

-%                               (see the input parameter target)

-%                           --> If ref_select == 3,

-%                               refine n_ref pairs closest to the imaginary axis

-%                           --> If ref_select == 4,

-%                               refine  n_ref pairs closest to the unit

-%                               circle

-%                           --> If ref_select == 5,

-%                               refine n_ref pairs closest to the real axis

-%

-% target (real/complex)     ::  Specifies a point of interest when selecting

-%                               which Ritz pairs to refine (see ref_select)

-%                               Only the Ritz values closest to target will 

-%                               be refined (and their Ritz vectors).  

-%

-% overwrite (string)        ::  Specifes whether the refined Ritz vectors will

-%                               overwrite the original ones, or returned in a

-%                               separate array

-%                           --> If overwrite == 'OverWrite' then the selected

-%                               and improved Ritz vectors overwite the

-%                               original ones. Otherwise, they are returned

-%                               in a separate array. 

-%

-% file_save (string)        ::  Specifies whether the selected variables

-%                               shold be saved in a file.

-%                           --> If file_save == 'NoSave', then no data is saved.

-%                           --> Otherwise, selected variables are saved in

-%                               the file file_save (given on input).

-% ======........................................................................

-% Output  

-% ======........................................................................

-% Z      (real/complex  2D-array) :: Ritz vectors, normalized so that 

-%                                    ||Z(:,i)||_2 = 1. If the refinement is

-%                                    requested with overwrite option, then some

-%                                    of the columns of Z are refined Ritz 

-%                                    vectors. Their indices are then listed in 

-%                                    the integer array RSI.

-%

-% Lambda (real/complex  1D-array) :: Ritz values

-%

-% rez    (real 1D-array)          :: 2-norms of the reziduals

-%                                    rez(i) = ||A*Z(:,i)-Lambda(:,i)*Z(:,i)||_2

-%

-% --> The following four arrays are void if the refiement is not requested.

-%     When calling this function, use tildas if they are not needed.

-%

-% RQ_ref (real/complex 1D-array)  :: Rayleigh quotients with refined Ritz

-%                                    vectors.

-%                                    

-% RSI (integer array)             :: If n_ref > 0, RSI contains the indices

-%                                    of those Ritz pairs that have been refined.

-%

-% Z_ref (real/complex 2D-array)   :: Refined Ritz vectors, in the case when

-%                                    they do not overwrite the orginal ones.

-%

-% rez_ref (real 1D-array)         :: The residuals of the refined Ritz pairs.

-%

-% --> The following two array arguments are returned for post-procesing such as

-%     e.g. refinement. Use tildas if they are not needed.

-%

-% U (real/complex 2-D array)      :: The left singular vectors of the matrix X.

-%

-% AxU_k (real/complex 2-D array)  :: The underlying operator A applied to the 

-%                                    k leading left singular vectors of X.

-%...............................................................................

-%...............................................................................

-% AIMdyn Inc.

-% Coded by Zlatko Drmac, Department of Mathematics, University of Zagreb.

-% drmac@math.hr

-% Version 1.0 ,  January  3, 2017.

-% Version 1.1 , November 11, 2017.

-%...............................................................................

-%...............................................................................

-%

-% The data matrix is n-by-m and it is assumed that the number of its rows 

-% is much larger than the number of its columns.

-%

-[ n, m ] = size( X ) ;

-if m > n 

-   warning('XY_DMD_R4 :(o -> the number of snapshots exceeds space dimension');

-end

-% Optionally, the data can be centered. This is a delicate issue. Only for R&D.

-if ( strcmp(to_center,'Center') )

-    X = X - ((1/m)*(X*ones(m,1))) * ones(1,m) ;

-    Y = Y - ((1/m)*(Y*ones(m,1))) * ones(1,m) ;

-end

-% The intermediate and the final results can be saved for post festum analysis.

-save_to_file = ~strcmp(file_save,'NoSave') ; 

-%

-if ( strcmp(to_scale,'Scale') )

-   % Column norms of the X-data matrix

-   D = ones(m,1) ; 

-   for i = 1 : m

-       D(i) = norm(X(:,i)) ;  

-   end

-   % Remove zero columns (unlikely event, but ..)

-   JX = find( D > 0 ) ; 

-   if ( length(JX) < m )

-      X = X(:,JX)   ; D = D(JX) ; Y = Y(:,JX) ; 

-      m = size(X,2) ;

-   end

-   %

-   % The data samples are scaled to the unit ell_2 sphere and then the SVD is

-   % computed. The numerical rank is determined for the scaled data (uness the 

-   % caller has hardcoded the number of left singular vectors).

-   % One has to be careful here, e.g. if SNR is small. Aditional scaling may

-   % be necessary with carefully chosen weights. If weighting is an issue,

-   % use the weighted version of DMD, see [1].

-   %

-   [ U, Sigma, V ] = svd( X*diag(1./D), 'econ' ) ; S = diag(Sigma) ; 

-else

-   [ U, Sigma, V ] = svd( X, 'econ' ) ; S = diag(Sigma) ; 

-end

-%

-% Determine the numerical rank k of X, based on the k_mode and tol

-if ( k_mode == -1 )

-   % classics

-   if ( tol < eps )

-       warning('XY_DMD_R4 :( -> parameter <tol> is too small')

-   end

-   k = find( S/S(1) > tol, 1, 'last' ) ;     

-elseif ( k_mode == -2 )

-   % more conservative, useful for R&D purposes  

-   if ( tol <= eps )

-       warning('XY_DMD_R4 :( -> parameter <tol> is too small')

-   end

-   %k = find( S(2:m)./S(1:m-1) < tol, 1, 'first' ) ;must be changed, for now use

-   k = find( S(2:m)./S(1:m-1) > tol, 1, 'last' ) ;

-elseif ( (k_mode > 0) && (k_mode <= min(m,n)) )

-   % the caller inputs k as the required number of the left singular

-   % vectors of X

-   k = k_mode ; 

-else

-    error('XY_DMD_R4 :( -> parameter <k_mode> had an illegal value')

-end

-%

-% The image of U(:,1:k) under the action of A

-if ( strcmp(to_scale,'Scale') )

-   AxU_k = Y*(diag(1./D)*( (V(:,1:k).*(ones(m,1)*(1./S(1:k))')))) ; 

-else

-   AxU_k = Y*(((V(:,1:k).*(ones(m,1)*(1./S(1:k))')))) ;

-end

-%

-%

-if ( n_ref > 0 ) 

-   % refinement of Ritz pairs requested 

-   % QR factorization of the matrix [ U(:,1:k), A*U(:,1:k)]

-   %

-     [ ~, R ] = qr( [ U(:,1:k), AxU_k ], 0 ) ;

-   %

-   % The Rayleigh quotient Sk = U_k' * A * U_k, expressed without using A

-   % Sk = (((U(:,1:k)' * Y)*diag(1./D))*V(:,1:k) ).*(ones(k,1)*(1./S(1:k))') ;

-   %

-     Sk = R(1:k,k+1:2*k) ;    % Recall that Sk = U(:,1:k)'*A*U(:,1:k) ;

-   %

-   % Caveat: Here we assume that the QR is computed so that even for complex

-   % matrices the diagonal of R is real, as is the Matlab's function qr().

-   % Otherwise, use linear algebra to determine proper scaling (easy).

-   for i = 1 : k

-       if ( R(i,i) < 0 )

-          Sk(i,:) = -Sk(i,:) ; 

-       end

-   end

-   if ( n_ref > k )

-       warning('XY_DMD_R4 :( -> <n_ref> exceeds the total number of pairs.')

-       n_ref = k ;

-   end

-else

-     % .. no refinement

-     % The Rayleigh quotient Sk = U_k' * A * U_k, expressed without using A

-     % Sk=(((U(:,1:k)' * Y)*diag(1./D))*V(:,1:k)).*(ones(k,1)*(1./S(1:k))')

-     Sk = U(:,1:k)'*AxU_k ; 

-end

-%

-%...............................................................................

-%

-% Ritz values

-%

-if ( strcmp(nobal,'NoBalance'))

-   [W, Lambda] = eig( Sk, 'vector', 'nobalance' ) ;

-else

-   [W, Lambda] = eig( Sk, 'vector' ) ;

-end

-%

-%   

-Z = U(:,1:k) * W  ; % Ritz vectors, normalized in l_2 norm.

-AxZ = AxU_k * W  ; 

-% Next, for each Ritz pair compute its residual.

-rez = zeros(k,1) ; 

-for i = 1 : k 

-    rez(i) = norm( AxZ(:,i) - Lambda(i)*Z(:,i) ) ; 

-end

-%

-if save_to_file

-    save( file_save, 'U','Sigma','V','k','AxU_k','Sk','Z','Lambda','rez' ) ; 

-end

-%...............................................................................

-%

-if ( n_ref > 0 )

-    % refinement requested ; select the specified pairs

-    switch ref_select

-        case 1,

-            % nref pairs with smallest residuals

-            [~, idummy] = sort(rez,'ascend') ;

-            RSI = idummy(1:n_ref) ; 

-        case 2,

-            % nref pairs closest to target (provided on input) 

-            [~,idummy] = sort(abs(Lambda-target) ,'ascend') ;

-            RSI = idummy(1:n_ref) ; 

-        case 3,

-            % nref pairs closest to the imaginary axis

-            [~,idummy] = sort(abs(real(Lambda)) ,'ascend') ;

-            RSI = idummy(1:n_ref) ;

-        case 4,

-            % nref pairs closest to the unit circle

-            [~,idummy] = sort(abs(abs(Lambda)-1) ,'ascend') ;

-            RSI = idummy(1:n_ref) ;

-        case 5,

-            % nref pairs closest to the real axis

-            [~,idummy] = sort(abs(imag(Lambda)) ,'ascend') ;

-            RSI = idummy(1:n_ref) ;

-        % more cases: define property and write a selection procedure 

-        % that extracts the corresponding indicies 

-        otherwise

-           error('XY_DMD_R4 :( -> parameter <ref_select> had an illegal value') 

-    end

-%...   

-%

-VR = zeros(k,n_ref) ;  RQ_ref = zeros(n_ref,1) ; rez_ref = zeros(n_ref,1) ;

-%

-% refinement of the selected pairs

-for ii = 1 : n_ref 

-    i = RSI(ii) ; 

-    RL = [ R(1:k,k+1:2*k) - Lambda(i)*R(1:k,1:k)  ; ...

-           R(k+1:2*k,k+1:2*k) ] ; 

-   [~,~,VL] = svd( RL, 'econ') ;

-% The singular vector of the smallest singular value is the representation 

-% of the refined eigenvector in the basis of U(:,1:k).

-   VR(:,ii) = VL(:,k) ;

-%

-   rez_ref(ii) =  norm(RL*VL(:,k)) ; % = k-th (smallest) singular value of RL

-%  More direct way is to use rez_ref(ii) = SL(k,k) with 

-%   [~,SL,VL] = svd( RL, 'econ') instead of [~,~,VL] = svd( RL, 'econ') above;

-%  This formula is intended for the case of using less accurate SVD method,

-%  when the smallest singular value may be computed inaccurately, but the

-%  corresponding singular vector is computed reasonably accurately. In an

-%  optimal implemenation (not considered here) a specially tailored method

-%  would comput just VL(:,k), without computig the whole SVD of RL.

-% 

-% Use the refined Ritz vectors to compute Rayleigh quotients as better 

-% eigenvalue approximations, with smaller residuals than with Lambda(1:k).

-% In most cases, the improvement is mild, but it is an improvement.

-   RQ_ref(ii)  = VL(:,k)' * Sk * VL(:,k) ; 

-end

-%

-% The matrix of the refined Ritz vectors

-if ( strcmp(overwrite,'OverWrite') )

-   Z(:,RSI) = U(:,1:k)*VR ; Z_ref = [] ;

-   Lambda(RSI) = RQ_ref  ;

-   rez(RSI)    = rez_ref ; 

-else

-   Z_ref = U(:,1:k)*VR ;

-end

-%

-else

-    RSI = [] ; RQ_ref = [] ; rez_ref = [] ; Z_ref = [] ;

-end

-U = U(:,1:k) ; 

-end

-%

-%===============================================================================

-% Copyright (c) 2016-2017 AIMdyn Inc.  

-% All right reserved. 

-% 

-% 3-Clause BSD License

-% 

-% Additional copyrights may follow.

-% 

-% 

-% Redistribution and use in source and binary forms, with or without  

-% modification, are permitted provided that the following conditions are met:

-% 

-% 1. Redistributions of source code must retain the above copyright notice,  

-%    this list of conditions and the following disclaimer.

-% 2. Redistributions in binary form must reproduce the above copyright notice, 

-%    this list of conditions and the following disclaimer in the documentation 

-%    and/or other materials provided with the distribution.

-% 3. Neither the name of the copyright holder nor the names of its contributors 

-%    may be used to endorse or promote products derived from this software  

-%    without specific prior written permission.

-% 

-% The copyright holder provides no reassurances that the source code

-% provided does not infringe any patent, copyright, or any other

-% intellectual property rights of third parties.  The copyright holder

-% disclaims any liability to any recipient for claims brought against

-% recipient by any third party for infringement of that parties

-% intellectual property rights.

-% 

-% THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS "AS IS" AND 

-% ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 

-% WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 

-% DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 

-% FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 

-% DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 

-% SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 

-% CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 

-% OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 

-% OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

-% 

-%===============================================================================

-

+function [ Z, Lambda, rez, RQ_ref, RSI, Z_ref, rez_ref, U, AxU_k ] = ...
+    XY_DDMD_R4( X, Y, to_scale, to_center, k_mode, tol, nobal, ...
+    n_ref, ref_select, target, overwrite, file_save)
+
+%       
+
+% Refined Rayleigh-Ritz procedure for Data Driven Dynamic Mode Decomposition
+
+% with data driven Residual bounds(DD-DMD-RRRR).
+
+% This is a sample code that provides algorthmic details of the method
+
+% described in
+
+% [1] Z. Drmac, I. Mezic, R. Mohr: Data driven modal decompositions: analysis
+
+%        and enhancements, SIAM Journal on Scientific Computing.
+
+%	     (arXiv preprint, https://arxiv.org/abs/1708.02685v1)
+
+% The code is prepared as a suplement to [1], only to facilitate easier
+
+% understanding of algorithmic details and experimenting.
+
+% Released under the 3-Clause BSD License, see below.
+
+% =====.........................................................................
+
+% Input
+
+% =====
+
+%
+
+% X (real/complex 2D-array) :: Data matrix. Assumed with more rows than columns. 
+
+%
+
+% Y (real/complex 2D-array) ::  Data matrix, Y = A*X for some otherwise
+
+%                               inaccessible matrix A. Both X and Y should
+
+%                               be tall matrices
+
+%                               Separate code (another function, more memory 
+
+%                               efficient) is used if Y is just shifted X
+
+%
+
+% to_scale (string)         ::  Specifies whether snapshot normalization is
+
+%                               needed. Equilibrating the snapshots in ell_2
+
+%                               is a good idea in this setting. It does not
+
+%                               apply if a forgetting factor is incorporated.
+
+%                           --> If to_scale == 'Scale'
+
+%                               then X = X*D, Y = Y*D where D is diagonal
+
+%                               matrix such that the new X has unit columns.
+
+%
+
+% to_center (string)        ::  Optional data centering. For experiments only!
+
+%                           --> If to_center == 'Center'
+
+%                               then X = X-(1/m)*X*ones(m,1)*ones(1,m) and
+
+%                               Y = Y-(1/m)*Y*ones(m,1)*ones(1,m). Note that
+
+%                               after centering, again, Y = A*X.
+
+%                           --> Set to_center = 'Ignore', except for R&D
+
+%                               experimenting.
+
+% 
+
+% k_mode (integer)          ::  Truncation mode for setting numerical rank k of
+
+%                               X, based on the singular values sigma(i).
+
+%                           --> If k_mode == -1, k is selected as
+
+%                               k = max( i : sigma(i) >= sigma(1)*tol. )
+
+%                           --> If k_mode == -2, k is selected as
+
+%                               k = max( i>1 : sigma(i) > sigma(i-1)*tol
+
+%                           --> If k > 0 and k <= min(m,n)
+
+%                               the value of k_mode is then understood
+
+%                               as the caller's request to use k=k_mode
+
+%                               dominant left singular vectors of X.
+
+%
+
+% tol (real, >0 )           ::  Tolerance threshold for truncating the 
+
+%                               singular values in defining numerical
+
+%                               rank of scaled data. (Only for the values
+
+%                               k_mode = -1, k_mode = -2.)
+
+%                              [!] the recommended value is tol=n*eps,
+
+%                               where eps is the round-off unit.
+
+%
+
+% nobal (string)            ::  Specifies whether to switch of the balancing
+
+%                               in the function eig() that is used to
+
+%                               compute the eigenvalues and eigenvectors
+
+%                               of the Rayleigh quotient.
+
+%                           --> If nobal == 'NoBalance', the balancing is
+
+%                               not used.
+
+%                           --> Set nobal = 'Ignore' if not sure what is
+
+%                               this about.
+
+%
+
+% n_ref (integer)           ::  Specifies how many of the Ritz pairs will be
+
+%                               refined. If n_ref exceds k, then n_ref is
+
+%                               set to k (=the total number of Ritz pairs).
+
+%                              [!!] In case of large dimensions, refining
+
+%                              too many Ritz pairs is time consuming.
+
+%
+
+% ref_select (integer)      ::  Specifies how to select n_ref Ritz pairs
+
+%                               that will be refined.
+
+%                           --> If ref_select == 1, 
+
+%                               refine n_ref pairs with smallest residuals
+
+%                           --> If ref_select == 2, 
+
+%                               refine n_ref pairs closest to a given target
+
+%                               (see the input parameter target)
+
+%                           --> If ref_select == 3,
+
+%                               refine n_ref pairs closest to the imaginary axis
+
+%                           --> If ref_select == 4,
+
+%                               refine  n_ref pairs closest to the unit
+
+%                               circle
+
+%                           --> If ref_select == 5,
+
+%                               refine n_ref pairs closest to the real axis
+
+%
+
+% target (real/complex)     ::  Specifies a point of interest when selecting
+
+%                               which Ritz pairs to refine (see ref_select)
+
+%                               Only the Ritz values closest to target will 
+
+%                               be refined (and their Ritz vectors).  
+
+%
+
+% overwrite (string)        ::  Specifes whether the refined Ritz vectors will
+
+%                               overwrite the original ones, or returned in a
+
+%                               separate array
+
+%                           --> If overwrite == 'OverWrite' then the selected
+
+%                               and improved Ritz vectors overwite the
+
+%                               original ones. Otherwise, they are returned
+
+%                               in a separate array. 
+
+%
+
+% file_save (string)        ::  Specifies whether the selected variables
+
+%                               shold be saved in a file.
+
+%                           --> If file_save == 'NoSave', then no data is saved.
+
+%                           --> Otherwise, selected variables are saved in
+
+%                               the file file_save (given on input).
+
+% ======........................................................................
+
+% Output  
+
+% ======........................................................................
+
+% Z      (real/complex  2D-array) :: Ritz vectors, normalized so that 
+
+%                                    ||Z(:,i)||_2 = 1. If the refinement is
+
+%                                    requested with overwrite option, then some
+
+%                                    of the columns of Z are refined Ritz 
+
+%                                    vectors. Their indices are then listed in 
+
+%                                    the integer array RSI.
+
+%
+
+% Lambda (real/complex  1D-array) :: Ritz values
+
+%
+
+% rez    (real 1D-array)          :: 2-norms of the reziduals
+
+%                                    rez(i) = ||A*Z(:,i)-Lambda(:,i)*Z(:,i)||_2
+
+%
+
+% --> The following four arrays are void if the refiement is not requested.
+
+%     When calling this function, use tildas if they are not needed.
+
+%
+
+% RQ_ref (real/complex 1D-array)  :: Rayleigh quotients with refined Ritz
+
+%                                    vectors.
+
+%                                    
+
+% RSI (integer array)             :: If n_ref > 0, RSI contains the indices
+
+%                                    of those Ritz pairs that have been refined.
+
+%
+
+% Z_ref (real/complex 2D-array)   :: Refined Ritz vectors, in the case when
+
+%                                    they do not overwrite the orginal ones.
+
+%
+
+% rez_ref (real 1D-array)         :: The residuals of the refined Ritz pairs.
+
+%
+
+% --> The following two array arguments are returned for post-procesing such as
+
+%     e.g. refinement. Use tildas if they are not needed.
+
+%
+
+% U (real/complex 2-D array)      :: The left singular vectors of the matrix X.
+
+%
+
+% AxU_k (real/complex 2-D array)  :: The underlying operator A applied to the 
+
+%                                    k leading left singular vectors of X.
+
+%...............................................................................
+
+%...............................................................................
+
+% AIMdyn Inc.
+
+% Coded by Zlatko Drmac, Department of Mathematics, University of Zagreb.
+
+% drmac@math.hr
+
+% Version 1.0 ,  January  3, 2017.
+
+% Version 1.1 , November 11, 2017.
+
+%...............................................................................
+
+%...............................................................................
+
+%
+
+% The data matrix is n-by-m and it is assumed that the number of its rows 
+
+% is much larger than the number of its columns.
+
+%
+
+[ n, m ] = size( X ) ;
+
+if m > n 
+
+   warning('XY_DMD_R4 :(o -> the number of snapshots exceeds space dimension');
+
+end
+
+% Optionally, the data can be centered. This is a delicate issue. Only for R&D.
+
+if ( strcmp(to_center,'Center') )
+
+    X = X - ((1/m)*(X*ones(m,1))) * ones(1,m) ;
+
+    Y = Y - ((1/m)*(Y*ones(m,1))) * ones(1,m) ;
+
+end
+
+% The intermediate and the final results can be saved for post festum analysis.
+
+save_to_file = ~strcmp(file_save,'NoSave') ; 
+
+%
+
+if ( strcmp(to_scale,'Scale') )
+
+   % Column norms of the X-data matrix
+
+   D = ones(m,1) ; 
+
+   for i = 1 : m
+
+       D(i) = norm(X(:,i)) ;  
+
+   end
+
+   % Remove zero columns (unlikely event, but ..)
+
+   JX = find( D > 0 ) ; 
+
+   if ( length(JX) < m )
+
+      X = X(:,JX)   ; D = D(JX) ; Y = Y(:,JX) ; 
+
+      m = size(X,2) ;
+
+   end
+
+   %
+
+   % The data samples are scaled to the unit ell_2 sphere and then the SVD is
+
+   % computed. The numerical rank is determined for the scaled data (uness the 
+
+   % caller has hardcoded the number of left singular vectors).
+
+   % One has to be careful here, e.g. if SNR is small. Aditional scaling may
+
+   % be necessary with carefully chosen weights. If weighting is an issue,
+
+   % use the weighted version of DMD, see [1].
+
+   %
+
+   [ U, Sigma, V ] = svd( X*diag(1./D), 'econ' ) ; S = diag(Sigma) ; 
+
+else
+
+   [ U, Sigma, V ] = svd( X, 'econ' ) ; S = diag(Sigma) ; 
+
+end
+
+%
+
+% Determine the numerical rank k of X, based on the k_mode and tol
+
+if ( k_mode == -1 )
+
+   % classics
+
+   if ( tol < eps )
+
+       warning('XY_DMD_R4 :( -> parameter <tol> is too small')
+
+   end
+
+   k = find( S/S(1) > tol, 1, 'last' ) ;     
+
+elseif ( k_mode == -2 )
+
+   % more conservative, useful for R&D purposes  
+
+   if ( tol <= eps )
+
+       warning('XY_DMD_R4 :( -> parameter <tol> is too small')
+
+   end
+
+   %k = find( S(2:m)./S(1:m-1) < tol, 1, 'first' ) ;must be changed, for now use
+
+   k = find( S(2:m)./S(1:m-1) > tol, 1, 'last' ) ;
+
+elseif ( (k_mode > 0) && (k_mode <= min(m,n)) )
+
+   % the caller inputs k as the required number of the left singular
+
+   % vectors of X
+
+   k = k_mode ; 
+
+else
+
+    error('XY_DMD_R4 :( -> parameter <k_mode> had an illegal value')
+
+end
+
+%
+
+% The image of U(:,1:k) under the action of A
+
+if ( strcmp(to_scale,'Scale') )
+
+   AxU_k = Y*(diag(1./D)*( (V(:,1:k).*(ones(m,1)*(1./S(1:k))')))) ; 
+
+else
+
+   AxU_k = Y*(((V(:,1:k).*(ones(m,1)*(1./S(1:k))')))) ;
+
+end
+
+%
+
+%
+
+if ( n_ref > 0 ) 
+
+   % refinement of Ritz pairs requested 
+
+   % QR factorization of the matrix [ U(:,1:k), A*U(:,1:k)]
+
+   %
+
+     [ ~, R ] = qr( [ U(:,1:k), AxU_k ], 0 ) ;
+
+   %
+
+   % The Rayleigh quotient Sk = U_k' * A * U_k, expressed without using A
+
+   % Sk = (((U(:,1:k)' * Y)*diag(1./D))*V(:,1:k) ).*(ones(k,1)*(1./S(1:k))') ;
+
+   %
+
+     Sk = R(1:k,k+1:2*k) ;    % Recall that Sk = U(:,1:k)'*A*U(:,1:k) ;
+
+   %
+
+   % Caveat: Here we assume that the QR is computed so that even for complex
+
+   % matrices the diagonal of R is real, as is the Matlab's function qr().
+
+   % Otherwise, use linear algebra to determine proper scaling (easy).
+
+   for i = 1 : k
+
+       if ( R(i,i) < 0 )
+
+          Sk(i,:) = -Sk(i,:) ; 
+
+       end
+
+   end
+
+   if ( n_ref > k )
+
+       warning('XY_DMD_R4 :( -> <n_ref> exceeds the total number of pairs.')
+
+       n_ref = k ;
+
+   end
+
+else
+
+     % .. no refinement
+
+     % The Rayleigh quotient Sk = U_k' * A * U_k, expressed without using A
+
+     % Sk=(((U(:,1:k)' * Y)*diag(1./D))*V(:,1:k)).*(ones(k,1)*(1./S(1:k))')
+
+     Sk = U(:,1:k)'*AxU_k ; 
+
+end
+
+%
+
+%...............................................................................
+
+%
+
+% Ritz values
+
+%
+
+if ( strcmp(nobal,'NoBalance'))
+
+   [W, Lambda] = eig( Sk, 'vector', 'nobalance' ) ;
+
+else
+
+   [W, Lambda] = eig( Sk, 'vector' ) ;
+
+end
+
+%
+
+%   
+
+Z = U(:,1:k) * W  ; % Ritz vectors, normalized in l_2 norm.
+
+AxZ = AxU_k * W  ; 
+
+% Next, for each Ritz pair compute its residual.
+
+rez = zeros(k,1) ; 
+
+for i = 1 : k 
+
+    rez(i) = norm( AxZ(:,i) - Lambda(i)*Z(:,i) ) ; 
+
+end
+
+%
+
+if save_to_file
+
+    save( file_save, 'U','Sigma','V','k','AxU_k','Sk','Z','Lambda','rez' ) ; 
+
+end
+
+%...............................................................................
+
+%
+
+if ( n_ref > 0 )
+
+    % refinement requested ; select the specified pairs
+
+    switch ref_select
+
+        case 1,
+
+            % nref pairs with smallest residuals
+
+            [~, idummy] = sort(rez,'ascend') ;
+
+            RSI = idummy(1:n_ref) ; 
+
+        case 2,
+
+            % nref pairs closest to target (provided on input) 
+
+            [~,idummy] = sort(abs(Lambda-target) ,'ascend') ;
+
+            RSI = idummy(1:n_ref) ; 
+
+        case 3,
+
+            % nref pairs closest to the imaginary axis
+
+            [~,idummy] = sort(abs(real(Lambda)) ,'ascend') ;
+
+            RSI = idummy(1:n_ref) ;
+
+        case 4,
+
+            % nref pairs closest to the unit circle
+
+            [~,idummy] = sort(abs(abs(Lambda)-1) ,'ascend') ;
+
+            RSI = idummy(1:n_ref) ;
+
+        case 5,
+
+            % nref pairs closest to the real axis
+
+            [~,idummy] = sort(abs(imag(Lambda)) ,'ascend') ;
+
+            RSI = idummy(1:n_ref) ;
+
+        % more cases: define property and write a selection procedure 
+
+        % that extracts the corresponding indicies 
+
+        otherwise
+
+           error('XY_DMD_R4 :( -> parameter <ref_select> had an illegal value') 
+
+    end
+
+%...   
+
+%
+
+VR = zeros(k,n_ref) ;  RQ_ref = zeros(n_ref,1) ; rez_ref = zeros(n_ref,1) ;
+
+%
+
+% refinement of the selected pairs
+
+for ii = 1 : n_ref 
+
+    i = RSI(ii) ; 
+
+    RL = [ R(1:k,k+1:2*k) - Lambda(i)*R(1:k,1:k)  ; ...
+
+           R(k+1:2*k,k+1:2*k) ] ; 
+
+   [~,~,VL] = svd( RL, 'econ') ;
+
+% The singular vector of the smallest singular value is the representation 
+
+% of the refined eigenvector in the basis of U(:,1:k).
+
+   VR(:,ii) = VL(:,k) ;
+
+%
+
+   rez_ref(ii) =  norm(RL*VL(:,k)) ; % = k-th (smallest) singular value of RL
+
+%  More direct way is to use rez_ref(ii) = SL(k,k) with 
+
+%   [~,SL,VL] = svd( RL, 'econ') instead of [~,~,VL] = svd( RL, 'econ') above;
+
+%  This formula is intended for the case of using less accurate SVD method,
+
+%  when the smallest singular value may be computed inaccurately, but the
+
+%  corresponding singular vector is computed reasonably accurately. In an
+
+%  optimal implemenation (not considered here) a specially tailored method
+
+%  would comput just VL(:,k), without computig the whole SVD of RL.
+
+% 
+
+% Use the refined Ritz vectors to compute Rayleigh quotients as better 
+
+% eigenvalue approximations, with smaller residuals than with Lambda(1:k).
+
+% In most cases, the improvement is mild, but it is an improvement.
+
+   RQ_ref(ii)  = VL(:,k)' * Sk * VL(:,k) ; 
+
+end
+
+%
+
+% The matrix of the refined Ritz vectors
+
+if ( strcmp(overwrite,'OverWrite') )
+
+   Z(:,RSI) = U(:,1:k)*VR ; Z_ref = [] ;
+
+   Lambda(RSI) = RQ_ref  ;
+
+   rez(RSI)    = rez_ref ; 
+
+else
+
+   Z_ref = U(:,1:k)*VR ;
+
+end
+
+%
+
+else
+
+    RSI = [] ; RQ_ref = [] ; rez_ref = [] ; Z_ref = [] ;
+
+end
+
+U = U(:,1:k) ; 
+
+end
+
+%
+
+%===============================================================================
+
+% Copyright (c) 2016-2017 AIMdyn Inc.  
+
+% All right reserved. 
+
+% 
+
+% 3-Clause BSD License
+
+% 
+
+% Additional copyrights may follow.
+
+% 
+
+% 
+
+% Redistribution and use in source and binary forms, with or without  
+
+% modification, are permitted provided that the following conditions are met:
+
+% 
+
+% 1. Redistributions of source code must retain the above copyright notice,  
+
+%    this list of conditions and the following disclaimer.
+
+% 2. Redistributions in binary form must reproduce the above copyright notice, 
+
+%    this list of conditions and the following disclaimer in the documentation 
+
+%    and/or other materials provided with the distribution.
+
+% 3. Neither the name of the copyright holder nor the names of its contributors 
+
+%    may be used to endorse or promote products derived from this software  
+
+%    without specific prior written permission.
+
+% 
+
+% The copyright holder provides no reassurances that the source code
+
+% provided does not infringe any patent, copyright, or any other
+
+% intellectual property rights of third parties.  The copyright holder
+
+% disclaims any liability to any recipient for claims brought against
+
+% recipient by any third party for infringement of that parties
+
+% intellectual property rights.
+
+% 
+
+% THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS "AS IS" AND 
+
+% ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 
+
+% WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 
+
+% DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 
+
+% FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 
+
+% DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 
+
+% SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 
+
+% CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 
+
+% OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 
+
+% OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+% 
+
+%===============================================================================
+
+
+
diff --git a/scripts/Koopman-ID/construct_eigfuncs.m b/scripts/Koopman-ID/construct_eigfuncs.m
new file mode 100644
index 0000000..db3d8ef
--- /dev/null
+++ b/scripts/Koopman-ID/construct_eigfuncs.m
@@ -0,0 +1,103 @@
+function [A_koop, z_eigfun] = construct_eigfuncs(n, N_basis, pow_eig_pairs,...
+                                                A_nom, B_nom, K_nom, X, X_dot)
+
+    %Identify lifted state space model using approximate Koopman invariant
+    %subspace
+
+    %Inputs:
+    %   n               - Number of states
+    %   m               - Number of control inputs
+    %   N_lambda        - Number of eigenvalue candidates to use
+    %   N_basis         - Number of boundary functions to use
+    %   basis_function  - Type of basis functions (only rbfs implemented)
+    %   rbf-type        - Type of rbf function 
+    %   center_type     - What centers to use for rbfs ('data'/'random') 
+    %   eps             - Width of rbf if rbf type has width parameter
+    %   plot_basis      - (true/false) determines to plot basis functions
+    %   xlim            - Plot limits x-axis
+    %   ylim            - Plot limits y-axis
+    %   A_edmd          - Passive dynamics matrix from EDMD
+    %   A_nom           - Nominal state space model passive dynamics 
+    %   B_nom           - Nominal state space model actuated dynamics
+    %   K_nom           - Nominal linear feedback control law gains
+    %   Xacc            - Trajectory data
+    %   Yacc            - Trajectory data at next time step
+    %   Uacc            - Control input for trajectory data
+    %   Xacc_c          - Trajectory data, perturbed control
+    %   Yacc_c          - Trajectory data at next time step, perturbed control
+    %   Uacc_c          - Control input for trajectory data, perturbed control
+
+    %Outputs:
+    %   A_edmd          - Passive dynamics matrix in lifted space
+    %   B_edmd          - Actuation matrix in lifted space
+    %   C_edmd          - Projection matrix from lifted space to output
+
+    
+    % Generate eigenfunctions and learn autonomous dynamics:
+    
+    disp('Starting autonomous dynamics learning...'); tic
+    
+    A_cl = A_nom + B_nom*K_nom;
+    cent = [5*rand(1,N_basis)-2.5;...
+        2*pi/3*rand(1,N_basis)-pi/3;...
+        4*rand(1,N_basis)-2;...
+        4*rand(1,N_basis)-2];
+    rbf_type = 'gauss';
+    eps_rbf = 1;
+        
+    % Set up nonlinear transformations and their gradients
+    zfun = @(xx) [xx; rbf(xx,cent,rbf_type,eps_rbf)];
+    zfun_grad = @(xx) [eye(n); rbf_grad(xx,cent,rbf_type,eps_rbf)];
+    zfun_dot = @(xx, xx_dot) [xx_dot; rbf_dot(xx,xx_dot,cent,rbf_type,eps_rbf)];
+    
+    % Set up Z and Z_dot matrices:
+    Z = zfun(X);
+    Z_dot = zfun_dot(X,X_dot);
+
+    %Set up constraint matrix
+    con1 = zfun_grad(zeros(n,1));
+
+    disp('Solving optimization problem...')
+    N = size(Z,1);
+    cvx_begin 
+        variable C(n,N);
+        minimize (norm(X_dot + C*Z_dot - A_cl*(X+C*Z),'fro') + 1e-2*norm(C,'fro')/N*n)
+        subject to
+            {C*con1 == zeros(n)}; 
+    cvx_end
+    fprintf('Solved, optimal value excluding regularization: %.10f, MSE: %.10f\n', norm(X_dot + C*Z_dot - A_cl*(X+C*Z),'fro'),immse(X_dot+C*Z_dot, A_cl*(X+C*Z)))
+    fprintf('Constraint violation: %.4f \n', sum(sum(abs(C*con1))))
+
+    yfun = @(xx) xx + [zeros(1,size(xx,2));...
+       C(2,:)*zfun(xx);...
+       zeros(1,size(xx,2));...
+       C(4,:)*zfun(xx)]; % Full learned diffeomorphism
+%    yfun = @(xx) xx + C*zfun(xx);
+    
+    
+    % Calculate eigenfunctions for linearized system
+    [V_a,D] = eig(A_cl);
+    [W_a,~] = eig(A_cl');
+
+    %Define powers (only implemented for n=2):
+    a = 0 : pow_eig_pairs;
+    combs = combnk(a,n);
+    powers = [];
+    for i = 1 : size(combs,1)
+        powers = [powers; perms(combs(i,:))];
+    end
+
+    linfunc = @(xx) (xx'*W_a)';%./diag(V_a'*W_a);
+    phifun = @(xx) (prod(linfunc(xx).^(powers')))';
+    lambd = prod(exp(diag(D)).^(powers'))';
+    lambd = log(lambd);
+
+    % Construct scaling function
+    gfun = @(xx) xx./[5 4 2*pi/3 4]'; %Scale state space into unit cube
+
+    % Construct eigenfunctions for nonlinear system
+    z_eigfun = @(xx) phifun_mat(phifun, gfun(yfun(xx))); 
+    A_koop = diag(lambd);
+end
+
+
diff --git a/scripts/Koopman-ID/generate_trajectory.m b/scripts/Koopman-ID/generate_trajectory.m
new file mode 100644
index 0000000..11ef673
--- /dev/null
+++ b/scripts/Koopman-ID/generate_trajectory.m
@@ -0,0 +1,37 @@
+function [q_t, u_t] = generate_trajectory(A,B,C,Qf,q0,qf,t,deltaT,lb,ub)
+
+% Extract relevant parameters
+Nsim = length(t);
+Ntraj = size(q0,2);
+N = size(A,1);
+n = size(C,1);
+m = size(B,2);
+
+% Discretize state space matrices
+continuous_model_nom = ss(A, B, C, 0);
+discrete_model_nom = c2d(continuous_model_nom,deltaT);
+A_d = discrete_model_nom.A;
+B_d = discrete_model_nom.B;
+C_d = discrete_model_nom.C;
+
+% Initialize variables
+q_t = zeros(n,Ntraj,Nsim);
+u_t = zeros(m,Ntraj,Nsim-1);
+% Generate trajectories for each initial-final condition pair
+for i = 1 : size(q0,2)
+    cvx_begin quiet
+        variable u(m,Nsim-1)
+        variable z(N,Nsim)
+        minimize ((C_d*z(:,Nsim)-qf(:,i))'*Qf*(C_d*z(:,Nsim)-qf(:,i)) + norm(u,'fro'));
+        subject to
+            {z(:,2:Nsim) == A_d*z(:,1:Nsim-1) + B_d*u};
+            {C_d*z(:,1) == q0(:,i)};
+            {z <= ub.*ones(N,Nsim)};
+            {z >= lb.*ones(N,Nsim)};
+    cvx_end;
+
+    q_t(:,i,:) = C*z; %Project back to state space if trajectory is based on lifted state
+    u_t(:,i,:) = u;
+end
+
+
diff --git a/scripts/Koopman-ID/koopman_eigen_id_laplace.m b/scripts/Koopman-ID/koopman_eigen_id_laplace.m
new file mode 100644
index 0000000..c0aad29
--- /dev/null
+++ b/scripts/Koopman-ID/koopman_eigen_id_laplace.m
@@ -0,0 +1,146 @@
+function [A_koop, B_koop, C_koop, phi_fun_v] = koopman_eigen_id_laplace(n,m, Ntraj, Ntime, ...
+    N_basis,basis_function, rbf_type, center_type, eps_rbf, N_lambda, ...
+    lambda_type, A_edmd, A_nom, B_nom, K_nom, Xacc, Yacc, Xstr, Xacc_c,...
+    Yacc_c, Uacc_c, Xstr_c, Ustr_c, timestr, deltaT, learn_type)
+
+    %Identify lifted state space model using approximate Koopman invariant
+    %subspace
+
+    %Inputs:
+    %   n               - Number of states
+    %   m               - Number of control inputs
+    %   N_lambda        - Number of eigenvalue candidates to use
+    %   N_basis         - Number of boundary functions to use
+    %   basis_function  - Type of basis functions (only rbfs implemented)
+    %   rbf-type        - Type of rbf function 
+    %   center_type     - What centers to use for rbfs ('data'/'random') 
+    %   eps             - Width of rbf if rbf type has width parameter
+    %   plot_basis      - (true/false) determines to plot basis functions
+    %   xlim            - Plot limits x-axis
+    %   ylim            - Plot limits y-axis
+    %   A_edmd          - Passive dynamics matrix from EDMD
+    %   A_nom           - Nominal state space model passive dynamics 
+    %   B_nom           - Nominal state space model actuated dynamics
+    %   K_nom           - Nominal linear feedback control law gains
+    %   Xacc            - Trajectory data
+    %   Yacc            - Trajectory data at next time step
+    %   Uacc            - Control input for trajectory data
+    %   Xacc_c          - Trajectory data, perturbed control
+    %   Yacc_c          - Trajectory data at next time step, perturbed control
+    %   Uacc_c          - Control input for trajectory data, perturbed control
+
+    %Outputs:
+    %   A_edmd          - Passive dynamics matrix in lifted space
+    %   B_edmd          - Actuation matrix in lifted space
+    %   C_edmd          - Projection matrix from lifted space to output
+
+    
+    % Generate eigenfunctions and learn autonomous dynamics:
+    
+    disp('Starting autonomous dynamics learning...'); tic
+    
+    A_c = A_nom + B_nom*K_nom;
+    A_c_d = expm(A_c*deltaT); %Discrete time dynamics matrix
+    cent = rand(n,50)*2*pi - pi;
+    rbf_type = 'gauss';
+    eps_rbf = 1;
+    
+    % Set up nonlinear transformations and their gradients
+    zfun = @(xx) [xx'; rbf(xx',cent,rbf_type,eps_rbf)];
+    zfun_grad = @(xx) [eye(2); rbf_grad(xx',cent,rbf_type,eps_rbf)];
+    
+    % Prepare data matrices
+    Y = [];
+    Z_mplus1 = [];
+    Z_m = [];
+    for i = 1 : Ntraj
+       X_mplus1 = reshape(Xstr(:,i,2:end),n,Ntime);
+       X_m = reshape(Xstr(:,i,1:end-1),n,Ntime);
+
+       % Set up Y matrix (targets), Y(:,i) = x_(i+1) - A_nom*x_i
+       Y = [Y X_mplus1-A_c_d*X_m]; 
+
+       % Set up Z_mplus1 matrix (inputs), phi(x_(i+1))
+       Z_mplus1 = [Z_mplus1 zfun(X_mplus1')];
+
+       % Set up Z_m matrix (inputs), phi(x_i)
+       Z_m = [Z_m zfun(X_m')]; 
+    end
+
+    %Set up constraint matrix
+    con1 = zfun_grad([0 0]);
+
+    disp('Solving optimization problem...')
+    N = size(Z_m,1);
+    cvx_begin
+        variable C(n,N);
+        minimize (norm(Y - (A_c_d*C*Z_m - C*Z_mplus1),'fro') + 1*norm(C,'fro'))
+        subject to
+            {C*con1 == zeros(n)}; 
+    cvx_end
+    fprintf('Solved, optimal value excluding regularization: %.10f, MSE: %.10f\n', norm(Y - (A_c_d*C*Z_m - C*Z_mplus1),'fro'),immse(Y, A_c_d*C*Z_m - C*Z_mplus1))
+    %fprintf('Constraint violation: %.4f \n', sum(sum(abs(C*con1))))
+
+    yfun = @(xx) xx + C*zfun(xx'); % Full learned diffeomorphism
+    
+    % Calculate eigenfunctions for linearized system
+    [~,D] = eig(A_c_d);
+    [V_a,~] = eig(A_c_d');
+
+    %Define powers (only implemented for n=2):
+    max_power = 5;
+    a = 0 : max_power;
+    [P,Q] = meshgrid(a,a);
+    c=cat(2,P',Q');
+    powers=reshape(c,[],2);
+
+    linfunc = @(xx) (xx'*V_a)';
+    phifun = @(xx) (prod(linfunc(xx).^(powers')))';
+    lambd = prod(diag(D).^(powers'))';
+
+    % Construct scaling function
+    gfun = @(xx) xx./pi; %Scale state space into unit cube
+
+    % Construct eigenfunctions for nonlinear system
+    z_eigfun = @(xx) [xx; phifun_mat(phifun, gfun(yfun(xx)))];
+    %z_eigfun = @(xx) phifun_mat(phifun, gfun(yfun(xx)));
+    
+    %A_koop = diag(lambd);
+    Z_x = z_eigfun(Xacc);
+    Z_y = z_eigfun(Yacc);
+    k=1;
+    A_koop = [eye(n/2) deltaT*eye(n/2) zeros(n/2,size(Z_x,1)-n)];
+    A_koop = [A_koop; (Z_x'\Z_y(n/2+1:n,:)')'];
+    %A_state = lasso(Z_x',Z_y(n/2+1:n,:)','Lambda',1e-4, 'Alpha', 0.8)';
+    %A_koop = [A_koop; A_state];
+    A_koop = [A_koop; [zeros(length(lambd),n) diag(lambd)]];
+
+    % Calculate projection of states onto eigenfunctions 
+    %C_koop = (Zacc'\Xacc')';
+    C_koop = zeros(n,size(A_koop,1));
+    C_koop(:,1:n) = eye(n);
+    fprintf('Learning autonomous dynamics done, training error %.2f, execution time: %1.2f s \n', immse(Yacc,C_koop*A_koop*Z_x), toc);
+
+    % Learn actuated dynamics:    
+    disp('Starting actuated dynamics learning...'); tic
+    [B_koop, phi_fun_v] = learn_B_koop(n, m, Ntraj, Ntime, z_eigfun, A_koop,...
+                    C_koop, Xstr_c, Ustr_c, Xacc_c, Yacc_c, Uacc_c, deltaT, learn_type);
+    fprintf('Learning actuated dynamics done, execution time: %1.2f s \n', toc);
+    
+    % Subtract effect of nominal controller from A_koop:
+    A_koop(:,1:n) = A_koop(:,1:n)-B_koop*K_nom;
+     
+    % Evaluate training error:
+%     Ypred_c = zeros(size(Xstr_c(:,:,2:end)));
+%     Zpred = phi_fun_v(Xstr_c(:,:,1));
+%     err = 0;
+%     for i = 1:Ntime
+%         Zpred = A_koop*Zpred+B_koop*Ustr_c(:,:,i);
+%         Ypred_c(:,:,i) = C_koop*Zpred;
+%         err = err + sum(sum((Ypred_c(:,:,i) - Xstr_c(:,:,i+1)).^2));
+%     end
+%     err = err/(Ntime*Ntraj);
+%     disp(err);
+end
+
+
diff --git a/scripts/Koopman-ID/main_comparison_pendulum.m b/scripts/Koopman-ID/main_comparison_pendulum.m
index a65692b..7200fe4 100644
--- a/scripts/Koopman-ID/main_comparison_pendulum.m
+++ b/scripts/Koopman-ID/main_comparison_pendulum.m
@@ -14,102 +14,114 @@
 %implementation.
 
 clear all; close all; clc;
-addpath('./Resources'); addpath('./utils')
+addpath('./Resources'); addpath('./utils'); addpath('Resources/qpOASES-3.1.0/interfaces/matlab/');
 rng(2141444)
 
 %% **************** Model and simulation parameters ***********************
 %Dynamics:
-m = 1; g = 9.81; l = 1;
-f_u =  @(t,x,u) [x(2,:); g/l*sin(x(1,:))+ u];
-n = 2; m = 1; %Number of states and control inputs
-Q = eye(2); R = 1; %LQR penalty matrices for states and control inputs
+M = .5; %Cart mass
+m = 0.2; %Pendulum mass
+g = 9.81; %Gravity
+l = 0.4; %Pole length
+f_th = 0.1; %Friction coefficient between pendulum and cart
+
+f_u = @(t,q,u) [q(3,:);...
+    q(4,:);...
+    (-m*g*sin(q(2,:)).*cos(q(2,:))+m*l*q(4,:).^2.*sin(q(2,:))+f_th*m*q(4,:).*cos(q(2,:)) + u)./(M+(1-cos(q(2,:)).^2)*m);...
+    ((M+m)*(g*sin(q(2,:))-f_th*q(2,:))-(l*m*q(4,:).^2.*sin(q(2,:)) + u).*cos(q(2,:)))./(l*(M+(1-cos(q(2,:)).^2)*m))];
+    
+n = 4; m = 1; %Number of states and control inputs
+ub = [2.5 pi/3 2 2]'; %State constraints
+lb = -ub; %State constraints
+Q = 50*eye(n); R = eye(m); %LQR penalty matrices for states and control inputs
 
 %Simulation parameters:
-Ntime = 200;    %Length of each trajectory (# of time steps)
-Ntraj = 20;     %Number of trajectories
+tspan = [0 2];
 deltaT = 0.01;  %Time step length of simulation
-X0_A = 2*rand(n,Ntraj)-1; %Sample initial points for each trajectory for learning A
-X0_A = 0.95*pi*X0_A./vecnorm(X0_A,2,1); %Normalize so initial points lie on unit circle
-Xf_A = zeros(n,Ntraj); %Terminal points for each trajectory for learning A
-X0_B = 0.6*(2*pi*rand(n,Ntraj)-pi); %Sample initial points for each trajectory for learning B
-Xf_B = 0.6*(2*pi*rand(n,Ntraj)-pi); %Terminal points for each trajectory for learning B
+Ntime = (tspan(2)-tspan(1))/deltaT+1; %Length of each trajectory (# of time steps)
+Ntraj = 20;     %Number of trajectories
+rng('shuffle'); %Set seed of random number generator
+q0 = [4*rand(1,Ntraj)-2;...
+    0.2*rand(1,Ntraj)-0.1;...
+    zeros(1,Ntraj);...
+    zeros(1,Ntraj)];
+qf = zeros(n,Ntraj); %Terminal points for each trajectory for learning A
 
 %E-DMD parameters:
-N_basis_edmd = 20; %Number of basis functions 
+N_basis_edmd = 60; %Number of basis functions 
 basis_function_edmd = 'rbf'; %Type of basis function
 rbf_type_edmd = 'thinplate'; %RBF type
 center_type_edmd = 'data'; %Centers of rbf ('data' - pick random points from data set, 'random' pick points uniformly at random
 eps_rbf_edmd = 1; %RBF width
-plot_basis_edmd = false; %Plots the basis functions if true
 
 %Koopman eigenfunction parameters:
-N_basis_koop = 20; %Number of basis functions 
-basis_function_koop = 'rbf'; %Type of basis function
-rbf_type_koop = 'thinplate'; %RBF type
-center_type_koop = 'data'; %Centers of rbf ('data' - pick random points from data set, 'random' pick points uniformly at random
-eps_rbf_koop = 1; %RBF width
-lambda_type = 'eDMD'; %Specifies how to choose candidate lambdas
-N_lambda = 10; %Number of candidate eigenvalues (when random lambdas are used)
-plot_basis_koop = false; %Plots the basis functions if true
+N_basis_diff = 200; %Number of RBFs to use for lifting when learning the diffeomorphism
+pow_eig_pairs = 4; %Highest power of eigenpairs to use when generating eigenfunctions for linearized system
 xlim = [-1, 1]; %Plot limits
 ylim = [-1, 1]; %Plot limits
-learn_type = 'single-step'; %Learn B-matrix with single step or multi-step prediction horizon
 
 %Test simulation parameters:
 Nsim = 5;
-Tsim = 1;
+Tsim = 2;
 plot_results = true;
-X0_sim = 2*rand(n,Nsim)-1;
+q0_sim = [1+rand(1,Nsim);...
+    0.2*rand(1,Nsim)-0.1;...
+    zeros(1,Nsim);...
+    zeros(1,Nsim)];
+qf_sim = zeros(n,Nsim); 
 
 %% ************************** Data Collection *****************************
-disp('Starting data collection...'); tic
-
 % Calculate nominal model with Jacobian and find nominal control gains:
+[A_nom, B_nom, C_nom, ~] = find_nominal_model_ctrl(n,m,f_u,Q,R,zeros(4,1)); %Linearization around 0.
+%Pole placement controller:
+p = [-10 -6 -3 -2];
+K_nom = -place(A_nom,B_nom,p);
 
-[A_nom, B_nom, K_nom] = find_nominal_model_ctrl(n,m,f_u,Q,R);
+% Generate trajectories for all initial and terminal points
+disp('Starting trajectory generation...'); tic
+t = tspan(1) : deltaT : tspan(2); Qf = 1e3*eye(n); % Time vector and penalty for final state deviations
+[q_traj, ~] = generate_trajectory(A_nom, B_nom, C_nom,Qf,q0,qf,t,deltaT,lb,ub);
+fprintf('Trajectory generation done, execution time: %1.2f s \n', toc);
 
 % Collect data to learn autonomous dynamics:
-autonomous_learning = true;
-U_perturb = 0.5*randn(Ntime,Ntraj); %Add normally distributed noise to nominal controller
-[Xstr, Xacc, Yacc, Ustr, Uacc, timestr]  = collect_data(n,m,Ntraj,...
-                  Ntime,deltaT,X0_A, Xf_A ,K_nom,f_u,U_perturb, autonomous_learning);
+disp('Starting data collection...'); tic
 
-% Collect data to learn controlled dynamics:
-autonomous_learning = false;
-U_perturb_c = 0.5*randn(Ntime,Ntraj); %Add normally distributed noise to nominal controller
-[Xstr_c, Xacc_c, Yacc_c, Ustr_c, Uacc_c, timestr_c] = collect_data(n,m,Ntraj,...
-                Ntime,deltaT,X0_B,Xf_B,K_nom,f_u,U_perturb_c, autonomous_learning);
+U_perturb = 0.5*randn(Ntime+1,Ntraj); %Add normally distributed noise to nominal controller
+[Xstr, Ustr, timestr]  = collect_data(n,m,Ntraj,...
+                  Ntime,deltaT,q0, qf ,K_nom,f_u,U_perturb,q_traj);
 
 fprintf('Data collection done, execution time: %1.2f s \n', toc);
-%% *********************** Model Identification ***************************
-           
 
+%% *********************** Model Identification ***************************
 % Identify model using E-DMD to get eigenvalues and model to use for
 % comparison:
 disp('Starting EDMD...'); tic
-[A_edmd, B_edmd, C_edmd, liftFun] = extendedDMD(n,m,N_basis_edmd,basis_function_edmd,...
-    rbf_type_edmd, center_type_edmd, eps_rbf_edmd, plot_basis_edmd, xlim,...
-    ylim, Xacc, Yacc, Uacc);
+
+[A_edmd, B_edmd, C_edmd, liftFun] = extendedDMD(n,m, Ntraj, Ntime, N_basis_edmd,basis_function_edmd,...
+    rbf_type_edmd, center_type_edmd, eps_rbf_edmd, Xstr, Ustr, qf, deltaT);
+
 fprintf('EDMD done, execution time: %1.2f s \n', toc);
 
-%Identify model using Koopman eigenfunctions:
+disp('Starting KEEDMD...'); tic
+
+[A_koop, B_koop, C_koop, phi_fun_v] = KEEDMD(n,m,Ntraj, Ntime, N_basis_diff,...
+        pow_eig_pairs, Xstr, Ustr, qf, A_nom, B_nom, K_nom, deltaT);
 
-[A_koop, B_koop, C_koop, phi_fun_v] = koopman_eigen_id(n, m, Ntraj, Ntime, N_basis_koop,...
-    basis_function_koop, rbf_type_koop, center_type_koop, eps_rbf_koop, ...
-    N_lambda, lambda_type, A_edmd, A_nom, B_nom, K_nom, Xacc, Xstr, Xacc_c,...
-    Yacc_c, Uacc_c, Xstr_c, Ustr_c, timestr, deltaT, learn_type);
+fprintf('KEEDMD done, execution time: %1.2f s \n', toc);
                     
 %% ************************ Analysis of Results ***************************
 
-[mse_edmd_avg, mse_koop_avg, mse_edmd_std, mse_koop_std,...
-    mse_nom_track, mse_edmd_track, mse_koop_track, E_nom, E_edmd, E_koop] = ...
-    evaluate_model_performance(K_nom, A_edmd, B_edmd, C_edmd, A_koop, B_koop, ...
-    C_koop, Nsim, X0_sim, Ntraj, Xstr, Tsim, deltaT, f_u, liftFun, phi_fun_v, plot_results);
- 
- fprintf('-----------------------------------------------------------------\n');
- fprintf('Predicition performance: \nAverage MSE: \n - EDMD:   %1.5f \n - Koopman eigenfunctions: %1.5f \nStandard Deviation MSE: \n - EDMD:  %1.5f \n - Koopman eigenfunctions: %1.5f \n', ...
-        mse_edmd_avg, mse_koop_avg, mse_edmd_std, mse_koop_std);
- fprintf('-----------------------------------------------------------------\n');
- fprintf('Closed loop performance (MPC): \nAverage MSE: \n - Nominal:   %1.5f \n - EDMD:   %1.5f \n - Koopman eigenfunctions: %1.5f \nEnergy consumption proxy: \n - Nominal:  %1.5f \n - EDMD:  %1.5f \n - Koopman eigenfunctions: %1.5f \n', ...
-        mse_nom_track, mse_edmd_track, mse_koop_track, E_nom, E_edmd, E_koop);
- fprintf('-----------------------------------------------------------------\n');
\ No newline at end of file
+t = 0 : deltaT : Tsim; Qf = 1e3*eye(n); % Time vector and penalty for final state deviations
+[q_traj_sim,~] = generate_trajectory(A_nom,B_nom,C_nom,Qf,q0_sim,qf_sim,t,deltaT,lb,ub);
+[q_traj_track,~] = generate_trajectory(A_nom,B_nom,C_nom,Qf,[2 0 0 0]',[0 0 0 0]',t,deltaT,lb,ub); 
+q_traj_track = reshape(q_traj_track,size(q_traj_track,1),size(q_traj_track,3));
+
+disp('Starting analysis of results...'); tic
+
+[mse_origin_avg, mse_edmd_avg, mse_koop_avg, mse_origin_std, mse_edmd_std, mse_koop_std,...
+    cost_nom, cost_edmd, cost_koop] = evaluate_model_performance(A_nom, B_nom, C_nom, K_nom, ...
+    A_edmd, B_edmd, C_edmd, A_koop, B_koop, C_koop, Nsim, q_traj_sim, q_traj_track, ...
+    Ntraj, Xstr, Tsim, deltaT, f_u, liftFun, phi_fun_v, plot_results);
+
+fprintf('Analysis done, execution time: %1.2f s \n', toc);
+ 
\ No newline at end of file
diff --git a/scripts/Koopman-ID/movingslope.m b/scripts/Koopman-ID/movingslope.m
new file mode 100644
index 0000000..6c8b2e8
--- /dev/null
+++ b/scripts/Koopman-ID/movingslope.m
@@ -0,0 +1,219 @@
+function Dvec = movingslope(vec,supportlength,modelorder,dt)
+% movingslope: estimate local slope for a sequence of points, using a sliding window
+% usage: Dvec = movingslope(vec)
+% usage: Dvec = movingslope(vec,supportlength)
+% usage: Dvec = movingslope(vec,supportlength,modelorder)
+% usage: Dvec = movingslope(vec,supportlength,modelorder,dt)
+%
+%
+% movingslope uses filter to determine the slope of a curve stored
+% as an equally (unit) spaced sequence of points. A patch is applied
+% at each end where filter will have problems. A non-unit spacing
+% can be supplied.
+%
+% Note that with a 3 point window and equally spaced data sequence,
+% this code should be similar to gradient. However, with wider
+% windows this tool will be more robust to noisy data sequences.
+%
+%
+% arguments: (input)
+%  vec - row of column vector, to be differentiated. vec must be of
+%        length at least 2.
+%
+%  supportlength - (OPTIONAL) scalar integer - defines the number of
+%        points used for the moving window. supportlength may be no
+%        more than the length of vec.
+%
+%        supportlength must be at least 2, but no more than length(vec)
+%
+%        If supportlength is an odd number, then the sliding window
+%        will be central. If it is an even number, then the window
+%        will be slid backwards by one element. Thus a 2 point window
+%        will result in a backwards differences used, except at the
+%        very first point, where a forward difference will be used.
+%
+%        DEFAULT: supportlength = 3
+%
+%  modelorder - (OPTIONAL) - scalar - Defines the order of the windowed
+%        model used to estimate the slope. When model order is 1, the
+%        model is a linear one. If modelorder is less than supportlength-1.
+%        then the sliding window will be a regression one. If modelorder
+%        is equal to supportlength-1, then the window will result in a
+%        sliding Lagrange interpolant.
+%
+%        modelorder must be at least 1, but not exceeding
+%        min(10,supportlength-1)
+%
+%        DEFAULT: modelorder = 1
+%
+%  dt - (OPTIONAL) - scalar - spacing for sequences which do not have
+%        a unit spacing.
+%
+%        DEFAULT: dt = 1
+%
+% arguments: (output)
+%  Dvec = vector of derivative estimates, Dvec will be of the same size
+%        and shape as is vec.
+% 
+%
+% Example:
+%  Estimate the first derivative using a 7 point window with first through
+%  fourth order models in the sliding window. Note that the higher order
+%  approximations provide better accuracy on this curve with no noise.
+%  
+%  t = 0:.1:1;
+%  vec = exp(t);
+%
+%  Dvec = movingslope(vec,7,1,.1)
+%  Dvec =
+%  Columns 1 through 7
+%    1.3657  1.3657  1.3657  1.3657  1.5093  1.668  1.8435
+%  Columns 8 through 11
+%    2.0373  2.0373  2.0373  2.0373
+%
+%  Dvec = movingslope(vec,7,2,.1)
+%  Dvec =
+%  Columns 1 through 7
+%    0.95747 1.0935  1.2296  1.3657  1.5093  1.668  1.8435
+%  Columns 8 through 11
+%    2.0373  2.2403  2.4433  2.6463
+%
+%  Dvec = movingslope(vec,7,3,.1)
+%  Dvec =
+%  Columns 1 through 7
+%    1.0027  1.1049  1.2206  1.3498  1.4918  1.6487  1.8221
+%  Columns 8 through 11
+%    2.0137  2.2268  2.4602  2.7138
+%
+%  Dvec = movingslope(vec,7,4,.1)
+%  Dvec =
+%    Columns 1 through 7
+%    0.99988 1.1052  1.2214  1.3498  1.4918  1.6487  1.8221
+%  Columns 8 through 11
+%    2.0137  2.2255  2.4597  2.7181
+%
+%
+% Example:
+%  Estimate the slope of a noisy curve, using a locally quadratic
+%  approximation. In this case, use a straight line so that we know
+%  the true slope should be 1. Use a wide window, since we have
+%  noisy data.
+%  
+%  t = 0:100;
+%  vec = t + randn(size(t));
+%  Dvec = movingslope(vec,10,2,1)
+%  mean(Dvec)
+%  ans = 
+%     1.0013
+%  std(Dvec)
+%  ans =
+%     0.10598
+%
+%  By way of comparison, gradient gives a much noisier estimate
+%  of the slope of this curve.
+%
+%  std(gradient(vec))
+%  ans =
+%     0.69847
+%
+%
+% Example:
+%  As a time test, generate random data vector of length 500000.
+%  Compute the slopes using a window of width 10.
+%
+%  vec = rand(1,500000);
+%  tic
+%  Dvec = movingslope(vec,10,2);
+%  toc
+%
+%  Elapsed time is 0.626021 seconds.
+%
+%
+% See also: gradient
+%
+%
+% Author: John D'Errico
+% e-mail: woodchips@rochester.rr.com
+% Release: 1.0
+% Release date: 10/19/07
+
+% how long is vec? is it a vector?
+if (nargin==0)
+  help movingslope
+  return
+end
+if ~isvector(vec)
+  error('vec must be a row or column vector')
+end
+n = length(vec);
+
+% supply defaults
+if (nargin<4) || isempty(dt)
+  dt = 1;
+end
+if (nargin<3) || isempty(modelorder)
+  modelorder = 1;
+end
+if (nargin<2) || isempty(supportlength)
+  supportlength = 3;
+end
+
+% check the parameters for problems
+if (length(supportlength)~=1) || (supportlength<=1) || (supportlength>n) || (supportlength~=floor(supportlength))
+  error('supportlength must be a scalar integer, >= 2, and no more than length(vec)')
+end
+if (length(modelorder)~=1) || (modelorder<1) || (modelorder>min(10,supportlength-1)) || (modelorder~=floor(modelorder))
+  error('modelorder must be a scalar integer, >= 1, and no more than min(10,supportlength-1)')
+end
+if (length(dt)~=1) || (dt<0)
+  error('dt must be a positive scalar numeric variable')
+end
+
+% now build the filter coefficients to estimate the slope
+if mod(supportlength,2) == 1
+  parity = 1; % odd parity
+else
+  parity = 0;
+end
+s = (supportlength-parity)/2;
+t = ((-s+1-parity):s)';
+coef = getcoef(t,supportlength,modelorder);
+
+% Apply the filter to the entire vector
+f = filter(-coef,1,vec);
+Dvec = zeros(size(vec));
+Dvec(s+(1:(n-supportlength+1))) = f(supportlength:end);
+
+% patch each end
+vec = vec(:);
+for i = 1:s
+  % patch the first few points
+  t = (1:supportlength)' - i;
+  coef = getcoef(t,supportlength,modelorder);
+  
+  Dvec(i) = coef*vec(1:supportlength);
+  
+  % patch the end points
+  if i<(s + parity)
+    t = (1:supportlength)' - supportlength + i - 1;
+    coef = getcoef(t,supportlength,modelorder);
+    Dvec(n - i + 1) = coef*vec(n + (0:(supportlength-1)) + 1 - supportlength);
+  end
+end
+
+% scale by the supplied spacing
+Dvec = Dvec/dt;
+% all done
+
+end % mainline end
+
+% =========================================================
+% subfunction, used to compute the filter coefficients
+function coef = getcoef(t,supportlength,modelorder)
+% Note: bsxfun would have worked here as well, but some people
+% might not yet have that release of matlab.
+A = repmat(t,1,modelorder+1).^repmat(0:modelorder,supportlength,1);
+pinvA = pinv(A);
+% we only need the linear term
+coef = pinvA(2,:);
+end % nested function end
diff --git a/scripts/Koopman-ID/normalize_data.m b/scripts/Koopman-ID/normalize_data.m
new file mode 100644
index 0000000..3ef4e0b
--- /dev/null
+++ b/scripts/Koopman-ID/normalize_data.m
@@ -0,0 +1,11 @@
+function [X_norm, X_mean, X_std] = normalize_data(X)
+    
+    %Make sure each row corresponds to a time step:
+    if size(X,1)>size(X,2)
+        X = X';
+    end
+    
+    X_mean = mean(X,2);
+    X_std = std(X,0,2);
+    X_norm = (X-X_mean)./X_std;
+end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/num_diff.m b/scripts/Koopman-ID/num_diff.m
new file mode 100644
index 0000000..22acd14
--- /dev/null
+++ b/scripts/Koopman-ID/num_diff.m
@@ -0,0 +1,7 @@
+function diff = num_diff(X, deltaT)
+    
+    diff = zeros(size(X));
+    for k = 1 : size(X,1)
+        diff(k,:) = movingslope(X(k,:),4,2,deltaT);
+    end
+end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/obsvk.m b/scripts/Koopman-ID/obsvk.m
new file mode 100755
index 0000000..66191ee
--- /dev/null
+++ b/scripts/Koopman-ID/obsvk.m
@@ -0,0 +1,6 @@
+function [ O ] = obsvk( A, C, k )
+O = [];
+for i = 1:k
+    O = [O;C*A^(i-1)]; % Inefficient
+end
+
diff --git a/scripts/Koopman-ID/phifun_mat.m b/scripts/Koopman-ID/phifun_mat.m
new file mode 100644
index 0000000..5c0c83a
--- /dev/null
+++ b/scripts/Koopman-ID/phifun_mat.m
@@ -0,0 +1,7 @@
+function res = phifun_mat(phifun,X)
+    [n,m] = size(X);
+    res = [];
+    for i = 1 : m
+        res = [res phifun(X(:,i))];
+    end 
+end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/rbf_dot.m b/scripts/Koopman-ID/rbf_dot.m
new file mode 100644
index 0000000..d8ef5c0
--- /dev/null
+++ b/scripts/Koopman-ID/rbf_dot.m
@@ -0,0 +1,53 @@
+% X: nxN
+% C: rbf center(s) - nx1 or nxK, if C is n x K, then the result is STACKED
+% over the centers K, i.e., of dimension K x N
+% eps: kernel width for Gaussian type rbfs (optional)
+% k: polyharmonic coefficient for polyharmonic rbfs (optional)
+% Note that all RBFs are symmetrical so evaluation of of a single point on
+% multiple centers can by done as evaluation of multiple points (the
+% centers in this case) on a single point (the center)
+
+function Y = rbf_dot(X,X_dot,C, type, eps, k )
+    type = lower(type);
+    if(~exist('eps','var') || isempty(eps))
+        eps = 1;
+    end
+    if(~exist('k','var') || isempty(k))
+        k = 1;
+    end
+    
+    N = size(X,2); %number of points
+    n = size(X,1);
+    Cbig = C;
+    Y = zeros(size(C,2),N);
+    for i = 1:size(Cbig,2)
+        C = Cbig(:,i);
+        C = repmat( C,1, N );
+        if (n>1)
+            r_squared = sum( (X - C).^2 );
+        else
+            r_squared = abs( (X - C).^2 );
+        end
+        switch type
+            case 'thinplate'
+                error('Gradient of RBF type not implemented')
+                %y = r_squared.*log(sqrt(r_squared));  % + 0.5*sqrt(r_squared); % makes nonnegative, try to impose nonnegativity of K
+                %y(isnan(y)) = 0;
+            case 'gauss'
+                y = -2*eps^2.*(X-C).*exp(-eps^2*r_squared);
+                %y = exp(-eps^2*r_squared);
+            case 'invquad'
+                error('Gradient of RBF type not implemented')
+                %y = 1 ./ (1+eps^2*r_squared);
+            case 'invmultquad'
+                error('Gradient of RBF type not implemented')
+                %y = 1 ./ sqrt((1+eps^2*r_squared));
+            case 'polyharmonic'
+                error('Gradient of RBF type not implemented')
+                %y = r_squared.^(k/2).*log(sqrt(r_squared));
+                %y(isnan(y)) = 0;
+            otherwise
+                error('RBF type not recognized')
+        end
+        Y(i,:) = sum(y.*X_dot,1);
+    end
diff --git a/scripts/Koopman-ID/rbf_grad.m b/scripts/Koopman-ID/rbf_grad.m
new file mode 100644
index 0000000..37baa32
--- /dev/null
+++ b/scripts/Koopman-ID/rbf_grad.m
@@ -0,0 +1,53 @@
+% X: nxN
+% C: rbf center(s) - nx1 or nxK, if C is n x K, then the result is STACKED
+% over the centers K, i.e., of dimension K x N
+% eps: kernel width for Gaussian type rbfs (optional)
+% k: polyharmonic coefficient for polyharmonic rbfs (optional)
+% Note that all RBFs are symmetrical so evaluation of of a single point on
+% multiple centers can by done as evaluation of multiple points (the
+% centers in this case) on a single point (the center)
+
+function Y = rbf_grad( X,C, type, eps, k )
+    type = lower(type);
+    if(~exist('eps','var') || isempty(eps))
+        eps = 1;
+    end
+    if(~exist('k','var') || isempty(k))
+        k = 1;
+    end
+    
+    N = size(X,2); %number of points
+    n = size(X,1);
+    Cbig = C;
+    Y = zeros(size(C,2),n);
+    for i = 1:size(Cbig,2)
+        C = Cbig(:,i);
+        C = repmat( C,1, N );
+        if (n>1)
+            r_squared = sum( (X - C).^2 );
+        else
+            r_squared = abs( (X - C).^2 );
+        end
+        switch type
+            case 'thinplate'
+                error('Gradient of RBF type not implemented')
+                %y = r_squared.*log(sqrt(r_squared));  % + 0.5*sqrt(r_squared); % makes nonnegative, try to impose nonnegativity of K
+                %y(isnan(y)) = 0;
+            case 'gauss'
+                y = -2*eps^2.*(X-C).*exp(-eps^2*r_squared);
+                %y = exp(-eps^2*r_squared);
+            case 'invquad'
+                error('Gradient of RBF type not implemented')
+                %y = 1 ./ (1+eps^2*r_squared);
+            case 'invmultquad'
+                error('Gradient of RBF type not implemented')
+                %y = 1 ./ sqrt((1+eps^2*r_squared));
+            case 'polyharmonic'
+                error('Gradient of RBF type not implemented')
+                %y = r_squared.^(k/2).*log(sqrt(r_squared));
+                %y(isnan(y)) = 0;
+            otherwise
+                error('RBF type not recognize')
+        end
+        Y(i,:) = y';
+    end
diff --git a/scripts/Koopman-ID/test_eigfunction_generation.m b/scripts/Koopman-ID/test_eigfunction_generation.m
new file mode 100644
index 0000000..3037319
--- /dev/null
+++ b/scripts/Koopman-ID/test_eigfunction_generation.m
@@ -0,0 +1,143 @@
+clear all; close all; clc;
+addpath('./Resources'); addpath('./utils')
+rng(2141444)
+
+%% **************** Model and simulation parameters ***********************
+%Dynamics:
+m = 1; g = 9.81; l = 1;
+f_u =  @(t,x,u) [x(2,:); g/l*sin(x(1,:))+ u];
+n = 2; m = 1; %Number of states and control inputs
+Q = eye(2); R = 1; %LQR penalty matrices for states and control inputs
+
+%Simulation parameters:
+Ntime = 200;    %Length of each trajectory (# of time steps)
+Ntraj = 10;     %Number of trajectories
+deltaT = 0.01;  %Time step length of simulation
+X0 = 2*rand(n,Ntraj)-1; %Sample initial points for each trajectory for learning A
+X0 = 0.95*pi*X0./vecnorm(X0,2,1); %Normalize so initial points lie on unit circle
+Xf = zeros(n,Ntraj); %Terminal points for each trajectory for learning A
+
+%% ************************** Data Collection *****************************
+disp('Starting data collection...'); tic
+
+% Calculate nominal model with Jacobian and find nominal control gains:
+
+[A_nom, B_nom, K_nom] = find_nominal_model_ctrl(n,m,f_u,Q,R,[0;0]);
+
+% Collect data to learn autonomous dynamics:
+autonomous_learning = true;
+U_perturb = 0.0*randn(Ntime+1,Ntraj); %Add normally distributed noise to nominal controller
+[Xstr, Xacc, Yacc, Ustr, U, timestr]  = collect_data(n,m,Ntraj,...
+                  Ntime,deltaT,X0, Xf ,K_nom,f_u,U_perturb);
+
+fprintf('Data collection done, execution time: %1.2f s \n', toc);
+
+%% ************** Prepare data and learn diffeomorphism h *****************
+    Xstr_shift = zeros(size(Xstr)); %Shift dynamics such that origin is fixed point
+    X = [];
+    X_dot = [];
+    U = [];
+    for i = 1 : Ntraj
+       Xstr_shift(:,i,:) = Xstr(:,i,:) - Xf(:,i);
+       X = [X reshape(Xstr_shift(:,i,:),n,Ntime+1)];
+       X_dot = [X_dot num_diff(reshape(Xstr_shift(:,i,:),size(Xstr_shift,1),size(Xstr_shift,3)),deltaT)];
+       U = [U reshape(Ustr(:,i,:),m,Ntime+1)];
+    end
+
+    N_basis = 100;
+    A_cl = A_nom + B_nom*K_nom;
+    cent = 2*pi*rand(n,N_basis) - pi;
+    rbf_type = 'gauss';
+    eps_rbf = 2;
+        
+    % Set up nonlinear transformations and their gradients
+    zfun = @(xx) [xx; rbf(xx,cent,rbf_type,eps_rbf)];
+    zfun_grad = @(xx) [eye(2); rbf_grad(xx,cent,rbf_type,eps_rbf)];
+    zfun_dot = @(xx, xx_dot) [xx_dot; rbf_dot(xx,xx_dot,cent,rbf_type,eps_rbf)];
+    
+   % Set up Z and Z_dot matrices:
+   Z = zfun(X);
+   Z_dot = zfun_dot(X,X_dot);
+
+    %Set up constraint matrix
+    con1 = zfun_grad([0; 0]);
+    
+    disp('Solving optimization problem...')
+    N = size(Z,1);
+    cvx_begin 
+        variable C(n,N);
+        minimize (norm(X_dot + C*Z_dot - A_cl*(X+C*Z),'fro') + 1*norm(C,'fro'))
+        subject to
+            {C*con1 == zeros(n)}; 
+    cvx_end
+    fprintf('Solved, optimal value excluding regularization: %.10f, MSE: %.10f\n', norm(X_dot + C*Z_dot - A_cl*(X+C*Z),'fro'),immse(X_dot+C*Z_dot, A_cl*(X+C*Z)))
+    %fprintf('Constraint violation: %.4f \n', sum(sum(abs(C*con1))))
+
+    yfun = @(xx) xx + C*zfun(xx); % Full learned diffeomorphism
+    
+    % Calculate eigenfunctions for linearized system
+    [V_a,D] = eig(A_cl);
+    [W_a,~] = eig(A_cl');
+
+    %Define powers (only implemented for n=2):
+    pow_eig_pairs = 3;
+    a = 0 : pow_eig_pairs;
+    [P,Q] = meshgrid(a,a);
+    c=cat(2,P',Q');
+    powers=reshape(c,[],2);
+
+    linfunc = @(xx) (xx'*W_a)'./diag(V_a'*W_a);
+    phifun = @(xx) (prod(linfunc(xx).^(powers')))';
+    lambd = prod(exp(diag(D)).^(powers'))';
+    lambd = log(lambd);
+
+    % Construct scaling function
+    gfun = @(xx) xx./pi; %Scale state space into unit cube
+
+    % Construct eigenfunctions for nonlinear system
+    z_eigfun = @(xx) phifun_mat(phifun, gfun(yfun(xx)));
+    %z_eigfun = @(xx) phifun(gfun(yfun(xx)));
+    A_koop = diag(lambd);
+%% Check how well the evolution of the eigenfunctions are described by the eigenvalues on linearized system:
+f_cl = @(t,x) A_cl*x;
+x0 = [1;1];
+[t_test, X_test] = ode45(f_cl,0:deltaT:2,x0);
+X_test = X_test';
+
+N = size(lambd,1);
+Z_test = phifun_mat(phifun,X_test);
+f_eig = @(t,z) A_koop*z;
+[~,Z_eig] = ode45(f_eig,0:deltaT:2,Z_test(:,1));
+Z_eig = Z_eig';
+
+afigure
+t = 0 : deltaT : Ntime*deltaT;
+for i = 1 : N
+    subplot(4,4,i)
+    plot(t,Z_eig(i,:),t,Z_test(i,:))
+    xlabel('Time (sec)');
+    ylabel('Observable value')
+    if i == N
+        legend('Eigval evolution', 'True')
+    end
+end
+    
+
+%% Check how well the evolution of the eigenfunctions are described by the eigenvalues:
+X_test = reshape(Xstr(:,1,:),n,Ntime+1);
+N = size(lambd,1);
+Z_test = z_eigfun(X_test);
+[~,Z_eig] = ode45(f_eig,0:deltaT:2,Z_test(:,1));
+Z_eig = Z_eig';
+
+afigure
+t = 0 : deltaT : Ntime*deltaT;
+for i = 1 : N
+    subplot(4,4,i)
+    plot(t,Z_eig(i,:),t,Z_test(i,:))
+    xlabel('Time (sec)');
+    ylabel('Observable value')
+    if i == N
+        legend('Eigval evolution', 'True')
+    end
+end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/test_trajectory_gen.m b/scripts/Koopman-ID/test_trajectory_gen.m
new file mode 100644
index 0000000..774b11e
--- /dev/null
+++ b/scripts/Koopman-ID/test_trajectory_gen.m
@@ -0,0 +1,85 @@
+clear all; close all; clc;
+%% **************** Model and simulation parameters ***********************
+%Dynamics:
+M = .5; %Cart mass
+m = 0.2; %Pendulum mass
+g = 9.81; %Gravity
+l = 0.4; %Pole length
+f_th = 0.1; %Friction coefficient between pendulum and cart
+
+% f_u = @(t,q,u) [q(2);...
+%     (-m*g*sin(q(3))*cos(q(3))+m*l*q(4)^2*sin(q(3))+f_th*m*q(4)*cos(q(3)) + u)/(M+(1-cos(q(3))^2)*m);...
+%     q(4);...
+%     ((M+m)*(g*sin(q(3))-f_th*q(4))-(l*m*q(4)^2*sin(q(3)) + u)*cos(q(3)))/(l*(M+(1-cos(q(3))^2)*m))];
+
+f_u = @(t,q,u) [q(3,:);...
+    q(4,:);...
+    (-m*g*sin(q(2,:)).*cos(q(2,:))+m*l*q(4,:).^2.*sin(q(2,:))+f_th*m*q(4,:).*cos(q(2,:)) + u)./(M+(1-cos(q(2,:)).^2)*m);...
+    ((M+m)*(g*sin(q(2,:))-f_th*q(2,:))-(l*m*q(4,:).^2.*sin(q(2,:)) + u).*cos(q(2,:)))./(l*(M+(1-cos(q(2,:)).^2)*m))];
+ 
+    
+n = 4; m = 1; %Number of states and control inputs
+
+%Simulation parameters:
+tSpan = [0 2];
+deltaT = 0.01;  %Time step length of simulation
+q0 = [2 0 0 0]';
+qf = [0 0 0 0]';
+
+%Linearize dynamics around x0
+Q = 20*eye(n); R = eye(m);
+[A_nom, B_nom, C_nom, ~] = find_nominal_model_ctrl(n, m, f_u, Q, R, qf);
+
+%Pole placement controller:
+p = [-10 -6 -3 -2];
+K_nom = -place(A_nom,B_nom,p);
+
+%% Test trajectory generator based on linearized model
+N = n; %Lifting dimension
+t = tSpan(1) : deltaT : tSpan(end);
+N_sim = length(t);
+R_traj = eye(m);
+Qf = 1000*eye(n);
+ub = [3 pi/3 10 10]';
+lb = -ub;
+
+[z_t,u_t] = generate_trajectory(A_nom,B_nom,C_nom,Qf,q0,qf,t,deltaT,lb,ub);
+
+afigure
+for i = 1 : n
+    subplot(2,2,i);
+    plot(t,z_t(i,:));
+end
+
+afigure
+plot(t(1:end-1),u_t(1,:))
+
+
+%% Test LQR controller
+u_lqr = zeros(m,N_sim-1);
+q_lqr = zeros(n,N_sim);
+q_lqr(:,1) = q0;
+
+u_lqr_t = zeros(m,N_sim-1);
+q_lqr_t = zeros(n,N_sim);
+q_lqr_t(:,1) = q0;
+
+for i = 1 : N_sim-1
+    u_lqr_t(:,i) = K_nom*(q_lqr_t(:,i)-z_t(:,i));
+    q_lqr_t(:,i+1) = sim_timestep(deltaT,f_u,0,q_lqr_t(:,i),u_lqr_t(:,i));
+%     u_lqr(:,i) = K_nom*(q_lqr(:,i)-qf);
+%     q_lqr(:,i+1) = sim_timestep(deltaT,f_u,0,q_lqr(:,i),u_lqr(:,i));
+end
+
+afigure
+for i = 1 : n
+    subplot(2,2,i);
+    plot(t,q_lqr_t(i,:),t,q_lqr(i,:),t,z_t(i,:));
+    legend('lqr traj', 'lqr','traj')
+end
+
+afigure
+plot(t(1:end-1),u_lqr_t(1,:),t(1:end-1),u_lqr(1,:),t(1:end-1),u_t(1,:))
+legend('lqr traj', 'lqr','traj')
+    
+    
\ No newline at end of file
diff --git a/scripts/Koopman-ID/utils/collect_data.m b/scripts/Koopman-ID/utils/collect_data.m
index d2aa418..5692e99 100644
--- a/scripts/Koopman-ID/utils/collect_data.m
+++ b/scripts/Koopman-ID/utils/collect_data.m
@@ -1,5 +1,5 @@
-function [Xstr, Xacc, Yacc, Ustr, Uacc, time_str] = collect_data(n,m,Ntraj,...
-                  Ntime,deltaT,X0,Xf_A,K_nom,f_u,U_perturb, autonomous_learning)
+function [Xstr, Ustr, time_str] = collect_data(n,m,Ntraj,...
+                  Ntime,deltaT,X0,Xf,K_nom,f_u,U_perturb,qtraj)
     %Collect data from the "true" system with nominal controller plus
     %perturbation
     %Inputs:
@@ -10,38 +10,26 @@
     %Outputs:
     %   Xstr    - (n x Ntraj x Ntime) struct with data from each timestep and 
     %              trajectory for each state
-    %   Xacc    - (n x Ntraj*Ntime) matrix with all trajectories added to it
-    %   Yacc    - (n x Ntraj*Ntime) matrix with all trajectories added to it
-    %              Yacc_k = Xacc_k+1
-    %   Uacc    - (m x Ntraj*Ntime) matrix with all trajectories added to it
-    %   Xstr    - (1 x Ntraj x Ntime) struct with the time stamp of each data 
+    %   Ustr    - (1 x Ntraj x Ntime) struct with the time stamp of each data 
     %              point from each trajectory
 
 
-    Xstr = zeros(n,Ntraj,Ntime+1); % *str is structure
-    Ustr = zeros(m,Ntraj,Ntime); 
-    Xacc = zeros(n,Ntraj*(Ntime)); % *acc is accumulated vectors
-    Yacc = zeros(n,Ntraj*(Ntime)); 
-    Uacc = zeros(m,Ntraj*(Ntime)); 
-
-    time_str = zeros(Ntraj,Ntime+1);
+    Xstr = zeros(n,Ntraj,Ntime); % *str is structure
+    Ustr = zeros(m,Ntraj,Ntime-1); 
+    time_str = zeros(Ntraj,Ntime);
+    
     Xcurrent = X0;
     Xstr(:,:,1) = X0;
-    for i = 1:Ntime
-        Ucurrent = K_nom*(Xcurrent-Xf_A)+U_perturb(i,:);
+    for i = 1:Ntime-1
+        Ucurrent = K_nom*(Xcurrent-qtraj(:,:,i))+U_perturb(i,:);
+        Ustr(:,:,i) = Ucurrent;
+        
         Xnext = sim_timestep(deltaT, f_u, 0, Xcurrent, Ucurrent);
         Xstr(:,:,i+1) = Xnext;
-        Xacc(:,Ntraj*(i-1)+1:Ntraj*(i)) = Xcurrent;
-        Yacc(:,Ntraj*(i-1)+1:Ntraj*(i)) = Xnext;
         Xcurrent = Xnext;
-        time_str(:,i+1) = i*deltaT*ones(Ntraj,1);
         
-        if autonomous_learning
-            Ustr(:,:,i) = Ucurrent;
-            Uacc(:,Ntraj*(i-1)+1:Ntraj*i) = Ucurrent;
-        else
-            Ustr(:,:,i) = Ucurrent - K_nom*Xcurrent;
-            Uacc(:,Ntraj*(i-1)+1:Ntraj*i) = Ucurrent - K_nom*Xcurrent;
-        end
+        time_str(:,i+1) = i*deltaT*ones(Ntraj,1);
     end
+    Ucurrent = K_nom*(Xcurrent-qtraj(:,:,i+1))+U_perturb(i+1,:);
+    Ustr(:,:,i+1) = Ucurrent;
 end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/utils/evaluate_model_performance.m b/scripts/Koopman-ID/utils/evaluate_model_performance.m
index b8e3e64..6645eaf 100644
--- a/scripts/Koopman-ID/utils/evaluate_model_performance.m
+++ b/scripts/Koopman-ID/utils/evaluate_model_performance.m
@@ -1,95 +1,168 @@
-function [mse_edmd_avg, mse_koop_avg, mse_edmd_std, mse_koop_std,...
-    mse_nom_track, mse_edmd_track, mse_koop_track, E_nom, E_edmd, E_koop] = ...
-    evaluate_model_performance(K_nom, A_edmd, B_edmd, C_edmd, A_koop, B_koop, ...
-    C_koop, Nsim, X0, Ntraj, Xstr, Tsim, deltaT, f_u, liftFun, phi_fun_v, plot_results)
+function [mse_origin_avg, mse_edmd_avg, mse_koop_avg, mse_origin_std, mse_edmd_std, mse_koop_std,...
+    cost_nom, cost_edmd, cost_koop] = evaluate_model_performance(A_nom, B_nom, C_nom, K_nom, ...
+    A_edmd, B_edmd, C_edmd, A_koop, B_koop, C_koop, Nsim, q_traj_sim, q_traj_track, Ntraj, ...
+    Xstr, Tsim, deltaT, f_u, liftFun, phi_fun_v, plot_results)
     
     n = size(C_koop,1);
     m = size(B_koop,2);
     n_edmd = size(A_edmd,1);
     n_koop = size(A_koop,1);
     Ntime = Tsim/deltaT;
-    A_koop_d = expm(A_koop*deltaT); %Discretize A_koop
     
     %Prediction task:
-    [x,x_edmd,x_koop,mse_edmd_avg,mse_koop_avg,mse_edmd_std,mse_koop_std]...
-            = sim_prediction(n,m,n_edmd,n_koop,Nsim,Ntime,deltaT,X0,f_u,liftFun,...
-            phi_fun_v, K_nom,A_edmd,B_edmd,C_edmd,A_koop_d,B_koop,C_koop);
+    [x,x_origin, x_edmd,x_koop, mse_origin_avg, mse_edmd_avg,mse_koop_avg,...
+    mse_origin_std, mse_edmd_std,mse_koop_std]...
+            = sim_prediction(n,m,n_edmd,n_koop,Nsim,Ntime,deltaT,q_traj_sim,f_u,liftFun,...
+            phi_fun_v,A_nom,B_nom,C_nom,K_nom,A_edmd,B_edmd,C_edmd,A_koop,B_koop,C_koop);
     
     %Closed loop task:
-    [x_track,x_edmd_track,x_koop_track,mse_nom_track, mse_edmd_track,...
-        mse_koop_track,t_plot,traj_d, E_nom, E_edmd, E_koop] = ...
-        sim_closed_loop_mpc(n,m,n_edmd,n_koop,Nsim,...
-        Ntime,deltaT,Tsim,f_u,liftFun,phi_fun_v, K_nom,A_edmd,B_edmd,...
-        C_edmd,A_koop_d,B_koop,C_koop);
+    q = [1 5 50];
+    r = [0.01 0.01 0.01];
+    
+    %Discretize nominal system:
+    continuous_model_nom = ss(A_nom, B_nom, C_nom, 0);
+    discrete_model_nom = c2d(continuous_model_nom,deltaT);
+    A_nom_d = discrete_model_nom.A;
+    B_nom_d = discrete_model_nom.B;
+    C_nom_d = discrete_model_nom.C;
+    
+    %Discretize EDMD system:
+    continuous_model_edmd = ss(A_edmd, B_edmd, C_edmd, 0);
+    discrete_model_edmd = c2d(continuous_model_edmd,deltaT);
+    A_edmd_d = discrete_model_edmd.A;
+    B_edmd_d = discrete_model_edmd.B;
+    C_edmd_d = discrete_model_edmd.C;
+    
+    %Discretize KEEDMD system:
+    continuous_model_koop = ss(A_koop, B_koop, C_koop, 0);
+    discrete_model_koop = c2d(continuous_model_koop,deltaT);
+    A_koop_d = discrete_model_koop.A;
+    B_koop_d = discrete_model_koop.B;
+    C_koop_d = discrete_model_koop.C;
+    
+    for i = 1 : length(q)
+        Q = q(i)*eye(n);
+        R = r(i)*eye(m);
+        [x_track{i},x_edmd_track{i},x_koop_track{i},mse_nom_track{i}, mse_edmd_track{i},...
+            mse_koop_track{i},traj_d, E_nom{i}, E_edmd{i}, E_koop{i},...
+            cost_nom{i}, cost_edmd{i}, cost_koop{i}] = ...
+            sim_closed_loop_mpc(n,m,n_edmd,n_koop,Nsim,...
+            Ntime,deltaT,Tsim,f_u,liftFun,phi_fun_v,A_nom_d,B_nom_d,C_nom_d,K_nom,A_edmd_d,B_edmd_d,...
+            C_edmd_d,A_koop_d,B_koop_d,C_koop_d,Q,R,q_traj_track);
+    end
 
     %Plot results
     if plot_results
-        lw=4;
+        lw=2;
+        t_plot = 0 : deltaT : Tsim;
+        font_size = 16;
+        set(gcf,'units','pixel')
         
+        close all; %TODO: Remove after figure design
         afigure(1)
-        h1 = subplot(2,2,1);
+        subplot(3,2,1);
         hold on
         for i = 1 %Only plot first trajectory
-            plot([0:Ntime]*deltaT,reshape(x(1,i,:),Ntime+1,1),'linewidth',lw); hold on
-            plot([0:Ntime]*deltaT,reshape(x_edmd(1,i,:),Ntime+1,1), '--r','linewidth',lw)
-            plot([0:Ntime]*deltaT,reshape(x_koop(1,i,:),Ntime+1,1), '--g','linewidth',lw-1)
+            plot([0:Ntime]*deltaT,reshape(x(1,i,:),Ntime+1,1),':b','linewidth',lw); hold on
+            plot([0:Ntime]*deltaT,reshape(x_origin(1,i,:),Ntime+1,1), 'k','linewidth',lw)
+            plot([0:Ntime]*deltaT,reshape(x_edmd(1,i,:),Ntime+1,1), 'r','linewidth',lw)
+            plot([0:Ntime]*deltaT,reshape(x_koop(1,i,:),Ntime+1,1), 'g','linewidth',lw)
         end
-        axis([0 Tsim min(x_edmd(1,1,:))-0.15 max(x_edmd(1,1,:))+0.15])
-        title('Predictor comparison, 1st trajectory - $x_1$','interpreter','latex'); xlabel('Time [s]','interpreter','latex');
-        set(gca,'fontsize',20)
-        LEG = legend('True','EDMD','Koopman e-func','location','southwest');
+        axis([0 Tsim min(x(1,1,:))-0.2 max(x(1,1,:))+0.2])
+        title('\textbf{Prediction Performance}','interpreter','latex'); 
+        xlabel('Time [s]','interpreter','latex');
+        ylabel('$x_1$','interpreter','latex');
+        set(gca,'fontsize',font_size)
+        LEG = legend({'True','Lin at 0', 'EDMD','KEEDMD'},'location','best', 'NumColumns',2);
         set(LEG,'interpreter','latex')
-        h1.Position = h1.Position + [0 0.1 0 -0.1]; %Modify size of 1st subplot
+        %h1.Position = h1.Position + [0 0.1 0 -0.1]; %Modify size of 1st subplot
         
-        h2 = subplot(2,2,2);
+        %h2 = subplot(2,2,2);
+        h2 = subplot(3,2,3);
         hold on
         for i = 1 %Only plot first trajectory
-            plot([0:Ntime]*deltaT,reshape(x(2,i,:),Ntime+1,1),'linewidth',lw); hold on
-            plot([0:Ntime]*deltaT,reshape(x_edmd(2,i,:),Ntime+1,1), '--r','linewidth',lw)
-            plot([0:Ntime]*deltaT,reshape(x_koop(2,i,:),Ntime+1,1), '--g','linewidth',lw-1)
+            plot([0:Ntime]*deltaT,reshape(x(3,i,:),Ntime+1,1),':b','linewidth',lw); hold on
+            plot([0:Ntime]*deltaT,reshape(x_origin(3,i,:),Ntime+1,1), 'k','linewidth',lw)
+            plot([0:Ntime]*deltaT,reshape(x_edmd(3,i,:),Ntime+1,1), 'r','linewidth',lw)
+            plot([0:Ntime]*deltaT,reshape(x_koop(3,i,:),Ntime+1,1), 'g','linewidth',lw)
         end
-        axis([0 Tsim min(x_edmd(2,1,:))-0.15 max(x_edmd(2,1,:))+0.15])
-        title('Predictor comparison, 1st trajectory - $x_2$','interpreter','latex'); xlabel('Time [s]','interpreter','latex');
-        set(gca,'fontsize',20)
+        axis([0 Tsim min(x(3,1,:))-0.2 max(x(3,1,:))+0.2])
+        %title('Prediction comparison, 1st trajectory - $x_2$','interpreter','latex'); 
+        xlabel('Time [s]','interpreter','latex');      
+        ylabel('$x_3$','interpreter','latex');
+        set(gca,'fontsize',font_size)
         %LEG = legend('True','EDMD','Koopman e-func','location','southwest');
         set(LEG,'interpreter','latex')
-        h2.Position = h2.Position + [0 0.1 0 -0.1]; %Modify size of 1st subplot
+        %h2.Position = h2.Position + [0 0.1 0 -0.1]; %Modify size of 1st subplot
         
-        % Plot trajectories Learning
-        h3 = subplot(2,2,3);
-        hold on
-        for i=1:Ntraj
-           scatter(Xstr(1,i,1),Xstr(2,i,1),'+')
-           p(i) = plot(squeeze(Xstr(1,i,:)),squeeze(Xstr(2,i,:)),'Color',[0.2,0.2,0.2],'linewidth',1);
-           %alpha(p(i),0.2)   
+        
+        %Closed loop plot:
+        h4 = subplot(3,2,2);
+        plot(t_plot,traj_d(1,:),'-b','linewidth',lw); hold on
+        lin_typ = {'--', '-', '-.'};
+        for i = 2% : length(q)
+            plot(t_plot,x_track{i}(1,:),strcat(lin_typ{i},'k'),'linewidth',lw); hold on
+            plot(t_plot,x_edmd_track{i}(1,:),strcat(lin_typ{i},'r'),'linewidth',lw); hold on
+            plot(t_plot,x_koop_track{i}(1,:),strcat(lin_typ{i},'g'),'linewidth',lw); hold on
         end
-        for i = 1 : Nsim
-            plot(reshape(x(1,i,:),Ntime+1,1),reshape(x(2,i,:),Ntime+1,1),'-r')
-            plot(reshape(x_edmd(1,i,:),Ntime+1,1),reshape(x_edmd(2,i,:),Ntime+1,1),'-b')
-            plot(reshape(x_koop(1,i,:),Ntime+1,1),reshape(x_koop(2,i,:),Ntime+1,1),'-g')
+            %axis([0 Tsim min(x_edmd(1,1,:))-0.15 max(x_edmd(1,1,:))+0.15])
+        title('\textbf{Closed Loop Tracking Performance}','interpreter','latex'); 
+        xlabel('Time [s]','interpreter','latex');
+        ylabel('$x_1$','interpreter','latex');
+        set(gca,'fontsize',font_size)
+%         LEG = legend({'Reference', ...
+%             strcat('Lin at 0$_{q/r=',num2str(q(1)/r(1)),'}$'),strcat('EDMD$_{q/r=',num2str(q(1)/r(1)),'}$'),strcat('KEEDMD$_{q/r=',num2str(q(1)/r(1)),'}$'),...
+%             strcat('Lin at 0$_{q/r=',num2str(q(2)/r(2)),'}$'),strcat('EDMD$_{q/r=',num2str(q(2)/r(2)),'}$'),strcat('KEEDMD$_{q/r=',num2str(q(2)/r(2)),'}$'),...
+%             strcat('Lin at 0$_{q/r=',num2str(q(3)/r(3)),'}$'),strcat('EDMD$_{q/r=',num2str(q(3)/r(3)),'}$'),strcat('KEEDMD$_{q/r=',num2str(q(3)/r(3)),'}$')},...
+%             'location','southeast','NumColumns',2);
+%         set(LEG,'Interpreter','latex')
+        %h4.Position = h4.Position + [0 0 0 0.125]; %Modify size of 4th subplot
+        
+        h4 = subplot(3,2,4);
+        plot(t_plot,traj_d(3,:),'-b','linewidth',lw); hold on
+        for i = 2% : length(q)
+            plot(t_plot,x_track{i}(3,:),strcat(lin_typ{i},'k'),'linewidth',lw); hold on
+            plot(t_plot,x_edmd_track{i}(3,:),strcat(lin_typ{i},'r'),'linewidth',lw); hold on
+            plot(t_plot,x_koop_track{i}(3,:),strcat(lin_typ{i},'g'),'linewidth',lw); hold on
         end
-        %scatter(cent(1,:),cent(2,:),'o')
-        %axis equal
-        xaxis([-5 5])
-        yaxis([-5 5])
-        xlabel('x')
-        xlabel('y')
-        %legend('True','EDMD','Koopman e-func')
-        title('Training and simulated trajectories','interpreter','latex');
-        h3.Position = h3.Position + [0 0 0 0.125]; %Modify size of 3rd subplot
-        set(gcf, 'Position',  [0, 0, 1670, 980]) %Modify size of figure window
+            %axis([0 Tsim min(x_edmd(1,1,:))-0.15 max(x_edmd(1,1,:))+0.15])
+        xlabel('Time [s]','interpreter','latex');
+        ylabel('$x_3$','interpreter','latex');
+        set(gca,'fontsize',font_size)
+%         LEG = legend({'Reference', ...
+%             strcat('Lin at 0$_{q/r=',num2str(q(1)/r(1)),'}$'),strcat('EDMD$_{q/r=',num2str(q(1)/r(1)),'}$'),strcat('KEEDMD$_{q/r=',num2str(q(1)/r(1)),'}$'),...
+%             strcat('Lin at 0$_{q/r=',num2str(q(2)/r(2)),'}$'),strcat('EDMD$_{q/r=',num2str(q(2)/r(2)),'}$'),strcat('KEEDMD$_{q/r=',num2str(q(2)/r(2)),'}$'),...
+%             strcat('Lin at 0$_{q/r=',num2str(q(3)/r(3)),'}$'),strcat('EDMD$_{q/r=',num2str(q(3)/r(3)),'}$'),strcat('KEEDMD$_{q/r=',num2str(q(3)/r(3)),'}$')},...
+%             'location','southeast','NumColumns',2);
+%         set(LEG,'Interpreter','latex')
+        %h4.Position = h4.Position + [0 0 0 0.125]; %Modify size of 4th subplot
+        set(gcf, 'Position',  [0, 0, 1200, 600]) %Modify size of figure window
         
-        %Closed loop plot:
-        h4 = subplot(2,2,4);
-        plot(t_plot,traj_d(1,:),':b','linewidth',lw); hold on
-        plot(t_plot,x_track(1,:),'k','linewidth',lw); hold on
-        plot(t_plot,x_edmd_track(1,:),'--r','linewidth',lw); hold on
-        plot(t_plot,x_koop_track(1,:),'--g','linewidth',lw); hold on
-        %axis([0 Tsim min(x_edmd(1,1,:))-0.15 max(x_edmd(1,1,:))+0.15])
-        title('Closed loop trajectory tracking - $x_1$','interpreter','latex'); xlabel('Time [s]','interpreter','latex');
-        set(gca,'fontsize',20)
-        LEG = legend('Reference', 'Nominal controller','EDMD','Koopman e-func','location','southeast');
-        set(LEG,'interpreter','latex')
-        h4.Position = h4.Position + [0 0 0 0.125]; %Modify size of 4th subplot
+        %Prediction table:
+        set(h2,'Units','pixels');
+        cnames = {'Mean MSE', 'Std MSE'};
+        rnames = {'Linearization at 0', 'EDMD', 'KEEDMD'};
+        data = [mse_origin_avg mse_edmd_avg mse_koop_avg;...
+            mse_origin_std mse_edmd_std mse_koop_std]';
+        height = 75;
+        pos = [h2.Position(1) h2.Position(2)-height-70 h2.Position(3) height];
+        col_width = {107 107};
+        pos_title = [pos(1) pos(2)+pos(4) pos(3) 20];
+        uitable('Data', data, 'ColumnName', cnames, 'RowName', rnames,'Position', pos, 'ColumnWidth', col_width,'FontSize',12);
+        txt_title = uicontrol('Style', 'text','Position', pos_title, 'String', 'Prediction performance over 5 trajectories','FontSize',14);
+        
+        %Closed loop table:
+        set(h4,'Units','pixels');
+        cnames = {strcat('q/r=',num2str(q(1)/r(1))) strcat('q/r=',num2str(q(2)/r(2))) strcat('q/r=',num2str(q(3)/r(3)))};
+        rnames = {'Linearization at 0', 'EDMD', 'KEEDMD'};
+        data = [cost_nom{1} cost_edmd{1} cost_koop{1};...
+            cost_nom{2} cost_edmd{2} cost_koop{2};...
+            cost_nom{3} cost_edmd{3} cost_koop{3}]';
+        height = 75;
+        pos = [h4.Position(1) h4.Position(2)-height-70 h4.Position(3) height];
+        col_width = {71 71};
+        pos_title = [pos(1) pos(2)+pos(4) pos(3) 20];
+        uitable('Data', data, 'ColumnName', cnames, 'RowName', rnames,'Position', pos, 'ColumnWidth', col_width,'FontSize',12);
+        txt_title = uicontrol('Style', 'text','Position', pos_title, 'String', 'Closed loop trajectory tracking performance (MPC Cost)','FontSize',14);
     end
 end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/utils/extendedDMD.m b/scripts/Koopman-ID/utils/extendedDMD.m
index f91d7ae..af916d2 100644
--- a/scripts/Koopman-ID/utils/extendedDMD.m
+++ b/scripts/Koopman-ID/utils/extendedDMD.m
@@ -1,8 +1,8 @@
-function [A_edmd, B_edmd, C_edmd, liftFun] = extendedDMD(n,m,N_basis,basis_function,...
-    rbf_type, center_type, eps, plot_basis, xlim, ylim, Xacc, Yacc, Uacc)
+function [A_edmd, B_edmd, C_edmd, liftFun] = extendedDMD(n,m,Ntraj, Ntime, N_basis,basis_function,...
+    rbf_type, center_type, eps, Xstr, Ustr,Xf,deltaT)
     %Identify lifted state space model using Extended Dynamic Mode
-
     %Decomposition
+    
     %Inputs:
     %   n               - Number of states
     %   m               - Number of control inputs
@@ -24,61 +24,48 @@
     %   C_edmd          - Projection matrix from lifted space to outputs
 
     
+    % ************************** Prepare data *****************************
+    Xstr_shift = zeros(size(Xstr)); %Shift dynamics such that origin is fixed point
+    X = [];
+    X_dot = [];
+    U = [];
+    for i = 1 : Ntraj
+       Xstr_shift(:,i,:) = Xstr(:,i,:)-Xf(:,i);
+       X = [X reshape(Xstr_shift(:,i,:),size(Xstr_shift,1),size(Xstr_shift,3))];
+       X_dot = [X_dot num_diff(reshape(Xstr_shift(:,i,:),size(Xstr_shift,1),size(Xstr_shift,3)),deltaT)];
+       U = [U reshape(Ustr(:,i,:),size(Ustr,1),size(Ustr,3))];
+    end
+    
+    % ******************************* Lift ***********************************
+    % Find RBF Centers
     switch center_type 
         case 'random'
             cent = rand(n,N_basis)*2 - 1;
         case 'data'
-            cent = datasample(Xacc',N_basis)'+0.05*(rand(n,N_basis)*2-1);
+            cent = datasample(X',N_basis)'+0.05*(rand(n,N_basis)*2-1);
     end
-
-    % ********************** Plot basis functions *************************
-    if plot_basis
-         [X,Y] = meshgrid(-xlim(1):0.1:xlim(2),ylim(1):0.1:ylim(2));
-         x_flat = reshape(X,[1,size(X,1)*size(X,2)]);
-         y_flat = reshape(Y,[1,size(X,1)*size(X,2)]);
-         xy_flat = [x_flat;y_flat];
-
-        afigure
-        hold on
-        for i=1:1
-            rbf_flat = rbf( xy_flat,cent(:,i), rbf_type, eps_rbf);
-            rbf_pack = reshape(rbf_flat,[size(X,1),size(X,2)]);
-            surf(X,Y,rbf_pack)
-        end
-        alpha 0.2
-        xlabel('x')
-        ylabel('y')
+    
+    switch basis_function
+        case 'rbf'
+            liftFun = @(xx)( [ones(1,size(xx,2)); xx;rbf(xx,cent,rbf_type,eps)]);
     end
 
-    % ******************************* Lift ***********************************
-
-    %disp('Starting LIFT GENERATION')
-    liftFun = @(xx)( [xx;rbf(xx,cent,rbf_type,eps)] );
-    Nlift = N_basis + n;
-    Xlift = liftFun(Xacc);
-    Ylift = liftFun(Yacc);
-    %fprintf('Lifting DONE, time = %1.2f s \n', toc);
-
+    Nlift = length(liftFun(zeros(n,1)));
+    Xlift = [];
+    Xlift_dot = [];
+    for i = 1 : Ntraj
+       Xlift_temp = liftFun(reshape(Xstr_shift(:,i,:),size(Xstr_shift,1),size(Xstr_shift,3)));
+       Xlift = [Xlift Xlift_temp];
+       Xlift_dot = [Xlift_dot num_diff(Xlift_temp,deltaT)];
+    end
     % ********************** Build predictor *********************************
 
-    %disp('Starting REGRESSION')
-    %tic
-    W = [Ylift ; Xacc];
-    V = [Xlift; Uacc];
+    W = [Xlift_dot ; X];
+    V = [Xlift; U];
     VVt = V*V';
     WVt = W*V';
     M = WVt * pinv(VVt); % Matrix [A B; C 0]
     A_edmd= M(1:Nlift,1:Nlift);
     B_edmd = M(1:Nlift,Nlift+1:end);
     C_edmd = M(Nlift+1:end,1:Nlift);
-
-    %fprintf('Regression done, time = %1.2f s \n', toc);
-
-    % Plot the eigenvalues
-    % cplot = @(r,x0,y0) plot(x0 + r*cos(linspace(0,2*pi,200)),y0 + r*sin(linspace(0,2*pi,200)),'-');
-    % afigure
-    % hold on
-    % plot(eig(Alift),'*')
-    % cplot(1,0,0)
-    % axis equal
 end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/utils/find_nominal_model_ctrl.m b/scripts/Koopman-ID/utils/find_nominal_model_ctrl.m
index 295dc5d..f93fd36 100644
--- a/scripts/Koopman-ID/utils/find_nominal_model_ctrl.m
+++ b/scripts/Koopman-ID/utils/find_nominal_model_ctrl.m
@@ -1,4 +1,4 @@
-function [A_nom, B_nom, K_nom] = find_nominal_model_ctrl(n, m, f_u, Q, R)
+function [A_nom, B_nom, C_nom, K_nom] = find_nominal_model_ctrl(n, m, f_u, Q, R, x0)
     %Linearize around origin and design linear feedback control law using LQR
     
     %Inputs:
@@ -13,8 +13,8 @@
     
     
     x = sym('x',[n;1]); u = sym('u',[m;1]);
-    A_nom = double(subs(jacobian(f_u(0,x,u),x),[x;u],[0;0;0]));
-    B_nom = double(subs(jacobian(f_u(0,x,u),u),[x;u],[0;0;0]));
-
+    A_nom = double(subs(jacobian(f_u(0,x,u),x),[x;u],[x0;0]));
+    B_nom = double(subs(jacobian(f_u(0,x,u),u),[x;u],[x0;0]));
+    C_nom = eye(n);
     K_nom = -lqr(A_nom,B_nom,Q,R);
 end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/utils/koopman_eigen_id.m b/scripts/Koopman-ID/utils/koopman_eigen_id.m
index 1c328ed..504a9fa 100644
--- a/scripts/Koopman-ID/utils/koopman_eigen_id.m
+++ b/scripts/Koopman-ID/utils/koopman_eigen_id.m
@@ -1,12 +1,11 @@
 function [A_koop, B_koop, C_koop, phi_fun_v] = koopman_eigen_id(n,m, Ntraj, Ntime, ...
     N_basis,basis_function, rbf_type, center_type, eps_rbf, N_lambda, ...
-    lambda_type, A_edmd, A_nom, B_nom, K_nom, Xacc, Xstr, Xacc_c,...
+    lambda_type, A_edmd, A_nom, B_nom, K_nom, Xacc, Yacc, Xstr, Xacc_c,...
     Yacc_c, Uacc_c, Xstr_c, Ustr_c, timestr, deltaT, learn_type)
 
     %Identify lifted state space model using approximate Koopman invariant
     %subspace
 
-    %Decomposition
     %Inputs:
     %   n               - Number of states
     %   m               - Number of control inputs
@@ -42,44 +41,86 @@
     
     %N_g0 = N_basis+n;
     %N_gen = N_lambda*N_g0;
-    cent_g0 = datasample(Xacc',N_basis)'+0.05*(rand(n,N_basis)*2-1);
-    gfun = @(xx) [xx'; rbf(xx',cent_g0,rbf_type,eps_rbf)];
+%     cent_g0 = datasample(Xacc',N_basis)'+0.05*(rand(n,N_basis)*2-1);
+%     gfun = @(xx) [xx'; rbf(xx',cent_g0,rbf_type,eps_rbf)];
 
-    switch lambda_type
+%     switch lambda_type
+%         case 'random'
+%             lambda = 2*(rand(N_lambda,1)*2-1) + 2i*(rand(N_lambda,1)*2-1);
+%         case 'DMDmixing'
+%             lambda_dmd = dmd(Xstr,deltaT);
+%             lambda = zeros(N_lambda,1);
+%             lambda(1:length(lambda_dmd)) = lambda_dmd; % start the first lambda to dmd ones
+%             dlambda = N_lambda - length(lambda_dmd); % how many more to add
+%             nlambda = length(lambda_dmd); % starting lambda number
+% 
+%             vlambda = [0,1
+%                        1,0
+%                        2,0
+%                        0,2
+%                        2,1
+%                        1,2
+%                        3,1
+%                        1,3
+%                        3,2
+%                        2,3];
+% 
+%             while nlambda < N_lambda % stop when we have enough lambda  
+%                 nlambda = nlambda+1;
+%                 for k =1:length(lambda_dmd) % combination of dmd lambdas            
+%                     lambda(nlambda) = lambda(nlambda) + lambda_dmd(k)*vlambda(nlambda,k);
+%                 end
+%             end
+%         case 'eDMD'
+%             lambda =  log(eig(A_edmd))/deltaT;
+%     end
+    center_type = 'random';
+    N = Ntraj*Ntime;
+    N_basis = 2*N;
+    
+    switch center_type 
         case 'random'
-            lambda = 2*(rand(N_lambda,1)*2-1) + 2i*(rand(N_lambda,1)*2-1);
-        case 'DMDmixing'
-            lambda_dmd = dmd(Xstr,deltaT);
-            lambda = zeros(N_lambda,1);
-            lambda(1:length(lambda_dmd)) = lambda_dmd; % start the first lambda to dmd ones
-            dlambda = N_lambda - length(lambda_dmd); % how many more to add
-            nlambda = length(lambda_dmd); % starting lambda number
-
-            vlambda = [0,1
-                       1,0
-                       2,0
-                       0,2
-                       2,1
-                       1,2
-                       3,1
-                       1,3
-                       3,2
-                       2,3];
-
-            while nlambda < N_lambda % stop when we have enough lambda  
-                nlambda = nlambda+1;
-                for k =1:length(lambda_dmd) % combination of dmd lambdas            
-                    lambda(nlambda) = lambda(nlambda) + lambda_dmd(k)*vlambda(nlambda,k);
-                end
-            end
-        case 'eDMD'
-            lambda =  log(eig(A_edmd))/deltaT;
+            cent = rand(n,N_basis)*2*pi - pi;
+        case 'data'
+            cent = datasample(Xacc',N_basis)'+0.05*(rand(n,N_basis)*2-1);
     end
 
     %N_lambda = length(lambda);
-    y_reg = @(x) [x(1),x(2)];
+%     y_reg = @(x) [x(1),x(2)];
+%     
+%     [phi_fun, A_koop, C_koop, phi_grid] = get_phi_A(Xstr, timestr, lambda, gfun,y_reg);
+    
+    zfun = @(xx) [xx'; rbf(xx',cent,rbf_type,eps_rbf)];
+    X = zfun(Xacc');
+    Y = zfun(Yacc');
+    
+    to_scale = 'Scale';
+    to_center = 'Ignore';
+    k_mode = -1;
+    tol = N_basis*eps;
+    nobal = 'Ignore';
+    n_ref = 0;
+    ref_select = 1;
+    target = [0 0]'; %Not used because ref_select not equal ot 2
+    overwrite = 'OverWrite';
+    file_save = 'NoSave';
+    
+    [Z, Lambda, rez, RQ_ref, RSI, Z_ref, rez_ref, U, AxU_k] = ...
+        XY_DDMD_R4( X, Y, to_scale, to_center, k_mode, tol, nobal, n_ref,...
+        ref_select, target, overwrite, file_save );
+    
+    z_eigfun = @(xx) Z'*zfun(xx);
+    A_koop = diag(Lambda); %Discrete time autonomous dynamics;
+    Zacc = z_eigfun(Xacc');
+    
+    %Compare z with its linear evolution
+    z = z_eigfun(Xacc(:,1)');
+    for i = 1 : 1000
+        z = [z A_koop*z(:,end)];
+    end
+        
     
-    [phi_fun, A_koop, C_koop, phi_grid] = get_phi_A(Xstr, timestr, lambda, gfun,y_reg);
+    C_koop = (Zacc/Xacc)';
     fprintf('Learning autonomous dynamics done, execution time: %1.2f s \n', toc);
 
     
diff --git a/scripts/Koopman-ID/utils/learn_B_koop.m b/scripts/Koopman-ID/utils/learn_B_koop.m
index cf17c2d..d2174e6 100644
--- a/scripts/Koopman-ID/utils/learn_B_koop.m
+++ b/scripts/Koopman-ID/utils/learn_B_koop.m
@@ -1,4 +1,4 @@
-function [B_koop, phi_fun_v] = learn_B_koop(n, m, Ntraj, Ntime, phi_fun, A_koop,...
+function [B_koop, phi_fun_v] = learn_B_koop(n, m, Ntraj, Ntime, phi_fun_v, A_koop,...
                     C_koop, Xstr_c, Ustr_c, Xacc_c, Yacc_c, Uacc_c, deltaT, learn_type)
 
     % Setup regression problem min ||Q * vec(Blift) - b||_2^2 corresponding to
@@ -12,49 +12,103 @@
     %Outputs:
     
     
-    A_koop_d = expm(A_koop*deltaT); %Discretize A_koop
-    phi_fun_v = @(x) phiFunction(phi_fun,x);
+    %phi_fun_v = @(x) phiFunction(phi_fun,x);
     
     switch learn_type
         case 'multi-step'
-            Obs =  obsv(A_koop_d,C_koop); % Not an efficient implementation
+%             Obs =  obsv(A_koop_d,C_koop); % Not an efficient implementation
+%             nc = size(C_koop,1);
+%             b = zeros(Ntime*nc,Ntraj);
+%             Q = zeros(Ntraj*Ntime*nc,size(A_koop_d,1)*m);
+%             X0 = Xstr_c(:,:,1);
+% 
+%             Obsx0 = Obs*phi_fun_v(X0);
+%             for j = 1:Ntime
+%                 b(nc*(j-1)+1:nc*j,:) = Obsx0( j*nc + 1 : (j+1)*nc, : ) - Xstr_c(:,:,j+1);
+%                 tmp = 0;
+%                 for k = 0 : j-1
+%                     kprime = j - k -1;
+%                     tmp = tmp + kron(Ustr_c(:,:,k+1)',Obs(kprime*nc + 1 : (kprime+1)*nc,:));
+%                 end        
+%                 Q(Ntraj*nc*(j-1)+1:Ntraj*nc*(j),:) = tmp;
+%             end
+%             b = reshape(-b,Ntraj*Ntime*nc,1);
+% 
+%             B_koop = pinv(Q)*b; % = vec(Blift) %Add regularization
+%             B_koop = reshape(B_koop,size(A_koop_d,1),m); % Unvectorize
+            
+            % Regression for B (Milan Korda implementation)
+            % Setup regression problem min ||Q * vec(Blift) - b||_2^2 corresponding to
+            % min_{Blift} \sum_{j=1}^{Nsim} ||xpred - xtrue||^2 where
+            % xpred = C*Alift^j*x0_lift + sum_{k=0}^{j-1} kron(u_k',C*Alift^{j-k-1})*vec(BLift)
+            % (And this sum over all trajectories)
+            Obs =  obsvk(sparse(A_koop),C_koop,Ntime+1); % Not an efficient implementation
+            b = [];
             nc = size(C_koop,1);
-            b = zeros(Ntime*nc,Ntraj);
-            Q = zeros(Ntraj*Ntime*nc,size(A_koop_d,1)*m);
-            X0 = Xstr_c(:,:,1);
-
-            Obsx0 = Obs*phi_fun_v(X0);
-            for j = 1:Ntime
-                b(nc*(j-1)+1:nc*j,:) = Obsx0( j*nc + 1 : (j+1)*nc, : ) - Xstr_c(:,:,j+1);
-                tmp = 0;
-                for k = 0 : j-1
-                    kprime = j - k -1;
-                    tmp = tmp + kron(Ustr_c(:,:,k+1)',Obs(kprime*nc + 1 : (kprime+1)*nc,:));
-                end        
-                Q(Ntraj*nc*(j-1)+1:Ntraj*nc*(j),:) = tmp;
+            Q = zeros(Ntraj*Ntime*nc,size(A_koop,1)*m);
+            for q = 1 : Ntraj
+                fprintf('Building regression matrices for B: %f percent complete \n', 100*q / Ntraj)
+                x0 = reshape(Xstr_c(:,q,1),n,1);
+                Obsx0 = Obs*phi_fun_v(x0);
+                for j = 1:Ntime
+                    b = [b ; Obsx0( j*nc + 1 : (j+1)*nc, : ) - reshape(Xstr_c(:,q,j+1),n,1)] ;
+                    tmp = 0;
+                    for k = 0 : j-1
+                        kprime = j - k - 1;
+                        tmp = tmp + kron(reshape(Ustr_c(:,q,k+1),m,1)',Obs(kprime*nc + 1 : (kprime+1)*nc,:));
+                    end
+                    Q((q-1)*Ntime*nc + (j-1)*nc + 1 : (q-1)*Ntime*nc + j*nc,:) = tmp;
+                end
             end
-            b = reshape(-b,Ntraj*Ntime*nc,1);
+            b = -b;
+
+            k = 1e-3;
+            B_koop = Q\b; % = vec(Blift)
+            %B_koop = (Q'*Q + k*eye(size(Q,2)))\(Q'*b); %Regularize B-learning
+            %B_koop = lasso(Q,b,'Lambda',1e-2);
+            B_koop = reshape(B_koop,size(A_koop,1),m); % Unvectorize
 
-            B_koop = pinv(Q)*b; % = vec(Blift) %Add regularization
-            B_koop = reshape(B_koop,size(A_koop_d,1),m); % Unvectorize
         case 'single-step'
             %Calculate residual after autonomous dynamics are subtracted
             Z_x = phi_fun_v(Xacc_c);
             Z_y = phi_fun_v(Yacc_c);
-            unusable_data_x = find(sum(isnan(Z_x),1) > 0);
-            unusable_data_y = find(sum(isnan(Z_y),1) > 0);
-            unusable_data = unique([unusable_data_x unusable_data_y]);
-            Z_x(:,unusable_data) = []; %Delete unsable elements
-            Z_y(:,unusable_data) = []; %Delete unsable elements
-            Uacc_c(:,unusable_data) = []; %Delete unsable elements
-
-            Racc = Z_y - A_koop_d*Z_x; %Calculate residual when subtracting autonomous dynamics
 
-            %Racc = phi_fun_v(Yacc_c) - A_koop*phi_fun_v(Xacc_c); 
-            B_koop = (Uacc_c'\Racc')';
+            Racc = Z_y-A_koop*Z_x; %Calculate residual when subtracting autonomous dynamics
+            
+            k = 10;
+            %B_koop = (Uacc_c'\Racc')';
+            B_koop = ((Uacc_c*Uacc_c' + k*eye(size(Uacc_c,1)))\(Uacc_c*Racc'))';
             if sum(sum(isnan(B_koop))) > 0
                 disp('Warning: NAN-values in B-matrix, terminating.')
                 return
             end
+        case 'multi-step_mod'
+            %Calculate residual after autonomous dynamics are subtracted
+%             Racc = [];
+%             Z_evol = zeros(length(phi_fun_v(zeros(n,1))),Ntime+1);
+%             for i = 1 : Ntraj
+%                 Z_evol(:,1) = phi_fun_v(reshape(Xstr_c(:,i,1),n,1));
+%                 Z_y = phi_fun_v(reshape(Xstr_c(:,i,2:end),n,Ntime));
+%                 for j = 2 : Ntime+1
+%                     Z_evol(:,j) = A_koop*Z_evol(:,j-1);
+%                 end
+%                 Racc = [Racc Z_evol(:,2:end)-Z_y]; 
+%             end
+%                     
+%             B_koop = (Uacc_c'\Racc')';
+            Racc = [];
+            U_r_acc = [];
+            Z_evol = zeros(length(phi_fun_v(zeros(n,1))),Ntime+1);
+            for i = 1 : Ntraj
+                Z_evol(:,1) = phi_fun_v(reshape(Xstr_c(:,i,1),n,1));
+                for j = 2 : Ntime+1
+                    Z_evol(:,j) = A_koop*Z_evol(:,j-1);
+                end
+                Yres = C_koop*Z_evol(:,2:end)-reshape(Xstr_c(:,i,2:end),n,Ntime) ;
+                Racc = [Racc phi_fun_v(Yres)]; 
+                U_r_acc = [U_r_acc reshape(Ustr_c(:,i,:),m,Ntime)];
+            end
+                    
+            B_koop = -(U_r_acc'\Racc')';
     end
 end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/utils/sim_closed_loop_mpc.m b/scripts/Koopman-ID/utils/sim_closed_loop_mpc.m
index bca6c13..d592132 100644
--- a/scripts/Koopman-ID/utils/sim_closed_loop_mpc.m
+++ b/scripts/Koopman-ID/utils/sim_closed_loop_mpc.m
@@ -1,25 +1,11 @@
-function [x_nom,x_edmd,x_koop,mse_nom,mse_edmd,mse_koop, t_plot, traj_d, E_nom, E_edmd, E_koop]...
+function [x_nom,x_edmd,x_koop,mse_nom,mse_edmd,mse_koop, q_traj, E_nom, E_edmd, E_koop,...
+            cost_nom, cost_edmd, cost_koop]...
             = sim_closed_loop_mpc(n,m,n_edmd,n_koop,Nsim,Ntime,deltaT,Tsim,...
-            f_u,liftFun,phi_fun_v, K_nom,A_edmd,B_edmd,C_edmd,A_koop_d,...
-            B_koop,C_koop)
-
-    %Set up trajectory to track:
-    t = 0 : deltaT : Tsim;
-    t1 = t(1:floor(Ntime/3));
-    t2 = t(floor(Ntime/3):floor(2*Ntime/3));
-    t3 = t(floor(2*Ntime/3):end);
-    t_plot = [t1 t2 t3];
-    x0 = -pi/2; %Start angle
-    x1 = 0.2; %Waypoint angle
-    x2 = -0.2; %Waypoint angle
-    x3 = pi/3; %Final angle
-    traj_d = [spline([t1(1) t1(end)], [0 [x0 x1] 0], t1) ...
-                spline([t2(1) t2(end)], [0 [x1 x2] 0], t2) ...
-                spline([t3(1) t3(end)], [0 [x2 x3] 0], t3)];
-    traj_d = [traj_d; diff([traj_d traj_d(end)])/deltaT];
+            f_u,liftFun,phi_fun_v,A_nom,B_nom,C_nom,K_nom,A_edmd,B_edmd,C_edmd,A_koop,...
+            B_koop,C_koop,Q,R,q_traj)
     
     %Initialize variables:
-    Ntime_track = size(traj_d,2)-1;
+    Ntime_track = size(q_traj,2)-1;
     x_nom = zeros(n,Ntime_track+1);
     u_nom = zeros(m,Ntime_track);
     z_edmd = zeros(n_edmd,Ntime_track+1);
@@ -28,36 +14,24 @@
     z_koop = zeros(n_koop,Ntime_track+1);
     x_koop = zeros(n,Ntime_track+1);
     u_koop = zeros(m,Ntime_track);
-    x_nom(:,1) = traj_d(:,1);
-    x_edmd(:,1) = traj_d(:,1);
-    x_koop(:,1) = traj_d(:,1);
-    z_edmd(:,1) = liftFun(traj_d(:,1));
-    z_koop(:,1) = phi_fun_v(traj_d(:,1));
+    x_nom(:,1) = q_traj(:,1);
+    x_edmd(:,1) = q_traj(:,1);
+    x_koop(:,1) = q_traj(:,1);
+    z_edmd(:,1) = liftFun(q_traj(:,1));
+    z_koop(:,1) = phi_fun_v(q_traj(:,1));
 
-    % Define Koopman controller
-    Q = eye(2); % Weight matrices
-    R = 0.01;
-    Tpred = 0.1; % Prediction horizon
+    % Define MPC parameters:
+    Tpred = Tsim;%0.1; % Prediction horizon
     Np = round(Tpred / deltaT);
-    traj_d = [traj_d traj_d(:,end).*ones(n,Np)]; %Add extra points to trajectory to allow prediction horizon
+    q_traj = [q_traj q_traj(:,end).*ones(n,Np)]; %Add extra points to trajectory to allow prediction horizon
     xlift_min = [];%[ymin ; nan(Nlift-1,1)]; % Constraints
     xlift_max = [];%[ymax ; nan(Nlift-1,1)];
 
     % Build Koopman MPC controller 
+    MPC_nom  = getMPC(A_nom,B_nom,C_nom,0,Q,R,Q,Np,-100, 100, xlift_min, xlift_max,'qpoases');
     MPC_edmd  = getMPC(A_edmd,B_edmd,C_edmd,0,Q,R,Q,Np,-100, 100, xlift_min, xlift_max,'qpoases');
-    MPC_koop  = getMPC(A_koop_d,B_koop,C_koop,0,Q,R,Q,Np,-100, 100, xlift_min, xlift_max,'qpoases');
+    MPC_koop  = getMPC(A_koop,B_koop,C_koop,0,Q,R,Q,Np,-100, 100, xlift_min, xlift_max,'qpoases');
     
-    % Get Jacobian of the true dynamics (for local linearization MPC)
-    x = sym('x',[2 1]); syms u;
-    f_ud_sym = sim_timestep(deltaT, f_u, 0, x, u);
-    Jx = jacobian(f_ud_sym,x);
-    Ju = jacobian(f_ud_sym,u);
-
-    %Save indices of infeasible MPC-problems (Should never happen with no
-    %state constraints)
-    wasinfeas= 0;
-    ind_inf = [];
-
     % Closed-loop simultion start
     for i = 1:Ntime_track
         if(mod(i,10) == 0)
@@ -65,18 +39,10 @@
         end
 
         % Prediction horizon of reference signal
-        yr = traj_d(:,i:i+Np-1);
+        yr = q_traj(:,i:i+Np-1);
 
         % Local linearization MPC
-        Aloc = double(subs(Jx,[x;u],[x_nom(:,i);u_nom(:,i)])); % Get local linearization
-        Bloc = double(subs(Ju,[x;u],[x_nom(:,i);u_nom(:,i)]));
-        Cloc = double(subs(f_ud_sym,[x;u],[x_nom(:,i);u_nom(:,i)])) - Aloc*x_nom(:,i) - Bloc*u_nom(:,i);
-        [U_nom,~,optval] = solveMPCprob(Aloc,Bloc,eye(2),[],Q,R,Q,Np,-100, 100,[],[],x_nom(:,i),yr); % Get control input
-        u_nom(:,i) = U_nom(1:m,1);
-        if(optval == Inf) % Detect infeasibility
-            ind_inf = [ind_inf i];
-            wasinfeas = 1;
-        end
+        u_nom(:,i) = MPC_nom(x_nom(:,i),yr);
         x_nom(:,i+1) = sim_timestep(deltaT, f_u, 0, x_nom(:,i), u_nom(:,i)); % Update true state
 
         % EDMD MPC
@@ -85,24 +51,20 @@
         x_edmd(:,i+1) = sim_timestep(deltaT, f_u, 0, x_edmd(:,i), u_edmd(:,i));
         
         % Koopman e-func MPC
-%         z_koop(:,i) = phi_fun_v(x_koop(:,i)); % Lift
-%         disp(sum(sum(isnan(z_koop(:,i)))));
-%         u_koop(:,i) = MPC_koop(z_koop(:,i),yr) + K_nom*x_koop(:,i); % Get control input
-%         x_koop(:,i+1) = sim_timestep(deltaT, f_u, 0, x_koop(:,i), u_koop(:,i));
-        
-        disp([u_nom(:,i), u_edmd(:,i) u_koop(:,i)])
-    end
-
-    if(isempty(ind_inf))
-        ind_inf = Ntime_track;
+        z_koop(:,i) = phi_fun_v(x_koop(:,i)); %Lift
+        u_koop(:,i) = MPC_koop(z_koop(:,i),yr);% Get control input
+        x_koop(:,i+1) = sim_timestep(deltaT, f_u, 0, x_koop(:,i), u_koop(:,i));
     end
     
     % Calculate corresponding predictions and MSE
-    traj_d = traj_d(:,1:end-Np); %Remove added elements from trajectory
-    mse_nom = immse(traj_d,x_nom);
-    mse_edmd = immse(traj_d,x_edmd);
-    mse_koop = immse(traj_d,x_koop);
+    q_traj = q_traj(:,1:end-Np); %Remove added elements from trajectory
+    mse_nom = immse(q_traj,x_nom);
+    mse_edmd = immse(q_traj,x_edmd);
+    mse_koop = immse(q_traj,x_koop);
     E_nom = norm(u_nom);
     E_edmd = norm(u_edmd);
     E_koop = norm(u_koop);
+    cost_nom = sum(diag((x_nom(:,2:end)-q_traj(:,2:end))'*Q*(x_nom(:,2:end)-q_traj(:,2:end)))) + sum(diag(u_nom'*R*u_nom));
+    cost_edmd = sum(diag((x_edmd(:,2:end)-q_traj(:,2:end))'*Q*(x_edmd(:,2:end)-q_traj(:,2:end)))) + sum(diag(u_edmd'*R*u_edmd));
+    cost_koop = sum(diag((x_koop(:,2:end)-q_traj(:,2:end))'*Q*(x_koop(:,2:end)-q_traj(:,2:end)))) + sum(diag(u_koop'*R*u_koop));
 end
\ No newline at end of file
diff --git a/scripts/Koopman-ID/utils/sim_prediction.m b/scripts/Koopman-ID/utils/sim_prediction.m
index 63d0e78..6471961 100644
--- a/scripts/Koopman-ID/utils/sim_prediction.m
+++ b/scripts/Koopman-ID/utils/sim_prediction.m
@@ -1,46 +1,62 @@
-function [x,x_edmd,x_koop,mse_edmd_avg,mse_koop_avg,mse_edmd_std,mse_koop_std]...
-            = sim_prediction(n,m,n_edmd,n_koop,Nsim,Ntime,deltaT,X0,f_u,liftFun,...
-            phi_fun_v, K_nom,A_edmd,B_edmd,C_edmd,A_koop_d,B_koop,C_koop)
+function [x,x_origin, x_edmd,x_koop, mse_origin_avg, mse_edmd_avg,mse_koop_avg,...
+    mse_origin_std, mse_edmd_std,mse_koop_std]...
+            = sim_prediction(n,m,n_edmd,n_koop,Nsim,Ntime,deltaT,q_traj,f_u,liftFun,...
+            phi_fun_v,A_nom,B_nom,C_nom,K_nom,A_edmd,B_edmd,C_edmd,A_koop,B_koop,C_koop)
         
-    u = 5*rand(m,Nsim,Ntime)-2.5;
+    u = zeros(m,Nsim,Ntime);
+    
+    % Define nominal and KEEDMD model:
+    f_nom = @(t,x,u) A_nom*x + B_nom*u;
+    f_edmd = @(t,x,u) A_edmd*x + B_edmd*u;
+    f_koop = @(t,x,u) A_koop*x + B_koop*u;
     
     x = zeros(n,Nsim,Ntime+1);
+    x_origin = zeros(n,Nsim,Ntime+1);
     z_edmd = zeros(n_edmd,Nsim,Ntime+1);
     z_koop = zeros(n_koop,Nsim,Ntime+1);
+    
+    X0 = q_traj(:,:,1);
     x(:,:,1) = X0;
+    x_origin(:,:,1) = X0;
     z_edmd(:,:,1) = liftFun(X0);
     z_koop(:,:,1) = phi_fun_v(X0);
     
-    
-
     % Simulate all systems and initial points
-    K_mod = K_nom + randn(size(K_nom));
     for i = 1:Ntime
+        %Controller:
+        u(:,:,i) = K_nom*(x(:,:,i)-q_traj(:,:,i)) + 5*rand()-2.5;
+        
         %True dynamics:
-        x(:,:,i+1) = sim_timestep(deltaT,f_u,0,x(:,:,i), K_mod*x(:,:,i)+u(:,:,i));
+        x(:,:,i+1) = sim_timestep(deltaT,f_u,0,x(:,:,i), u(:,:,i));
+        
+        %Linearization at origin:
+        x_origin(:,:,i+1) = sim_timestep(deltaT,f_nom,0,x_origin(:,:,i), u(:,:,i));
         
         %EDMD predictor:
-        z_edmd(:,:,i+1) = A_edmd*z_edmd(:,:,i) + B_edmd*(K_mod*C_edmd*z_edmd(:,:,i) + u(:,:,i));
+        z_edmd(:,:,i+1) = sim_timestep(deltaT,f_edmd,0,z_edmd(:,:,i), u(:,:,i));
 
         %Koopman eigenfunction predictor:
-        z_koop(:,:,i+1) = A_koop_d*z_koop(:,:,i) + ...
-            B_koop*(K_mod*C_koop*z_koop(:,:,i)+u(:,:,i)-K_nom*C_koop*z_koop(:,:,i));
+        z_koop(:,:,i+1) = sim_timestep(deltaT,f_koop,0,z_koop(:,:,i), u(:,:,i));
     end
     
     % Calculate corresponding predictions and MSE
     x_edmd = zeros(n,Nsim,Ntime+1);
     x_koop = zeros(n,Nsim,Ntime+1);
+    mse_origin = zeros(Nsim,1);
     mse_edmd = zeros(Nsim,1);
     mse_koop = zeros(Nsim,1);
     for i = 1 : Nsim
         x_edmd(:,i,:) = C_edmd * reshape(z_edmd(:,i,:),n_edmd,Ntime+1); %EDMD predictions
         x_koop(:,i,:) = C_koop * reshape(z_koop(:,i,:),n_koop,Ntime+1); %Koopman eigenfunction predictions
+        mse_origin(i) = immse(reshape(x(:,i,:),n,Ntime+1), reshape(x_origin(:,i,:),n,Ntime+1));
         mse_edmd(i) = immse(reshape(x(:,i,:),n,Ntime+1), reshape(x_edmd(:,i,:),n,Ntime+1));
         mse_koop(i) = immse(reshape(x(:,i,:),n,Ntime+1), reshape(x_koop(:,i,:),n,Ntime+1));
     end
     
+    mse_origin_avg = mean(mse_origin);
     mse_edmd_avg = mean(mse_edmd);
     mse_koop_avg = mean(mse_koop);
+    mse_origin_std = std(mse_origin);
     mse_edmd_std = std(mse_edmd);
     mse_koop_std = std(mse_koop);
 end
\ No newline at end of file