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           = 10  ;           % 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          = 2.9e-4 ;        % radius               [m]
T0          = 40 ;            % tension              [N]
sig0T       = 0.1 ;           % sig0 Loss transv     [1/s]
sig1        = 0.0004 ;        % sig1 Loss transv     [m^2/s]
sig0L       = 0.2 ;           % sig0 Loss long       [1/s]


%-- output parameters
rpL         = 0.32 ;          % readout point        [frac of L]
rpR         = 0.76 ;

%--raised cosine conditions
t0          = 0.001 ;         % force start time      [s]
Amp         = 2.5 ;
muq         = 2 ;             % 1 = half-raised; 2 = raised
tw          = 0.0008 ;        % contact duration      [s]
xh          = 0.72 ;          % source location       [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      = L/pi*sqrt(pi)*sqrt((-T0/2/pi + sqrt(T0^2/4/pi^2 + 1/k^2*EI*rA))/(EI))


hbar        = L/modMax/1.05 ;
th          = 0.5 + 0.5 * (T0*k^2*hbar^2 + 4*EI*k^2) / (rA*hbar^4) ;
hsq         = ( T0*k^2 + 4*rA*sig1*k ) + sqrt( ( T0*k^2 + 4*rA*sig1*k )^2 + 16*(2*th-1)*rA*EI*k^2 ) ;
hsq         = hsq / 2 / rA / (2*th - 1);
h           = 1.05 * sqrt(hsq) ;            %-- grid spacing (from stability condition)
M           = floor(L / h)  ;               %-- number of grid points
h           = L / M;
rpL         = round(rpL*L/h) ;
rpR         = round(rpR*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);
D2          = 1/h^2*spdiags([Dv, -2*Dv, Dv],-1:1,M-1,M-1);
D4          = D2^2;
Dxp         = 1/h*spdiags([-Dv, Dv],0:1,M-1,M);
Dxm         = -Dxp.';
Mat00       = (rA/k^2)*h*(Id+(1-th)/2*h^2*D2) + rA*Id/k*sig0T*h ;


%-- number of modes
om          = 0;
Nmodes      = 0;
MatCos      = [];
MatSin      = [];
omv         = [] ;
while om < 5000 * 2 * pi 
    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) ;
dd          = sqrt(diag(-2*h/L*MatSin.'*D2*MatSin)) ;
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.' ;

%-- source spatial distribution
Jh          = sparse(M-1,1) ;
lh          = xh*L/h ;
lhInt       = floor(lh) ;
frach       = lh - lhInt ;
if lhInt < M - 1
    Jh(lhInt)   = (1-frach) ;
    Jh(lhInt+1) = frach ;
else
    Jh(M-1) = 1 ;
end

%--raised cosine build
x0          = floor(t0*fs) ;
x           = x0 : floor(tw*fs) + x0 ;
xW          = length(x) ;
fRW         = 0.5 * Amp * (1 - cos( muq * pi * ( x - x0 ) / xW ) );

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


%--------------------------------------------------------
% Initialise
ovec        = ones(M,1) ;
u0          = zeros(M-1,1) ;
s0          = zeros(Nmodes,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)) ;

Et          = zeros(Ts,1);

outu        = zeros(Ts,2);
outuL       = zeros(Ts,1);
outuR       = zeros(Ts,1);
outz        = zeros(Ts,1);
outuWE      = zeros(Ts,1);
ff          = zeros(Ts,1) ;
ff(x)       = fRW ;

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


%--------------------------------------------------------
% Main Loop
tic
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-th)/2*h^2*D2)*(2*u0 - um) + h*D2*(T0*u0+2*rA*sig1/k*(u0-um)) - EI*h*D4*u0 + rA/k*sig0T*h*um + Mat12*sm + h*Dxp*(chim.*g) + Jh*ff(n);
    b2          = M0Mod*(2*s0 - sm) - vecMod.*s0 + rA*sig0L/k*L/2*sm + 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
    outuL(n)    = u0(rpL) ;
    outuR(n)    = u0(rpR) ;
    outu(n,:)   = [outuL(n),outuR(n)] ; %-- this is the output (displacement)
    tt =  MatSin*sp ; outz(n) = tt(rpL)  ;
    
    %-- 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([-0.01,0.01]); xlim([0,M*h]); grid on;
            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]);
            grid on; xlabel('x [m]'); ylabel('z(x) [m]'); drawnow;
        end
    end
end



time_elapsed    = toc
real_time_frac  = toc/T

%-- play sound
v = diff(outu) ;
soundsc(v,fs) ;


%-- spectrogram
lwindow = floor(0.018*fs) ;
overlap = floor(0.5*lwindow) ;
ww      = kaiser(lwindow) ;
[s,f,t] = spectrogram(diff(outuL),lwindow,overlap,[],fs,'yaxis');
spc = 20*log10(abs(s)) ;
[T,F] = meshgrid(t,f) ;
surf(T,F,spc,'EdgeColor','none','FaceColor','interp');
grid on;
view(2); axis tight ; colormap gray;
%xlim([0,2]);
ylim([0,24000])

%-- time domain plots
figure ;
subplot(2,1,1)
plot(tv,outu(:,1)/1e-3,'k') ;
xlabel('t (s)'); ylabel('u (mm)') ;
subplot(2,1,2)
plot(tv,outz/1e-3,'k') ;
xlabel('t (s)'); ylabel('v (mm)') ;
