%Cross-well example from Chapter 10 of Aster, Borchers, and Thurber, 2004
clear

% convergence enhancement factor for raypath plotting
XFAC=1.5;

%number of iterations in ray bending
NIT=50;

%Travel-time convergence parameter
CONV=2e-6;

%scale of problem (m)
SCALE=1600;

%scale of problem (receivers, sources and velocity nodes)
PSCALE=8;

%  source and receiver coordinate matrices
scz=linspace(0,SCALE,PSCALE); scx=zeros(size(scz));
sc=[scx;scz]';
rcz=linspace(0,SCALE,PSCALE); rcx=SCALE*ones(size(scz));
rc=[rcx;rcz]';

%  node positions xn, zn
xn=linspace(-150,SCALE+150,PSCALE)';
zn=xn;

%construct the true model
vbackground=2900;
for i=1:PSCALE
  for j=1:PSCALE
    vtrue(i,j)=vbackground*...
(1+0.10*exp(-.00004*((xn(i)-xn(3*PSCALE/4))^2+(zn(j)-zn(3*PSCALE/4))^2)))*...
(1-0.15*exp(-.00004*((xn(i)-xn(PSCALE/2))^2+(zn(j)-zn(PSCALE/2))^2)));
  end
end

%linerly interpolate the velocity field for plotting
for i=1:PSCALE
  xnm(i,:)=xn';
  znm(:,i)=zn;
end

xni=[-150:5:SCALE+150]'; zni=[-150:5:SCALE+150]';
for i=1:length(xni)
  xnim(i,:)=xni';
  znim(:,i)=zni;
end

vi=interp2(xnm,znm,vtrue,xnim,znim,'spline');

%plot velocity structure and ray paths (Figure 10.1).
figure(1)
bookfonts;
imagesc(xni,zni,vi)
colormap(gray)
caxis([2600 3200])
H=colorbar;
set(H,'FontSize',18);

%now plot the travel times for the model

hold on
%run the forward modeling and plotting function
disp('True velocity model:')
pause
tstor=plotraypaths(SCALE,PSCALE,NIT,CONV,XFAC,xn,zn,vtrue,sc,rc);
hold off
% another view
disp('A plot of velocity:')
figure(10)
[X,Z] = meshgrid(xn,zn);
meshc(X,Z,vtrue)

%save noise-free data
%save tstor.mat tstor vtrue
%inverse modeling of the cross-well example, with refracted rays

%this section reproduces the rest of the figures (and uses the true model from above to make comparisons).
%clear
%clear variables that are reused in this section
clear xnim
clear znim
clear xni

% convergence enhancement factor for subsequent subroutines
XFAC=1.2;

%Number of GN linearization steps
MAXITER=5;

%Travel-time convergence parameter
CONV=1e-4;

%noise level (s)
NOISE = 0.001;

%scale of problem (m)
SCALE=1600;

%scale of problem (receivers, sources and velocity nodes)
PSCALE=8;

%  source and receiver coordinate matrices
scz=linspace(0,SCALE,PSCALE); scx=zeros(size(scz));
sc=[scx;scz]';
rcz=linspace(0,SCALE,PSCALE); rcx=SCALE*ones(size(scz));
rc=[rcx;rcz]';

%  node positions xn, zn
xn=linspace(-150,SCALE+150,PSCALE)';
zn=xn;

%interpolated node positions
xni=[-150:5:SCALE+150]'; zni=[-150:5:SCALE+150]';

ivec=[1:length(xni)];
xnim(ivec,:)=xni;
znim(:,ivec)=zni';

%construct the 2-d 2nd order roughening matrix (with model wrap-around)
%
% Construct a two-dimensional, second-order roughening matrix
% for a PSCALE by PSCALE square nodal model (with 
% model wrap-around)
%
for i=1:PSCALE
  for j=1:PSCALE
    L((j-1)*PSCALE+i,(j-1)*PSCALE+i)=-4;
%
% Up
%
    if (i>1)
      L((j-1)*PSCALE+i,(j-1)*PSCALE+i-1)=1;
    else
      L((j-1)*PSCALE+i,(j-1)*PSCALE+PSCALE)=1;
    end
%
% Down
%
    if (i<PSCALE)
      L((j-1)*PSCALE+i,(j-1)*PSCALE+i+1)=1;
    else
      L((j-1)*PSCALE+i,(j-1)*PSCALE+1)=1;
    end
%
% Left
%
    if (j>1)
      L((j-1)*PSCALE+i,(j-2)*PSCALE+i)=1;
    else
      L((j-1)*PSCALE+i,(PSCALE-1)*PSCALE+i)=1;
    end
%
% Right
%
    if (j<PSCALE)
      L((j-1)*PSCALE+i,j*PSCALE+i)=1;
    else
      L((j-1)*PSCALE+i,i)=1;
    end
  end
end

%piter corresponds oto a particular choice of alpha
piter=0;
for piterexp=-2:0.25:1.75
piter=piter+1;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%now evaluate the travel time data for this true model
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%starting model
vbackground=2900;
v=vbackground*ones(PSCALE,PSCALE);
%convert velocity to slowness
slow=1./v;

%iterative slowness model vector
m=reshape(slow,PSCALE*PSCALE,1);

wv=zeros(4);

%noise free data and true velocity model
%load tstor.mat
%generate another realization of observed travel times with noise
%ttobs=tstor+NOISE*randn(size(tstor));
%save tstorenoise.mat ttobs

%...or load a particular realization used in the textbook
load tstorenoise.mat

mtrue=reshape(1./vtrue,PSCALE*PSCALE,1);

vstartvec=reshape(v,PSCALE*PSCALE,1);

%path segments
NSEG=PSCALE*2;

% start of iterative model loop
calcrays=0;
rpinit=0;

for iter=1:MAXITER

ttcal=zeros(PSCALE,PSCALE);

%get the Jacobian matrix, J, and forward modeled travel times for the working model, v
  [J,ttcal,rpstore]=getj(SCALE,PSCALE,NIT,CONV,XFAC,xn,zn,v,sc,rc,calcrays,rpinit);

%save raypaths to start the next iteration
  calcrays=1;
  rpinit=rpstore;

%set the regularization tradeoff parameter here

  if (iter == 1)
    rparam=norm(J)*10^piterexp;
  end

  alphasq(piter)=rparam;

%calculate the travel-time residual vector for this iteration
  r=ttcal-ttobs;
  rvec=reshape(r,PSCALE*PSCALE,1);
  rms(iter)=sqrt(rvec'*rvec);

%GN, explicit zeroth order regularization
  J1= J'*J+rparam*L'*L;
  rhs=-(J'*rvec+rparam*L'*L*m);
  dm=J1\rhs;

%model update
  m=m+dm;
  slow=reshape(m,PSCALE,PSCALE);
  mnslow=mean(mean(slow));
  v=1./slow;


%run the forward model to document variance reduction at this step
  [K,tttry,rpstore]=getj(SCALE,PSCALE,NIT,CONV,XFAC,xn,zn,v,sc,rc,calcrays,rpinit);
  rtry=ttobs-tttry;
  dtry=reshape(rtry,PSCALE*PSCALE,1);
  rmstry=sqrt(dtry'*dtry);
  disp('alpha, iteration, rmsnew, sqrt(chi^2)')
  [sqrt(alphasq(piter)) iter rmstry NOISE*PSCALE]
  misfit(piter,iter)=rmstry;
  mnorm(piter,iter)=norm(L*m);


%rms difference wrt true model
  mrms(piter,iter)= norm((mtrue-m));

%figure 2 shows the separate model for each iteration as the code runs, for different regularization parameters (alpha).
%the final model in each case is assembled into figure 10.2
  figure(2)
  bookfonts
%linerly interpolate the velocity field for plotting
  for i=1:PSCALE
    xnm(i,:)=xn';
    znm(:,i)=zn;
  end

  vi=interp2(xnm,znm,v,xnim,znim,'spline');

  caxis([2300 3200])
  imagesc(xni,zni,vi)
  colormap('gray')
  H=colorbar;
  set(H,'FontSize',18);
  title(['\alpha = ',num2str(sqrt(alphasq(piter)))]);
  drawnow;
  pause(0.1);

%  end of tomography inversion loop
end

vstore(piter,:,:)=vi;

% end of regularization parameter loop
end


%show results for the various regularization parameters

%this figure (not used in text) shows the residual stats for each iteration and alpha value
%figure(3)
%bookfonts
%plot(1:iter,misfit)
%xlabel('Iteration')
%ylabel('Residual Norm, ||G(m)-d||_2')

%this figure (not used in text) shows the seminorm stats for each iteration and alpha value
%figure(4)
%bookfonts
%semilogy(1:iter,mnorm)
%xlabel('Iteration')
%ylabel('Solution Seminorm, ||Lm||_2')

%L-curve (figure 10.3)
figure(5)
bookfonts;
loglog(misfit(1:piter,MAXITER),mnorm(1:piter,MAXITER),'ok-')
xlabel('Residual Norm, ||G(m)-d||_2')
ylabel('Solution Seminorm, ||Lm||_2')
axis([.001 .1 .00001 .001])
hold on
loglog([NOISE*PSCALE,NOISE*PSCALE],[.00001,.001],'k--')
for i=1:2:piter
text(misfit(i,MAXITER),mnorm(i,MAXITER),['   ',num2str(sqrt(alphasq(i)))]);
end
text(NOISE*PSCALE,.000015,'   \delta=0.008')
hold off


%Models compared to true model (Figure 10.5) 
figure(6)
bookfonts;
semilogx(sqrt(alphasq),mrms(1:piter,MAXITER),'ok-')
xlabel('\alpha')
ylabel('||m_{true}-m||_2')
hold on
semilogx(sqrt(alphasq(9)),mrms(9,MAXITER),'*')
for i=[1,9,16]
text(sqrt(alphasq(i)),mrms(i,MAXITER)+.000002,['  ',num2str(sqrt(alphasq(i)))]);
end
hold off


%Suite of models (Figure 10.2)
figure(7)
for i=1:piter
subplot(4,4,i)
vtmp1=vstore(i,:,:);
vtmp(:,:)=vtmp1;
imagesc(vtmp)
caxis([2600 3200])
axis square
set(gca,'xticklabel','')
set(gca,'yticklabel','')
title(['\alpha = ',num2str(sqrt(alphasq(i)))])
colormap('gray')
end

%"best" model (Figure 10.4)
disp('Best model: ')
pause
figure(8)
bookfonts
vtmp1=vstore(9,:,:);
vtmp(:,:)=vtmp1;
imagesc(xn,zn,vtmp)
caxis([2600 3200])
axis square
xlabel('m')
ylabel('m')
colormap('gray')
H=colorbar;
set(H,'FontSize',18);
