%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Math 3520: Differential Equations and Dynamical Systems
%   Nonlinear first-order systems: X'=F(X)
%   Plot trajectories for the Lorenz system.
%   Trajectories using Runge-Kuta 4th order solver
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%% Runge-Kutta 4th order solver in 3D:
function X = rk4b(X0,F,steps,h)
 % Input
 %  X0 = [x0,y0,z0]   Row vector, initial position
 %  F  = function handle, computes Xdot(x,y,z)
 %  steps = number of time steps to compute (e.g., 100)
 %  h = time increment per step (e.g., 1/steps )
 % Output
 %  X = (1+steps)x(3) matrix of (x,y,z) values along
 %                    the trajectory starting at X0
  X = zeros(1+steps,3);  % Solution to be computed
  X(1,:) = X0;  % Initial point
  for n=1:steps
    f1 = F(X(n,:));
    f2 = F(X(n,:)+f1*h/2);
    f3 = F(X(n,:)+f2*h/2);
    f4 = F(X(n,:)+f3*h);
    X(n+1,:) = X(n,:) +h*(f1+2*f2+2*f3+f4)/6;
  end
end

%%%% From Strogatz, 3rd edition, p.353
X0=[0,1,0]; % Lorenz's initial value
sigm=10; % Prandtl number, from Lorenz
rayl=28; % Rayleigh number, from Lorenz
b=8/3;   % unnamed aspect ratio, from Lorenz
Xdot = @(V) [sigm*(V(2)-V(1)), (rayl-V(3))*V(1)-V(2), V(1)*V(2)-b*V(3)];

%%%% Trajectories
h=0.01;   % Use same time step for same accuracy
clf;
X = rk4b(X0, Xdot, 5000,h);  % Complete trajectory to t=50.
init=1:100;   plot3(X(init,1),X(init,2),X(init,3));  % 0<t<1;
hold on; pause(2);
next=101:1000; plot3(X(next,1),X(next,2),X(next,3)); pause(2);
next=1001:2000; plot3(X(next,1),X(next,2),X(next,3)); pause(2);
next=2001:3000; plot3(X(next,1),X(next,2),X(next,3)); pause(2);
next=3001:5001; plot3(X(next,1),X(next,2),X(next,3)); pause(2);
xlabel("X"); ylabel("Y"); zlabel("Z");
title("Trajectory for the Lorenz Equation")

%%% Individual components of a trajectory on new graphs
figure; plot(X(:,1)); title("X Coordinate");
figure; plot(X(:,2)); title("Y Coordinate");
figure; plot(X(:,3)); title("Z Coordinate");

%%% Plot (X,Y) projection of the trajectory
figure; plot(X(:,1),X(:,2));
xlabel("X"); ylabel("Y");
title("Projection onto X,Y plane of Lorenz trajectory.");



%%%% Lorenz map: find the local maximum values in z=X(:,3)

%%%% Function to find local arg-maxima and maximum values:
function [Xs,Ys] = lmax(t,z)
% Input
%  t = Row vector of independent variable samples
%  z = row vector of dependent variable samples
% Output
%  X = row vector of local maxima
%  Y = row vector local maximum values
 dz = diff(z);  % differences z(k+1)-z(k), k=1,2,...
 n = 0;         % initial local max index
 for k=1:(length(dz)-1)
   if ( dz(k)>0 && dz(k+1)<0 )  % found local max at k+1
      n = n+1;      % increment the count of local maxes
      x1=t(k); x2=t(k+1); x3=t(k+2);
      y1=z(k); y2=z(k+1); y3=z(k+2);
      xnum=x1*x1*(y3-y2)+x2*x2*(y1-y3)+x3*x3*(y2-y1);
      xden=x1*(y3-y2)+x2*(y1-y3)+x3*(y2-y1);
      Xs(n) = 0.5*xnum/xden;  % append nth arg max
      d12=x1-x2; d31=x3-x1; d23=x2-x3;
      ynum1= 2*(y1*y1*d23^4+y2*y2*d31^4+y3*y3*d12^4);
      ynum2= (y1*d23^2+y2*d31^2+y3*d12^2)^2;
      yden = 4*d12*d23*d31*(y1*d23+y2*d31+y3*d12);
      Ys(n) = (ynum1-ynum2)/yden; % append nth max val
   end  % endif
 end  % endfor
 return
end % endfunction

% Compute the Lorenz map from a particular trajectory
%%%% From Strogatz, 3rd edition, p.353
X0=[0,1,0]; % Lorenz's initial value
sigm=10; % Prandtl number, from Lorenz
rayl=28; % Rayleigh number, from Lorenz
b=8/3;   % unnamed aspect ratio, from Lorenz
Xdot = @(V) [sigm*(V(2)-V(1)), (rayl-V(3))*V(1)-V(2), V(1)*V(2)-b*V(3)];

X = rk4b(X0, Xdot, 10000,h);  % Complete trajectory to t=100.
t = (0:10000)*h;     % time values
z = X(:,3);          % z-coordinate values
figure; plot(t,z);   % display the trajectory's z-coordinate
[Xs,Ys] = lmax(t,z); % find the local maximum values in z()
Zn = Ys(1:length(Ys)-1);
Znp1 = Ys(2:length(Ys));
figure;  plot(Zn,Znp1,"o");  % point plot of z_n vs. z_{n+1}
hold on;
zz=floor(min(Ys)):ceil(max(Ys)); plot(zz,zz);  % plot Y=X line
title("Lorenz map, r=28"); xlabel("Z_n"); ylabel("Z_{n+1}");

%%%% Try another Rayleigh number, smaller than rH
X0=[0,1,0]; % Lorenz's initial value
sigm=10; % Prandtl number, from Lorenz
rayl=21; % Rayleigh number, from Lorenz
b=8/3;   % unnamed aspect ratio, from Lorenz
Xdot = @(V) [sigm*(V(2)-V(1)), (rayl-V(3))*V(1)-V(2), V(1)*V(2)-b*V(3)];
X = rk4b(X0, Xdot, 10000,h);  % Complete trajectory to t=100.
t = (0:10000)*h;     % time values
z = X(:,3);          % z-coordinate values
figure; plot(t,z);   % display the trajectory's z-coordinate
[Xs,Ys] = lmax(t,z); % find the local maximum values in z()
Zn = Ys(1:length(Ys)-1);
Znp1 = Ys(2:length(Ys));
figure;  plot(Zn,Znp1,"o");  % point plot of z_n vs. z_{n+1}
hold on;
zz=floor(min(Ys)):ceil(max(Ys)); plot(zz,zz);  % plot Y=X line
title("Lorenz map, r=21"); xlabel("Z_n"); ylabel("Z_{n+1}");

%%%% Try another Rayleigh number, much bigger than rH
X0=[0,1,0]; % Lorenz's initial value
sigm=10; % Prandtl number, from Lorenz
rayl=99; % Rayleigh number, from Lorenz
b=8/3;   % unnamed aspect ratio, from Lorenz
Xdot = @(V) [sigm*(V(2)-V(1)), (rayl-V(3))*V(1)-V(2), V(1)*V(2)-b*V(3)];
X = rk4b(X0, Xdot, 10000,h);  % Complete trajectory to t=100.
t = (0:10000)*h;     % time values
z = X(:,3);          % z-coordinate values
figure; plot(t,z);   % display the trajectory's z-coordinate
[Xs,Ys] = lmax(t,z); % find the local maximum values in z()
Zn = Ys(1:length(Ys)-1);
Znp1 = Ys(2:length(Ys));
figure;  plot(Zn,Znp1,"o");  % point plot of z_n vs. z_{n+1}
hold on;
zz=floor(min(Ys)):ceil(max(Ys)); plot(zz,zz);  % plot Y=X line
title("Lorenz map, r=99"); xlabel("Z_n"); ylabel("Z_{n+1}");