diff --git a/functions/optimal_bound_2d.m b/functions/optimal_bound_2d.m
new file mode 100644
index 0000000000000000000000000000000000000000..a7e85f434012856af9e505f4fdf590749af5dab3
--- /dev/null
+++ b/functions/optimal_bound_2d.m
@@ -0,0 +1,155 @@
+function V = optimal_bound_2d(A,q11,q22,J,x1vec,x2vec,h,plot,print)
+
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+%
+% Written by: Bo Bernhardsson and Marcus Thelander Andrén, 2017
+%
+% This function numerically computes the value function V,
+% that determines the optimal reset bound. The idea is to use the
+% simularity to the diffusion-convection equation, and simulate it to
+% stationarity while at each time-step enforcing an inequality on V.
+%
+% The partial-differential-equation we simulate is:
+%
+% (x'Qx - J) + x'A'grad(V) + 1/2laplace(V) = dV/dt
+% "Production" "Convection" "Diffusion"
+%
+% while at each time step we enforce the inequality:
+%
+% 0 >= V(x) (assumes V(0) = rho)
+%
+% Without loss of generality, the system is assumed to be of the form
+% Q = diag(q11, q22) and R = I (the system state can always be transformed
+% to achieve this).
+%
+% The integration is performed numerically using the Finite Difference
+% Method. We use the implicit scheme, with a first-order upwind in time
+% approximation, and a central difference in space approximation.
+%
+% For more information, see the related article:
+%
+% M. Thelander Andrén, B. Bernhardsson, A. Cervin and K. Soltesz,
+% "On Event-Based Sampling for LQG-Optimal Control" in Proc. 56th IEEE
+% Conf. on Decision and Control, Melbourne, Australia, 2017
+%
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+
+%A-matrix
+a11 = A(1,1);
+a12 = A(1,2);
+a21 = A(2,1);
+a22 = A(2,2);
+
+%Gridding (NOTE: Assuming equidistant grid!)
+dx1 = abs(x1vec(2)-x1vec(1));
+dx2 = abs(x1vec(2)-x1vec(1));
+nx1 = length(x1vec);
+nx2 = length(x2vec);
+
+%Grid up relative value function V
+V = zeros(nx1,nx2);
+
+%Make grid of state space
+[X1,X2] = meshgrid(x1vec, x2vec);
+X1 = X1';
+X2 = X2'; %To get it in the right direction
+X1cut = X1(2:end-1,2:end-1);
+X2cut = X2(2:end-1,2:end-1);
+
+%"Production" term at each (i,j)-location for duration h
+prod = h*(q11*X1.^2 + q22*X2.^2 - J);
+prod = prod(2:end-1,:);
+prod = prod(:,2:end-1);
+prodvec = reshape(prod,[],1);
+
+%Dirichlet Boundary condition vector (V(x) = 0 at boundary)
+bc = zeros(nx2-2,nx1-2);
+
+% Coefficient matrices for the implicit scheme
+Ex1 = sparse(2:nx1-2,1:nx1-3,1,nx1-2,nx1-2);
+Ax1 = Ex1+Ex1'-2*speye(nx1-2);
+Ex2 = sparse(2:nx2-2,1:nx2-3,1,nx2-2,nx2-2);
+Ax2 = Ex2+Ex2' - 2*speye(nx2-2);
+Amat = kron(Ax2/dx2^2,speye(nx1-2)) + kron(speye(nx2-2),Ax1/dx1^2);
+D = speye((nx1-2)*(nx2-2))-0.5*h*Amat;
+
+diffx1 = Ex1'-Ex1;
+diffx2 = Ex2'-Ex2;
+Dx1 = kron(speye(nx2-2),diffx1/(2*dx1));
+Dx2 = kron(diffx2/(2*dx2),speye(nx1-2));
+
+fx1 = reshape(a11*X1cut + a12*X2cut,[],1);
+fx2 = reshape(a21*X1cut + a22*X2cut,[],1);
+
+%Setup simulation parameters
+maxiter = 5e4;
+rel_stop_crit = 1e-4*h;
+
+%Setup plot
+if plot
+ figure()
+ colormap('jet')
+ %Optimal bound
+ subplot(1,2,2)
+ %Suppress warning of constant Z-data
+ warning('off','MATLAB:contour:ConstantData')
+ [~, h1] =contour(X1,X2,V,[0 0],'b','ZDataSource','V');
+ axis([x1vec(1) x1vec(end) x2vec(1) x2vec(end)])
+ xlabel('x_1')
+ ylabel('x_2')
+ title('Optimal Reset Bound')
+ grid on
+ axis equal
+
+ %Surface plot of V
+ subplot(1,2,1)
+ h2 = surf(X1,X2,V,'EdgeColor','none','ZDataSource','V');
+ shading interp
+ xlabel('x_1')
+ ylabel('x_2')
+ zlabel('V')
+ title('Value Function')
+ axis tight
+end
+
+%Main loop
+for it=1:maxiter
+ Vn = V;
+
+ %Time-step
+ v = Vn;
+ v = v(2:end-1,2:end-1);
+ v = reshape(v+bc,[],1);
+ v = D\(v + prodvec + h*fx1.*(Dx1*v) + h*fx2.*(Dx2*v));
+ v = reshape(v,nx2 - 2, nx1 - 2);
+ V(2:nx2-1,2:nx1-1) = v;
+
+ %Enforce boundary conditions and V <= 0
+ V(1,:) = 0;
+ V(nx2,:) = 0;
+ V(:,1) = 0;
+ V(:,nx1) = 0;
+ V = min(V,0);
+
+ %Check for steady-state
+ Vdiff = abs(V - Vn);
+ maxVdiff = max(max(Vdiff));
+ if maxVdiff < rel_stop_crit
+ disp('Stopping criteria met')
+ return
+ end
+
+ %Print and plot
+ if mod(it,10) == 0
+ if print
+ disp(['Iterations : ' num2str(it),' out of ',num2str(maxiter), ' | Maximum Vdiff : ', num2str(maxVdiff),' V(0) = ', num2str(min(min(V)))])
+ end
+ if plot
+ refreshdata(h1, 'caller')
+ refreshdata(h2, 'caller')
+ drawnow
+ end
+ end
+end
+disp('Maximum number of iterations reached')
+end
diff --git a/optimal_bound_2d.m b/optimal_bound_2d.m
deleted file mode 100644
index 785956d31f141fe3df973e3b7ce6e9f7480c20cb..0000000000000000000000000000000000000000
--- a/optimal_bound_2d.m
+++ /dev/null
@@ -1,184 +0,0 @@
-%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
-%
-%Authors: Bo Bernhardsson and Marcus Thelander Andrén, 2017.
-%
-%Script for numerically finding the value function V,
-%which determines the optimal trigger bound. The idea is to use the
-%simularity to the diffusion-convection equation, and simulate it to
-%stationarity while at each time-step enforcing an inequality on V.
-
-%The PDE we simulate is:
-%
-% (x'Qx - J) + x'A'grad(V) + 1/2laplace(V) = dV/dt
-% "Production" "Convection" "Diffusion"
-%
-%while at each time step we enforce the inequality:
-%
-% 0 >= V(x,t) (assumes V(0,t) = rho)
-
-%Here we assume that the state cost matrix Q=diag(...) and the noise
-%intensity matrix R=I.
-%Through linear transformations, we can always transform the system to this
-%setting, i.e., there is no loss of generality.
-
-%The integration is performed numerically using the Backward Time Central
-%Space (BTCS) Finite Difference method.
-%
-%
-%The grid in space is defined as:
-%
-%
-%
-% x2|^ *(i,j+1)
-% |
-% | *(i-1,j) *(i,j) *(i+1,j)
-% |_______________________________> x1
-%
-%
-%
-%For more information, please see the related submission:
-%
-% M.T Andén, B. Bernhardsson, A. Cervin and K. Soltesz, "On Event-Based
-%Sampling for H2-Optimal Control", Submitted to 56th Conf. on Decision and
-%Control, 2017.
-%
-%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
-%%
-clear all
-close all
-%% Define Parameters
-
-%A-matrix
-a11 = 1;
-a12 = -1;
-a21 = 0;
-a22 = -70;
-Amat = [a11 a12;a21 a22];
-
-%Q-matrix and J
-q11 = 1;
-q22 = 1;
-J = 1;
-
-%State-space grid parameters
-x1min = -2;
-x1max = 2;
-x2min = -2;
-x2max = 2;
-
-nx1 = 200;
-nx2 = 200;
-dx1 = (x1max-x1min)/(nx1-1);
-dx2 = (x2max-x2min)/(nx2-1);
-
-x1 = x1min:dx1:x1max;
-x2 = x2min:dx2:x2max;
-
-%Time discretization and stopping criteria
-dt = 2.5e-4;
-maxiter = 5e4; %Stop when # of time-steps = maxiter
-rel_stop_crit = 1e-1*dt; %Stop when abs(max(V(t+1) - V(t)) <= rel_stop_crit
-
-%% Gridding
-
-%Grid up solution V
-V = zeros(nx2,nx1); %Initial condition V(x,0) = 0
-Vn = zeros(nx2,nx1);
-
-%Make grid of state space
-[X1,X2] = meshgrid(x1,x2);
-X2=flipud(X2);
-X1cut = X1(2:end-1,2:end-1);
-X2cut = X2(2:end-1,2:end-1);
-
-%"Production" term at each (i,j)-location for duration dt
-prod = dt*(q11*X1.^2 + q22*X2.^2 - J);
-prod(1,:) = [];
-prod(end,:) = [];
-prod(:,1) = [];
-prod(:,end) = [];
-prodvec = reshape(prod,[],1);
-
-%Dirichlet Boundary condition vector (V(x) = 0 at boundary)
-bc = zeros(nx2-2,nx1-2);
-
-%% Coefficient matrices for BTCS scheme
-Ex1 = sparse(2:nx1-2,1:nx1-3,1,nx1-2,nx1-2);
-Ax1 = Ex1+Ex1'-2*speye(nx1-2);
-Ex2 = sparse(2:nx2-2,1:nx2-3,1,nx2-2,nx2-2);
-Ax2 = Ex2+Ex2' - 2*speye(nx2-2);
-A = kron(Ax2/dx2^2,speye(nx1-2)) + kron(speye(nx2-2),Ax1/dx1^2);
-D = speye((nx1-2)*(nx2-2))-0.5*dt*A;
-
-diffx1 = Ex1'-Ex1;
-diffx2 = Ex2'-Ex2;
-Dx1 = kron(speye(nx2-2),diffx1/(2*dx1));
-Dx2 = kron(diffx2/(2*dx2),speye(nx1-2));
-
-fx1 = reshape(a11*X1cut + a12*X2cut,[],1);
-fx2 = reshape(a21*X1cut + a22*X2cut,[],1);
-
-%% Simulating the distribution to steady-state
-maxVdiff = inf;
-Vplot = V;
-
-%Plotting
-h1=figure(1);
-
-subplot(1,2,1)
-h2=surf(x1,x2,Vplot,'EdgeColor','none','ZDataSource','Vplot');
-shading interp
-title('Value function V(x)')
-axis([x1min x1max x2min x2max -0.15 0.05])
-axis square
-xlabel('x1')
-ylabel('x2')
-zlabel('V(x)')
-
-subplot(1,2,2)
-[~,h3]=contour(x1,x2,Vplot,[0 0],'b','ZDataSource','Vplot');
-axis([x1min x1max x2min x2max])
-axis square
-title({['Optimal Trigger Bound with A = ', mat2str(Amat), ', Q = diag(',num2str(q11),',',num2str(q22),')']})
-xlabel('x1')
-ylabel('x2')
-grid on
-
-for it=0:maxiter
- Vn = V;
- if mod(it,10) == 0
- disp(['Iterations : ' num2str(it),' out of ',num2str(maxiter), ' | Maximum Vdiff : ', num2str(maxVdiff)])
-
- %Update plot
- refreshdata(h2)
- refreshdata(h3)
- drawnow;
- end
-
- %Time-step
- v = Vn;
- v(1,:) = [];
- v(end,:) = [];
- v(:,1) = [];
- v(:,end) = [];
- v = reshape(v+bc,[],1);
- v = D\(v + prodvec + dt*fx1.*(Dx1*v) + dt*fx2.*(Dx2*v));
- v = reshape(v,nx2 - 2, nx1 - 2);
- V(2:nx2-1,2:nx1-1) = v;
-
- %Enforce boundary conditons
- V(1,:) = 0;
- V(nx2,:) = 0;
- V(:,1) = 0;
- V(:,nx1) = 0;
- V = min(V,0);
- Vplot=V;
-
- %Check for steady-state
- Vdiff = abs(V - Vn);
- maxVdiff = max(max(Vdiff));
- if maxVdiff <= rel_stop_crit
- disp('Stopping Criteria Met')
- break;
- end
-end
diff --git a/script_optimal_bound_2d.m b/script_optimal_bound_2d.m
new file mode 100644
index 0000000000000000000000000000000000000000..1a75203f1a6fcb3e681f23d218f2f578b18f027f
--- /dev/null
+++ b/script_optimal_bound_2d.m
@@ -0,0 +1,78 @@
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+%
+% Written by: Marcus Thelander Andrén, 2017
+%
+% This script demonstrates for the 2D case how to numerically compute the
+% value function V that determines the optimal reset bound.
+%
+% Without loss of generality, the system is assumed to be of the form
+% Q = diag(q11, q22) and R = I (the system state can always be transformed
+% to achieve this).
+%
+% For more information, see the related article:
+%
+% M. Thelander Andrén, B. Bernhardsson, A. Cervin and K. Soltesz,
+% "On Event-Based Sampling for LQG-Optimal Control" in Proc. 56th IEEE
+% Conf. on Decision and Control, Melbourne, Australia, 2017
+%
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+
+%% Preamble
+clear all
+close all
+addpath('./functions')
+
+%% Setup parameters
+
+%Define system
+A = [-1 20; 20 -1];
+q11 = 1;
+q22 = 1;
+J=1;
+
+%Gridding
+dx = 3e-2;
+x1min=-2;
+x1max = 2;
+x2min= -2;
+x2max = 2;
+x1vec = x1min:dx:x1max;
+x2vec = x2min:dx:x2max;
+h = 5e-4;
+
+%Enable/disable plotting and printing
+plot_progress = 0;
+print_progress = 1;
+plot_result = 1;
+
+%% Compute V
+disp('Computing V...')
+V = optimal_bound_2d(A, q11, q22, J, x1vec, x2vec, h,...
+ plot_progress, print_progress);
+disp('Done computing V')
+
+%% Plot the end result
+if plot_result && ~plot_progress
+ figure()
+ colormap('jet')
+
+ %Surface plot of V
+ subplot(1,2,1)
+ surf(x1vec,x2vec,V,'EdgeColor','none');
+ shading interp
+ xlabel('x_1')
+ ylabel('x_2')
+ zlabel('V')
+ title('Value Function')
+ axis tight
+
+ %Optimal bound
+ subplot(1,2,2)
+ contour(x1vec,x2vec,V,[0 0],'b')
+ axis([x1vec(1) x1vec(end) x2vec(1) x2vec(end)])
+ xlabel('x_1')
+ ylabel('x_2')
+ title('Optimal Reset Bound')
+ grid on
+ axis equal
+end