% fits 2D data with differential equation as prior information

clear all;

% want to know m on Nx by Ny grid
Nx = 21;
dx = 1.0;
xmin = 0.0;
xmax = xmin+dx*(Nx-1);
x = xmin+dx*[0:Nx-1]';

Ny = 21;
dy = 1.0;
ymin = 0.0;
ymax = ymin+dy*(Nx-1);
y = ymin+dy*[0:Ny-1]';

% but data d measured at only 10% of the Nx by Ny grip points
Nd=round(Nx*Ny/10);
rowindex=unidrnd(Nx,Nd,1);
colindex=unidrnd(Ny,Nd,1);
xd=xmin+dx*rowindex;
yd=ymin+dy*colindex;

% true data is sinusoidal
fx=0.02;
fy=0.04;
zd=zeros(Nd,1);
for p = [1:Nd]
    zd(p) = sin(2*pi*fx*xd(p))*sin(2*pi*fy*yd(p))';
end

% array of measured data for plotting purposes
figure(1);
clf;
z = zeros(Nx,Ny);
for p = [1:Nd]
    z(rowindex(p),colindex(p)) = zd(p);
end

% plot measured data
imagesc( [xmin, xmax], [ymin, ymax], z);
hold on;
colorbar;
xlabel('y');
ylabel('x');
title( 'measured field z' );

M=Nx*Ny;
m=zeros(M,1);
% order unknowns: index = (ix-1)*Ny + iy

% top part of matrix is m = d
N=Nd;
G=zeros(N,M);
d=zeros(N,1);
for p = [1:Nd]
    q = (rowindex(p)-1)*Ny+colindex(p);
    G(p,q) = 1.0;
    d(p)=zd(p);
end

% bottom part of matrix is differential equation
% and boundary conditions

epsilon = 1e-6;
a=epsilon*(1.0);
b=epsilon*(-4.0);
c=epsilon*(-1.0);

Nc = (Nx-2)*(Ny-2) + 2*Nx + 2*Ny;
D=zeros(Nc,M);

ic=0;

% top edge
for iy = [1:Ny]
    ic=ic+1;
    ix=1;
    q1 = (ix-1)*Ny+iy;
    q2 = (ix+1-1)*Ny+iy;
    D(ic,q1)=a;
    D(ic,q2)=c;
end

% bottom edge
for iy = [1:Ny]
    ic=ic+1;
    ix=Nx;
    q1 = (ix-1-1)*Ny+iy;
    q2 = (ix-1)*Ny+iy;
    D(ic,q1)=a;
    D(ic,q2)=c;
end

% left edge
for ix = [1:Nx]
    ic=ic+1;
    iy=1;
    q1 = (ix-1)*Ny+iy;
    q2 = (ix-1)*Ny+iy+1;
    D(ic,q1)=a;
    D(ic,q2)=c;
end

% right edge
for ix = [1:Nx]
    ic=ic+1;
    iy=Ny;
    q1 = (ix-1)*Ny+iy-1;
    q2 = (ix-1)*Ny+iy;
    D(ic,q1)=a;
    D(ic,q2)=c;
end

% interior
for ix = [2:Nx-1]
for iy = [2:Ny-1]
    ic=ic+1;
    q0 = (ix-1)*Ny+iy;
    ql = (ix-1)*Ny+iy-1;
    qr = (ix-1)*Ny+iy+1;
    qt = (ix-1-1)*Ny+iy;
    qb = (ix-1+1)*Ny+iy;
    D(ic,q0)=b;
    D(ic,ql)=a;
    D(ic,qr)=a;
    D(ic,qb)=a;
    D(ic,qt)=a;
end
end

% standard least-squares solution
m= inv(G'*G + D'*D)*G'*d;

% rearrange solution m back into image
z2 = zeros(Nx,Ny);
for ix = [1:Nx]
for iy = [1:Ny]
    q0 = (ix-1)*Ny+iy;
    z2(ix,iy) = m(q0);
end
end

% plot solution
figure(2);
clf;
imagesc( [xmin, xmax], [ymin, ymax], z2);
hold on;
colorbar;
xlabel('y');
ylabel('x');
title( 'reconstructed field z' );

% error at observation points, for plotting purposes
z3 = zeros(Nx,Ny);
for p = [1:Nd]
    z3(rowindex(p),colindex(p)) = zd(p);
end

% plot error
figure(3);
clf;
imagesc( [xmin, xmax], [ymin, ymax], z3);
hold on;
colorbar;
xlabel('y');
ylabel('x');
title( 'error at observation points' );