clear all
close all
clc

%--------------------------------------------------------
%              Geometrically Exact 
%            Nonlinear Wave Equation:
%       Fully-Implicit Scheme (Newton-Raphson)
%               Raised Cosine Source
%                Dr M Ducceschi
%            University of Bologna
%                  22 Dec 2021
%--------------------------------------------------------

% Custom Parameters (user can modify these)
%--------------------------------------------------------

fs0         = 48e3 ;           % base sample rate     [Hz]
T           = 0.01  ;           % time of simulation   [s]
OF          = 4 ;             % oversampling factor

%-- string parameters
rho         = 8000 ;          % string density       [Kg/m^3]
L           = 1.0 ;           % string length        [m]
E           = 2e11 ;          % Young's modulus      [Pa]
rd          = 2.9e-4 ;        % radius               [m]
T0          = 40 ;            % tension              [N]


%-- output parameters
rp          = 0.72 ;          % readout point        [frac of L]


%--raised cosine conditions
Amp         = 2e-3 ;             % maximum force         [N]
sig         = 0.1 ;

%-- flags
flagLin     = 0 ;             % 1 = linear string dynamics
flagPlot    = 0 ;             % 1 = liveplotting
plotRefresh = 5*OF ;          % plot refresh rate    [samples]

NITER       = 5 ;
% End Custom Parameters
%--------------------------------------------------------


% Derived Parameters (do not modify below)
%--------------------------------------------------------
A           = pi * rd * rd ;                    %-- area
I           = pi / 4 * rd * rd * rd * rd ;      %-- moment of inertia
EI          = 0 * E * I ;
EA          = E * A ;
rA          = rho * A;

fs          = OF * fs0 ;                    %-- sample rate
k           = 1 / fs;                       %-- time step
Ts          = floor(fs * T);                %-- number of samples
tv          = (0:Ts-1)*k ;                  %-- time array


hsq         = ( T0*k^2  ) + sqrt( ( T0*k^2 )^2 + 16*rA*EI*k^2 ) ;
hsq         = hsq / 2 / rA ;
h           = 1.05 * sqrt(hsq) ;            %-- grid spacing (from stability condition)
h           = sqrt(EA / rA) * k ;

M           = floor(L / h)  ;               %-- number of grid points
h           = L / M;
rp          = round(rp*M) ;




%--matrices
Dv          = ones(M,1);
Dxp         = 1/h*spdiags([-Dv, Dv],0:1,M-1,M);
Dxm         = -Dxp.';
D2          = Dxp*Dxm;
D2d         = Dxm*Dxp;
D4d         = Dxm*D2*Dxp;
Id          = speye(M-1,M-1);
IdNT        = speye(M,M);
Mat00       = (2*rA*Id/k) ;

%--initial conditions
x           = 1 : M - 1;
fRW         = zeros(1,M - 1);
xR          = floor((0.5-sig) * (M - 1)) : ceil((0.5+sig) * (M - 1));
x0          = 0.5 * (M - 1);
xW          = 2*(M - 1)*sig;
fRW(xR)     =  0.5 * Amp * (1 + cos( 2 * pi * ( xR - x0 ) / xW ) );

mf          = max(abs(fRW));  %% this is used in plot ranges
mfq         = max(abs(Dxm*fRW.'));

if Amp == 0
    mf = fAmp;
    mfq = mf;
end


%--------------------------------------------------------


%--------------------------------------------------------
% Initialise
z0          = ones(M - 1, 1).*rand(M - 1,1)*1e-15;
u0          = fRW.';

um          = u0;
zm          = z0;

EtNT        = zeros(Ts,1);
outuNT      = zeros(Ts,1);
outzNT      = zeros(Ts,1);

u0NT        = u0;
umNT        = um;
z0NT        = z0;
zmNT        = zm;

qmNT = Dxm*umNT;
pmNT = Dxm*zmNT;
q0NT = Dxm*u0NT;
p0NT = Dxm*z0NT;

time_elapsed = zeros(NITER,1);
%--------------------------------------------------------


%--------------------------------------------------------
% Main Loop
for ncount = 1 : NITER
    ncount
    tic
    for n = 1 : Ts
        
        %%--- Newton-Raphson
        rqNT        = q0NT-qmNT ;
        rpNT        = p0NT-pmNT ;
        avecq0p0    = sqrt((1+p0NT).^2 + q0NT.^2);
        avecqmp0    = sqrt((1+p0NT).^2 + qmNT.^2);
        avecq0pm    = sqrt((1+pmNT).^2 + q0NT.^2);
        phiqmp0     = 0.5 * (EA - T0) * (avecqmp0 - 1).^2;
        phiq0pm     = 0.5 * (EA - T0) * (avecq0pm - 1).^2;
        
        nNt = 0 ;
        tol = 10 ;
        while tol > 1e-13 && nNt < 100
            
            nNt = nNt + 1 ;
            
            avecqpp0    = sqrt((1+p0NT).^2 + (rqNT+qmNT).^2);
            avecq0pp    = sqrt((1+rpNT+pmNT).^2 + q0NT.^2);
            
            phiqpp0     = 0.5 * (EA - T0) * (avecqpp0 - 1).^2;
            phiq0pp     = 0.5 * (EA - T0) * (avecq0pp - 1).^2;
            phiPqpp0    = (EA - T0) * (avecqpp0 - 1) .* (rqNT + qmNT) ./ avecqpp0;
            phiPq0pp    = (EA - T0) * (avecq0pp - 1) .* (rpNT + pmNT + 1) ./ avecq0pp;
            
            VecqTemp    =  (rqNT .* phiPqpp0 - phiqpp0 + phiqmp0) ./ (rqNT).^2 ;
            VecpTemp    =  (rpNT .* phiPq0pp - phiq0pp + phiq0pm) ./ (rpNT).^2 ;
            
            
            VecqTemp(isnan(VecqTemp)) = 0;
            VecpTemp(isnan(VecpTemp)) = EA ;
            VecqTemp(isinf(VecqTemp)) = 0;
            VecpTemp(isinf(VecpTemp)) = EA ;
            
            
            MatqTemp    = D2d * sparse( diag ( VecqTemp ) ) ;
            MatpTemp    = D2d * sparse( diag ( VecpTemp ) ) ;
            
            JacqNT      = rA/k^2*IdNT - (MatqTemp);
            JacpNT      = rA/k^2*IdNT - (MatpTemp);
            
            
            VecqTemp2    =  (phiqpp0 - phiqmp0) ./ (rqNT) ;
            VecpTemp2    =  (phiq0pp - phiq0pm) ./ (rpNT) ;
            
            VecpTemp2(isnan(VecpTemp2)) = (EA - T0) * (avecq0p0(isnan(VecpTemp2)) - 1) .* (p0NT(isnan(VecpTemp2)) + 1) ./ avecq0p0(isnan(VecpTemp2));
            VecpTemp2(isinf(VecpTemp2)) = (EA - T0) * (avecq0p0(isinf(VecpTemp2)) - 1) .* (p0NT(isinf(VecpTemp2)) + 1) ./ avecq0p0(isinf(VecpTemp2));
            
            VecqTemp2(isnan(VecqTemp2)) = 0;
            VecqTemp2(isinf(VecqTemp2)) = 0;
            
            fqNT        = rA/k^2*(rqNT-2*q0NT+2*qmNT) - D2d * VecqTemp2 - T0 * D2d * q0NT + EI * D4d * q0NT ;
            fpNT        = rA/k^2*(rpNT-2*p0NT+2*pmNT) - D2d * VecpTemp2 - T0 * D2d * p0NT;
            
            rqNTNew        = rqNT - JacqNT\fqNT;
            rpNTNew        = rpNT - JacpNT\fpNT;
            
            tol = abs((rqNTNew.')*(rqNTNew) - (rqNT.')*(rqNT)) + abs((rpNTNew.')*(rpNTNew) - (rpNT.')*(rpNT));
            
            rqNT = rqNTNew ;
            rpNT = rpNTNew ;
            
            
        end
        
        %nNt
        
        qpNT = (rqNT+qmNT);
        ppNT = (rpNT+pmNT);
        upNT = Dxm\qpNT;
        zpNT = Dxm\ppNT;
        
        
        %%--- record output
        
        outuNT(n)   = upNT(rp) ;
        outzNT(n)   = zpNT(rp) ;
        
        
        
        %%--- update
        
        
        qmNT        = q0NT ;
        q0NT        = qpNT ;
        pmNT        = p0NT ;
        p0NT        = ppNT ;
        umNT        = u0NT ;
        u0NT        = upNT ;
        zmNT        = z0NT ;
        z0NT        = zpNT ;
        
        
    end
    
    time_elapsed(ncount)  = toc ;
end


mean(time_elapsed)

