% MATLAB implementation of Example 7.15 from 
% Ralph C. Smith, Uncertainty Quantification, SIAM, 2014
% z''(t) + C*z'(t) + K*z(t) = 0; z(0)=2; z'(0) = -C.

 t = linspace(0,5,500)';     % sample times

 sig2 = 0.01;
% recover C only using nonlinear least squares
 C = 1.5;                     % true value of C
 K = 20.5;                    % true value of K (assume known for now)

 Y0 = get_shm(t,0);           % get samples; no noise
 Y  = get_shm(t,sig2,C,K);    % get samples; variance sig2


 Chat = lsqnonlin(@(C) (Y-2*exp(C*t/2).*cos(sqrt(K-C^2/4)*t)),[1]);

% build X, the matrix of partial derivatives 
 n = length(t);
 p = 1;
 X = exp(-C*t/2).*((C*t/sqrt(4*K-C^2)).*sin(sqrt(K-C^2/4)*t) ...
            - t.*cos(sqrt(K-C^2/4)*t)); 

% construct R, the residual vector
 Rhat = Y-2*exp(C*t/2).*cos(sqrt(K-C^2/4)*t);

% estimate variance of the noise
 sig2hat = Rhat'*Rhat/(n-p);

% estimate variance of C
 Vq = sig2hat*inv(X'*X);

 figure(1), clf
 
 plot(t, Y0, 'b-','linewidth',2)
 hold on
 plot(t, Y, 'r.','markersize',16)
 
 fprintf('Experiment 1:  C unknown\n');
 fprintf('---------------------------------------------------\n');
 fprintf('         estimated value of C:   %15.12f\n', Chat)
 fprintf('              true value of C:   %15.12f\n\n', C)
 fprintf('  estimated variance of noise:   %15.12f\n', sig2hat)
 fprintf('       true variance of noise:   %15.12f\n\n', sig2)
 fprintf('      estimated variance of C:   %15.12f\n', Vq)

% repeat the same procedure, but now recover K and C
% note: nonlinear solver is sensitive to initial guess of K!

 Q = lsqnonlin(@(Q) (Y-2*exp(Q(2)*t/2).*cos(sqrt(Q(1)-Q(2)^2/4)*t)),[1 20]);
 Khat = Q(1);
 Chat = Q(2);

% build X, the matrix of partial derivatives 
 XK = [exp(-C*t/2).*t.*sin(sqrt(K-C^2/4)*t)/sqrt(K-C^2/4)];
 XC = [ exp(-C*t/2).*((C*t/sqrt(4*K-C^2)).*sin(sqrt(K-C^2/4)*t) ...
            - t.*cos(sqrt(K-C^2/4)*t))]; 
 X = [XK XC];
 
% construct R, the residual vector
 Rhat = Y-2*exp(C*t/2).*cos(sqrt(K-C^2/4)*t);

% estimate variance of the noise
 sig2hat = Rhat'*Rhat/(n-p)

% estimate variance of K and C
 Vq = sig2hat*inv(X'*X)

 fprintf('Experiment 2:  C and K unknown\n');
 fprintf('---------------------------------------------------\n');
 fprintf('         estimated value of C:   %15.12f\n', Chat)
 fprintf('              true value of C:   %15.12f\n\n', C)
 fprintf('         estimated value of K:   %15.12f\n', Khat)
 fprintf('              true value of K:   %15.12f\n\n', K)
 fprintf('  estimated variance of noise:   %15.12f\n', sig2hat)
 fprintf('       true variance of noise:   %15.12f\n\n', sig2)
 fprintf('   est Covar matrix for C & K:   %10.7f  %10.7f\n', Vq(1,1),Vq(1,2))
 fprintf('                                 %10.7f  %10.7f\n', Vq(2,1),Vq(2,2))
