% Two-way Wave Equation Solver
% 1D with fixed boundaries
%18.086 - Gilbert Strang
%written by: Fred Pearce
%February 20, 2006

xbnd   = [0 50];                                 % x-boundaries of model (m)
bndtyp = [0 1];                                  % 0 = fixed (Dirichlet) & 1 = free (Neumann) w/ bndtyp(1) at x(1)
ntstps = 600;                                    % # of time steps to solve

cp     = 200;                                    % compressional wavespeed (m/s)
r      = 1;                                      % courant #
ppw    = 25;                                     % # of points per wavelength (caculate dx assuming "characteristic" frequency)
fc     = 40;                                     % "characteristic" frequency (Hz)

srctyp = 1;                                      % 0 = delta function, 1 = ricker wavelet, 2 = half-cycle cosine, 3 = square wave
srcamp = 1;                                      % source amplitude
srcx   = [25];                                   % location of source in x-dimension


%--------------------------------------------------------------------------------------------------------------------------------
% define time vector (t), space vector (x), & boundary flag

dx = cp./fc./ppw                                 % spatial step (m)
x  = [xbnd(1):dx:xbnd(2)].';
nx = length(x);

dt = r./fc./ppw                                  % time step (s)
t  = [0:dt:(ntstps-1).*dt].';
nt = length(t);

if (bndtyp(1) == 0 & bndtyp(2) == 0)
    bndflg = 1;
elseif (bndtyp(1) == 0 & bndtyp(2) == 1)
    bndflg = 2;
elseif (bndtyp(1) == 1 & bndtyp(2) == 1)
    bndflg = 3;
end

% define source

[trsh,srcxind] = min(abs(srcx-x));

if srctyp == 0                                   % delta function
    src = srcamp.*[1;zeros(nt-1,1)];
elseif srctyp == 1                               % ricker wavelet = 2nd derivative of Gaussian
    src = get_ricker(t,srcamp,fc);
elseif srctyp == 2                               % half-cycle cosine w/ fc
    ts  = [-.25./fc:dt:.25./fc].';
    src = srcamp.*[cos(2.*pi.*fc.*ts);zeros(nt-length(ts),1)];
elseif srctyp == 3                               % square-wave w/ fc
    nts = round(.5./fc./dt);
    src = srcamp.*[ones(nts,1);-ones(nts,1);zeros(nt-2.*nts,1)];
end

n1_src = sum(abs(src));                          % 1st norm to measure of total source energy

% initialize variables

cpr    = cp.*r;
ux     = zeros(nx,1);
vx     = zeros(nx,1);
exx    = zeros(nx+1,1);
n1_vx  = zeros(nt,1);

% solve wave equation

fh1 = figure;

for nn = 1:nt
    
    vx  = cpr.*(exx(2:nx+1)-exx(1:nx)) + vx;     % update particle velocities (j,n) *Note - cp&exx defined at (j+1/2)
    
    vx(srcxind) = vx(srcxind) + src(nn);         % add source term - prescribed particle velocity
    
    ux  = vx.*dt + ux;                           % update particle diplacements (j,n+1/2)
    
    exx(2:nx) = (ux(2:nx)-ux(1:nx-1))./dx;       % update gradient in ux (strains) (j+1/2,n+1/2)
        
    if bndflg == 1                               % fixed-fixed boundary (strain field symmetric about boundary)
        exx([1 nx+1]) = [exx(2) exx(nx)];       
    elseif bndflg == 2                           % fixed-free boundary
        exx([1 nx+1]) = [exx(2) -exx(nx)];
    elseif bndflg == 3
        exx([1 nx+1]) = [-exx(2) -exx(nx)];      % free-free boundary (strain field anti-symmetric about boundary)
    end
    
    n1_vx(nn) = sum(abs(vx));                    % 1st norm to measure total system energy
    
    figure(fh1), plot(x,vx)
    if r <= 1
        set(gca,'YLim',[-1.01.*srcamp./r 1.01.*srcamp./r])
    end
    drawnow
    
end

figure(fh1);
xlabel('x (m)')
ylabel('Particle Velocity (m/s)')

fh2 = figure;
ah2 = axes('NextPlot','add');
plot(t,repmat(n1_src,nt,1),'r--','LineWidth',2)
plot(t,n1_vx,'k','LineWidth',2)
xlabel('time (s)')
ylabel('Total Energy (m/s)')
legend('Total Source Energy','Total System Energy',4)