function [ ts, ys ] = butterfly(init, tMax)
% butterfly: Simulate Lorenz model of atmospheric circulation.
%
%   [ ts, ys ] = butterfly(init, tMax)
%
% Demonstrates the simple Lorenz model of atmospheric circulation from Heath
% computer problem 9.6.
%
% System dynamics:
%
%     \dot y_1      = \sigma (y_2 - y_1)
%     \dot y_2      = r y_1 - y_2 - y_1 y_3
%     \dot y_3      = y_1 y_2 - b y_3
%
%   where nominal parameters are \sigma = 10, b = 8/3, r = 28.
%
% Input Parameters:
%
%   init: Initial state condition.  May be a 3 element column vector (only
%   one initial condition) or a 3 x n array (n initial conditions).  Default
%   is [ 0; 1; 0 ].
%
%   tMax: Final time for integrations.  Default = 50.
%
% Output Parameters:
%
%   ts: Time steps for each run.  A column vector (such as would be returned
%   by ode45) if there was only one initial condition, otherwise a cell
%   vector of column vectors.
%
%   ys: State vector for each run.  An array with three columns and one row
%   for each timestep if there was only one initial condition, otherwise a
%   cell vector of such arrays.
%
% Ian Mitchell for CPSC 303, 1/7/07.

%---------------------------------------------------------------------------
% Set default parameters
if(nargin < 1)
  init = [ 0; 1; 0 ];
end

num_init = size(init, 2);

if(nargin < 2)
  tMax = 50;
end

%---------------------------------------------------------------------------
% Nominal parameters from Ascher & Petzold.
params = [ 10; 8/3; 28 ];

% What integrator to use?
intfunc = @ode45;

% What integration options to use?
options = odeset('RelTol', 1e-6);
options = odeset;

%---------------------------------------------------------------------------
% Indices for the various components of the ODE.
yIndex = 1:3;

%---------------------------------------------------------------------------
% Cell vector return parameters.
ts = cell(num_init, 1);
ys = cell(num_init, 1);
labels = cell(num_init, 1);

%---------------------------------------------------------------------------
% Generate trajectories.
for i = 1 : num_init
  [ ts{i}, ys{i} ] = feval(intfunc, @lorenz, [ 0 tMax ], init(:,i), ...
                           options, params);
  labels{i} = sprintf('IC: [ %f, %f, %f ]', init(:,i));
end
  
%---------------------------------------------------------------------------
% Plot the time histories.
tf = figure;
% These styles are not effective for B&W printing, but work fine on screen.
styles = { 'b-'; 'r-'; 'k-'; 'b--'; 'r--'; 'k--' };
for state = 1 : 3
  subplot(3,1,state)
  hold on;
  for i = 1 : num_init
    plot(ts{i}, ys{i}(:,state), styles{rem(i-1, length(styles)) + 1});
  end
  ylabel([ 'y_' num2str(state) ]);
end
legend(labels);
xlabel('t');

% Plot the phase space results.
pf = figure;
hold on;
for i = 1 : num_init
  plot3(ys{i}(:,1), ys{i}(:,2), ys{i}(:,3), ...
        styles{rem(i-1, length(styles)) + 1});
end
xlabel('y_1');
ylabel('y_2');
zlabel('y_3');
legend(labels);
box on;
grid on;

% We want to show the time histories first, so move that figure back to
% the front.
figure(tf);

% If the user only asked for one initial condition, don't return the
% outputs in cell vectors.
if(num_init == 1)
  ts = ts{1};
  ys = ys{1};
end


%---------------------------------------------------------------------------
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%---------------------------------------------------------------------------
function ydot = lorenz(t, y, params)
% lorenz: Simple model of atmospheric circulation (proposed by Lorenz).
%
%   ydot = lorenz(t, y, params)
%
% ODE function for the simple Lorenz model of atmospheric circulation,
%   and possibly the model's sensitivity matrix.
%   Model adapted from Heath computer problem 9.6.
%
% The length of y determines whether to include the sensitivity ODE:
%   If length(y) = 3, return just the model's dynamics.
%   If length(y) = 12, return dynamics and sensitivity.
%   Otherwise, throw an error.
%
% Ian Mitchell for CPSC 403, 2/6/05.

ydot = zeros(size(y));

if((length(y) ~= 3) && (length(y) ~= 12))
  error('lorenz ODE built only for ODE with 3 or 12 dimensions');
end

%---------------------------------------------------------------------------
% System states.
y1 = y(1);
y2 = y(2);
y3 = y(3);

% System parameters
sigma = params(1);
b = params(2);
r = params(3);

%---------------------------------------------------------------------------
% System dynamics.
ydot(1) = sigma * (y2 - y1);
ydot(2) = r * y1 - y2 - y1 * y3;
ydot(3) = y1 * y2 - b * y3;

%---------------------------------------------------------------------------
% Sensitivity matrix dynamics.
if(length(y) == 12)

  % Unpack current sensitivity matrix.
  P = reshape(y(4:12), 3, 3);

  % Jacobian of dynamics wrt state.
  fy = [ -sigma, +sigma, 0; r - y3, -1, -y1; +y2, +y1, -b ];

  % Jacobian of dynamics wrt parameters.
  fp = diag([ y2 - y1; y1; -y3 ]);

  % Matrix ODE
  Pdot = fy * P + fp;

  % Put it back into vector form.
  ydot(4:12) = Pdot(:);

end