# L1-3_modeling.py - Lecture 1.2 MATLAB calculations
# RMM, 23 Sep 2012
# 
# This file contains the code to generate the plots contained in CDS
# 101/110 Lecture 1.3 (modeling).

import numpy as np
import matplotlib.pyplot as mpl
from scipy.integrate import odeint

#
# Spring mass system
#

def springmass(y, t, A, omega):
    # Set the parameters
    k1 = 50.; k2 = 50.; k3 = 50. # spring constants
    m1 = 250.; m2 = 250.         # masses
    b = 10.                      # damping

    # compute the input to drive the system
    u = A * np.cos(omega*t)

    # compute the time derivative of the state vector
    dydt = (y[2], y[3], 
            -(k1+k2)/m1*y[0] + k2/m1*y[1],
            k2/m2*y[0] - (k2+k3)/m2*y[1] - b/m2*y[3] + k3/m2*u)

    return dydt

# Call ode45 routine (MATLAB 6 format; help ode45 for details)
tspan = np.linspace(0, 500, 1000)       # time range for simulation
y0 = (0, 0, 0, 0);			# initial conditions
A = 0.00315; omega = 0.75               # amplitude of forcing
sol = odeint(springmass, y0, tspan, (A, omega))
t = tspan

# Plot the input and outputs over entire period
mpl.figure(1); mpl.clf()
mpl.plot(t, A*np.cos(omega*t), t, sol[:,0], t, sol[:,1]);
mpl.show()

# Now plot the data for the final 10# (assuming this is long enough...)
endlen = len(t)/10;                    # last 10% of data record
trange = range(len(t)-endlen, len(t)); # create vector of indices
tend = t[trange];

mpl.figure(2); mpl.clf()
mpl.plot(tend, A*np.cos(omega*tend), tend, sol[trange,0], tend, sol[trange,1]);
mpl.legend(('u', 'q1', 'q2'))
mpl.show()

# Compute the relative phase and amplitude of the signals
#
# This code shows how to compute the magnitude gain and phase change
# for one of he masses (q1) at a single frequency.  You can write a loop
# over different frequencies to plot the full frequency response.
#
# We make use of the fact that we have a sinusoid in steady state,
# as well as its derivative.  This allows us to compute the magnitude
# and phase the sinusoid using simple trigonometry.

# Start by getting the signals and their derivatives (= sin and cos)
u = A*np.cos(omega*tend); udot = -A*omega*np.sin(omega*tend)
y = sol[trange, 0]; ydot = sol[trange, 2]

print "At omega = " + str(omega) + ":"

# Magnitude: use the fact that (A sin(wt))^2 + (A cos(wt))^2 = A^2
# If u = A sin(wt) then udot = A omega cos(wt) => can compute A
ampu = np.mean( np.sqrt(u * u + (udot/omega) * (udot/omega)) )
ampy = np.mean( np.sqrt(y * y + (ydot/omega) * (ydot/omega)) )
print "  Gain = " + str(ampy/ampu)

# Phase: use the atan2(A sin(theta), A cos(theta)) = theta
# If u = A sin(wt) and y = B sin(wt + phi), we can compute phi by
# computing the differences in the phases (watching out for wrapping)
phaseu = np.arctan2(u, udot/omega);
phasey = np.arctan2(y, ydot/omega);
phase_diff = phasey - phaseu
phase_norm = np.arctan2(np.sin(phase_diff), np.cos(phase_diff))
print "  Phase = " + str(np.mean(phase_norm)*180./np.pi) + " deg"

# phase = np.mean( np.atan2(u * u, (udot/omega) * (udot/omega) )
# print "Phase = " + str(ampu*100) + " deg"

#
# Predator prey system
#

# Set up parameters (note that c = a in the model below)
br = 0.6; df = 0.7; a = 0.014;
nperiods = 365;				# simulate each day
duration = 90;				# number of years for simulation

# Set up the initial state
H = np.zeros(duration*nperiods); H[0] = 10; 
L = np.zeros(duration*nperiods); L[0] = 10;

# For simplicity, keep track of the year as well
year = np.zeros(duration*nperiods); year[0] = 1845;

# Iterate the model
Ha = np.zeros(duration); La = np.zeros(duration); 
for k in range(duration*nperiods-1):
    b = br;					# constant food supply
    # b = br*(1+0.5*sin(2*pi*k/(4*nperiods)));	# varying food supply (try it!)
    H[k+1] = H[k] + (b*H[k] - a*L[k]*H[k])/nperiods;
    L[k+1] = L[k] + (a*L[k]*H[k] - df*L[k])/nperiods;
    year[k+1] = year[k] + 1/nperiods;

    if (np.mod(k, nperiods) == 1):
        # Store the annual population
        Ha[k/nperiods] = H[k];
        La[k/nperiods] = L[k];

# Store the final population
Ha[duration-1] = H[duration*nperiods-1];
La[duration-1] = L[duration*nperiods-1];

# Plot the populations of rabbits and foxes versus time
mpl.figure(3); mpl.clf();
mpl.plot(range(1845, 1845 + duration), Ha, '.-', \
         range(1845, 1845 + duration), La, '.--');

# Adjust the parameters of the plot
mpl.xlabel('Year');
mpl.ylabel('Population');

# Now reset the parameters to look like we want
mpl.legend(('hares', 'lynxes'), loc='upper left')
mpl.axis([1845, 1925, 0, 250]);
mpl.show();