%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%
%                 SS2INI
%
%    Solve the equations for the 
%           Swinging Spring
%    for the 2 dimensional case. 
%        (Polar Coordinates)
%
%    Exercises to test initializations.
%
%    This version has solver inline.
% 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%      Basic spring parameters.
%        mass of bob : m
%        gravity     : g
%        stiffness   : k
%        length      : l
%
% Length of integration (seconds) : tmax
%
%  CODE BETWEEN THE LINES MARKED
%       vvvv  and  ^^^^^  
%  SHOULD NOT BE CHANGED !
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Clear memory and plot window
clear; clf; hold off;

% Define the basic (constant) parameters.
m=1;
g = pi^2;
k=100*pi^2;
l=1;

epssq=(m*g)/(k*l);
epsilon=sqrt(epssq)
l_0=(1-epssq)*l;
omega_R = sqrt(g/l);  omega_E = sqrt(k/m);
tau_R = 2*pi/omega_R, tau_E = 2*pi/omega_E

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%  LOOP FOR DIFFERENT INITIALIZATIONS
%%%    LOOP=1: Linear Initialization
%%%    LOOP=2: Nonlinear Initialization.

for LOOP = 1:2

% Set the initial conditions

  if ( LOOP == 1 )       % Linear Initialization:
    theta_init    = 1.0; 
    thetadot_init = 0.0;
    r_init        = l; 
    rdot_init     = 0.0;
  end

  if ( LOOP == 2 )       % Non-linear Initialization.
    theta_init    = 1.0; 
    thetadot_init = 0.0;
    r_init        = (l_0 + epssq*l*cos(theta_init) ) / ...
                    ( 1 - m*thetadot_init^2/k ) ;  %  check.
    rdot_init     = 0.0;
  end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%  Length of integration (seconds)
   tmax=6;

%====================================
%VVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVV
%  CODE BETWEEN THE LINES MARKED
%       vvvv  and  ^^^^^  
%  SHOULD NOT BE CHANGED !
%VVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVV

% Initial positions and momenta.
  x0(1) = theta_init; 
  x0(2) = thetadot_init; 
  x0(3) = r_init;
  x0(4) = rdot_init;

% Solve the set of four odes.


%  tspan = [0, tmax];
%  [tvec,xvec] = ode45('SS2roc',tspan,x0);

%  The numerical integration is handled by the method of ODE23P (inline).
    t0=0;
    tfinal=tmax;
    pow = 1/3;
    tol = 0.0001;  %%% reduced by factor of 10. 
    t = t0;
    hmax = (tfinal - t)/5;
    hmin = (tfinal - t)/200000;
    h = (tfinal - t)/100;
    x = x0(:);
    tau = tol * max(norm(x,'inf'),1);

    tvec = [];
    xvec = [];
    
% The main loop
while (h >= hmin)
  if t + h > tfinal, h = tfinal - t; end

% s1 = (feval('ss2roc', t, x))';
    xx = x;
    xdot(1) = xx(2)/(m*xx(3)*xx(3));
    xdot(2) = - m*g*xx(3)*sin(xx(1));
    xdot(3) = xx(4)/m;
    xdot(4) = xx(2)^2/(m*xx(3)^3)-k*(xx(3)-l_0)+m*g*cos(xx(1));
    s1 = xdot';
% s2 = (feval('ss2roc', t+h, x+h*s1))';
    xx = x+h*s1;
    xdot(1) = xx(2)/(m*xx(3)*xx(3));
    xdot(2) = - m*g*xx(3)*sin(xx(1));
    xdot(3) = xx(4)/m;
    xdot(4) = xx(2)^2/(m*xx(3)^3)-k*(xx(3)-l_0)+m*g*cos(xx(1));
    s2 = xdot';
% s3 = (feval('ss2roc', t+h/2, x+h*(s1+s2)/4))'; 
    xx = x+h*(s1+s2)/4;
    xdot(1) = xx(2)/(m*xx(3)*xx(3));
    xdot(2) = - m*g*xx(3)*sin(xx(1));
    xdot(3) = xx(4)/m; 
    xdot(4) = xx(2)^2/(m*xx(3)^3)-k*(xx(3)-l_0)+m*g*cos(xx(1));
    s3 = xdot';

% Estimate the error and the acceptable error
   delta = norm(h*(s1 - 2*s3 + s2)/3,'inf');
   tau = tol*max(norm(x,'inf'),1.0);
       	
% Update the solution only if the error is acceptable
   tvec = [tvec,t];
   xvec = [xvec,x];
   if delta <= tau
      t = t + h;
      x = x + h*(s1 + 4*s3 + s2)/6;
   end

% Update the step size
   if delta ~= 0.0
     h = min(hmax, 0.9*h*(tau/delta)^pow);
   end

end;  % Main loop ... 

% Calculate the solution a regular grid
% ( Note: this is required for the fft).
nt = 512; nv = size(tvec,2);
treg = linspace(t0,tmax,nt);
xreg=zeros(4,nt);
for n=2:nt-1
  for nn=2:nv
    if(tvec(nn)>=treg(n))
      w = (treg(n)-tvec(nn-1))/(tvec(nn)-tvec(nn-1));
      xreg(:,n)=w*xvec(:,nn)+(1-w)*xvec(:,nn-1);
      break
    end
  end
end
xreg(:,1)=x0'; 
xreg(:,nt)=xvec(:,nv);
 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Extract the solution.
  theta    = xreg(1,:);
  thetadot = xreg(2,:);
  r        = xreg(3,:);
  rdot     = xreg(4,:);
  
  rdash    = r-l;
  zz       = -r.*cos(theta);

%  TOTAL ENERGY
Energy=0.5*m*(rdot.^2+(r.*thetadot).^2)+0.5*k*(r-l_0).^2+m*g*zz;
Energy0=+0.5*k*(l-l_0).^2+m*g*(-l); 
Energy=Energy-Energy0;

% T_Spring = 0.5*m*(rdot.^2); V_Spring = 0.5*k*zzdash.^2;
% E_Spring = T_Spring+V_Spring;
% E_Swing = Energy-E_Spring;

%^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
%  CODE BETWEEN THE LINES MARKED
%       vvvv  and  ^^^^^  
%  SHOULD NOT BE CHANGED !
%^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
%====================================

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Plot some results.
  subplot(2,2,1) % Upper left panel
  title('theta')
  hold on
  if(LOOP==1) plot(treg,theta,':r'); end
  if(LOOP==2) plot(treg,theta,'b'); end 
  
  subplot(2,2,3) % Lower left panel
  title('rdash')
  hold on
  if(LOOP==1) plot(treg,rdash,':r'); end 
  if(LOOP==2) plot(treg,rdash,'b'); end 

  subplot(2,2,2) % Upper right panel
  title('theta spectrum')
  hold on
  th=theta-mean(theta);
  tpower = abs(fft(th));
  range = 1:50; % range=1:nt;
  freq = range/(tmax);
  if(LOOP==1) plot(freq,tpower(range),':r'); end
  if(LOOP==2) plot(freq,tpower(range),'b'); end

  subplot(2,2,4) % Lower right panel
  title('r spectrum')
  hold on
  rr=rdash-mean(rdash);
  rpower = abs(fft(rr));
  range = 1:50; % range=1:nt;
  freq = range/(tmax);
  if(LOOP==1) plot(freq,rpower(range),':r'); end
  if(LOOP==2) plot(freq,rpower(range),'b'); end
%  pause

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

end % End of loop.

% Plot a legend
subplot(2,2,2); legend('LNMI','NNMI');
subplot(2,2,4); legend('LNMI','NNMI');
 
hold off

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%