clear all
close all
clc

%--------------------------------------------------------
%              Geometrically Exact
%            Nonlinear Wave Equation:
%        Linearly-Implicit, Conservative Scheme
%               Mixed FD / Modal
%               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          = 1 ;             % oversampling factor

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


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


%-- flags
flagLin     = 0 ;             % 1 = linear string dynamics
flagPlot    = 0 ;             % 1 = liveplotting
plotRefresh = 5*OF ;          % plot refresh rate    [samples]
% 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          = 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

k0          = 1/fs0 ;
modMax      = sqrt((-T0/L^2/rA+sqrt(T0^2/L^2/rA^2+4/k^2*EI/L^4*pi^2/rA))/(2*EI/L^4/rA*pi^2)) ;
hbar        = L/modMax/1.05 ;
thu         = 0.5 + 0.5 * (T0*k^2*hbar^2 + 4*EI*k^2) / (rA*hbar^4) ;
hsq         = ( T0*k^2 ) + sqrt( ( T0*k^2  )^2 + 16*(2*thu-1)*rA*EI*k^2 ) ;
hsq         = hsq / 2 / rA / (2*thu - 1);
h           = 1.05 * sqrt(hsq) ;            %-- grid spacing (from stability condition)
M           = floor(L / h)  ;               %-- number of grid points
h           = L / M;
rp          = round(rp*L/h) ;



%-- some constants
rAhhLksq    = 0.25*rA/k/k*L ;
rAksq       = rA/k/k ;

%-- matrices
Id          = speye(M-1,M-1);
Dv          = ones(M,1);
Dxx         = 1/h^2*spdiags([Dv, -2*Dv, Dv],-1:1,M-1,M-1);
Dxxxx       = Dxx^2;
Dxp         = 1/h*spdiags([-Dv, Dv],0:1,M-1,M);
Dxm         = -Dxp.';
Mat00       = (rA/k^2)*h*(Id+(1-thu)/2*h^2*Dxx)  ;


%-- number of modes
om          = 0;
Nmodes      = 0;
MatCos      = [];
MatSin      = [];
omv         = [] ;
while om < 2 * fs
    Nmodes      = Nmodes + 1 ;
    om          = sqrt(EA / rA) * Nmodes * pi / L ;
    omv         = [omv;om] ;
    vec         = sin(Nmodes*pi*h*((1:M-1).')/L);
    MatSin      = [MatSin,vec];
end

MatCos      = Dxm*MatSin ;
omv         = omv(1:end-1) ;
IdMod       = speye(Nmodes,Nmodes) ;
%Lambda      = diag((1:Nmodes)*pi/L) ; %-- old lambda
dd          = sqrt(diag(-2*h/L*MatSin.'*Dxx*MatSin)) ;
%  dd = ones(length(dd),1) ; %-- uncomment to test energy conservation in a non-consistent case
Lambda      = diag(dd) ;            %-- reflects reviewer's comment
M0Mod       = L/2*(rA*IdMod/k^2 - (EA-T0)/7*Lambda^2) ;
MatSig0L    = IdMod ;
Mat00Mod    = M0Mod  ;
vecMod      = T0*L/2*(dd.^2) ;
vecMatMod   = sparse(diag(vecMod)) ;
MatCosP     = MatCos.' ;

% End Derived Parameters
%--------------------------------------------------------


%--------------------------------------------------------
% Initialise
ovec            = ones(M,1) ;
M3              = floor(M/3) ;
u0              = zeros(M-1,1) ;
sigRaised = 0.1 ;
for m = 1 : M - 1 
    if m*h >= 0.5*L-sigRaised && m*h <= 0.5*L + sigRaised
        u0(m) = 0.5*Amp * (1 + cos(2*pi*(m*h-0.5*L)/2/sigRaised/L)) ;
    end
end
    
s0          = zeros(Nmodes,1) ;
v0          = zeros(M-1,1) ;
um          = u0;
sm          = s0;

q0          = Dxm*u0;
p0          = MatCos*s0;
qm          = Dxm*um;
pm          = MatCos*sm;

avec        = sqrt( (ovec + 0.5*(p0+pm)).^2 + (0.5*q0+0.5*qm).^2 ) ;
phi         = 0.5 * (EA-T0) * ((avec-ovec).^2 ) ;
chim        = sqrt(2*abs(phi)) ;



outu        = zeros(Ts,1);
outz        = zeros(Ts,1);
Ek          = zeros(Ts,1) ;
Epl         = zeros(Ts,1) ;
Epnl        = zeros(Ts,1) ;

% End Initialise
%--------------------------------------------------------


%--------------------------------------------------------
% Main Loop

for n = 1 : Ts
    
    %-- build g's
    
    %--(nonlinear WE)
    q0          = Dxm*(u0) ;
    p0          = MatCos*s0 ;
    avec        = sqrt( (ovec+p0).^2 + q0.^2 ) ;
    
    g           = sqrt(EA - T0) * (  q0 ./ avec )  ;
    f           = sqrt(EA - T0) * ( (ovec + p0) ./ avec )   ;
    gsq         = g.*g ;
    fsq         = f.*f ;
    
    if flagLin == 1
        g = 0*g ;
        f = 0*f ;
    end
    
    
    %-- Build Loop Matrices
    Lg          = sparse(diag(g)) ;
    Lf          = sparse(diag(f)) ;
    
    B1          = Dxp*Lg ;
    B2          = -MatCosP*Lf ;
    
    Mat1s       = 0.25*h*B1*(B1.')    ;
    Mat2s       = 0.25*h*B2*(B2.') ;
    Mat11       = Mat00 + Mat1s ;
    Mat12       = 0.25*h*B1*(B2.') ;
    Mat21       = Mat12.' ;
    Mat22       = Mat00Mod + Mat2s ;
    
    %-- Block LU decomposition
    z10         = (rA/k^2)*h*(Id+(1-thu)/2*h^2*Dxx)*(2*u0 - um) + h*Dxx*(T0*u0) - EI*h*Dxxxx*u0 + Mat12*sm + h*Dxp*(chim.*g) ;
    b2          = M0Mod*(2*s0 - sm) - vecMod.*s0 + Mat21*um + Mat2s*sm  - h*(MatCosP)*(chim.*f ) ;
    z1          = z10 + Mat1s*um  ;
    y1          = Mat11\z1 ;
    z2          = -Mat21*y1 + b2 ;
    dt          = (Mat11\Mat12) ;
    
    %-- solutions
    sp          = (Mat22-Mat21*dt)\z2 ;
    up          = y1 - dt*sp ;
    
    chip        = 0.5*g.*( Dxm*(up-um) ) + 0.5*f.*( MatCos*(sp-sm) ) + chim ;
    
    
    %-- record output
    outu(n)    = u0(rp) ;
    vp =  MatSin*sp ; outz(n) = vp(rp)  ;
    
    %-- Energy
    Ek(n)   = 0.5 * h * (rAksq*((up-u0).'*(up-u0))  - rAksq*0.5*(1-thu)*h^2*(Dxm*(up-u0)).'*Dxm*(up-u0) ) + ...
        rAhhLksq*((sp-s0).'*(sp-s0)) - L/4*(EA-T0)/7*sum((Lambda*(sp-s0)).^2) ;
    Epl(n)  = 0.5 * h * ( T0 * (Dxm*u0).'* Dxm*up + EI * (Dxx*u0).'* Dxx*up ) + ...
        + 0.5*sum(vecMod.*sp.*s0) ;
    Epnl(n) = h/2 * (chip.') * chip ;
    

    %-- update
    um          = u0 ;
    u0          = up ;
    sm          = s0 ;
    s0          = sp ;
    chim        = chip ;
    
    
    %-- live plotting
    if flagPlot == 1
        if mod(n,plotRefresh) == 0 || n == 1
            
            subplot(2,1,1)
            cla ;
            hold on ;
            plot((0:M)*h,[0;um;0],'linewidth',1.5,'color','k')
            hold off;
            ylim([-Amp,Amp]); xlim([0,M*h]); 
            xlabel('x [m]'); ylabel('u(x) [m]');  drawnow;
            
            subplot(2,1,2)
            plot((0:M)*h,[0;MatSin*sm;0],'linewidth',1.5,'color','k');
            ylim([-0.001,0.001]); xlim([0,M*h]);
            xlabel('x [m]'); ylabel('z(x) [m]'); drawnow;
        end
    end
end

Et = Ek+Epl+Epnl ;

subplot(2,2,1)
plot(tv,outu(:,1)/1e-3,'k') ;
xlabel('t (s)'); ylabel('u (mm)') ;
title('transv disp')
subplot(2,2,2)
plot(tv,outz/1e-3,'k') ;
xlabel('t (s)'); ylabel('v (mm)') ;
title('long disp')
subplot(2,2,3)
plot(tv,Et,'k') ; hold on ;
plot(tv,Ek,'k','marker','o','linestyle','none','markersize',2) ;
plot(tv,Epl,'k','marker','+','linestyle','none','markersize',2) ;
plot(tv,Epnl,'k','marker','x','linestyle','none','markersize',2) ;
legend('H','Ek','Epl','Epnl')
xlabel('t (s)'); ylabel('H (J)') ;
title('Energy')
subplot(2,2,4)
plot(tv,1-Et/Et(1),'k','marker','+','linestyle','none','markersize',2,'linestyle','none') ;
xlabel('t (s)'); ylabel('\Delta H') ;
title('Energy error') 
