#!/usr/bin/env python
# Numerically computes Lyapunov Characteristic Exponents for a n-dimensional
# system Copyright (C) 2007 Dale Lukas Peterson The author may be reached at:
# hazelnusse@gmail.com http://www.dlpeterson.com/

# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 51 Franklin
# Street, Fifth Floor, Boston, MA 02110-1301, USA.


import numpy as N
import scipy.integrate as SInt
#import psyco
#psyco.full()

def UnitVec(x):
   """Return a vector with length 1.0 with respect to the n dimensional
   Euclidean 2-norm.
   """
   return x/N.sqrt((x*x).sum())

def InnerProduct(x,y):
   """Standard inner product.
   """
   return (x*y).sum()

def Norm(x):
   """n dimensional Euclidean 2-norm.
   """
   return N.sqrt((x*x).sum())

def GramSchmidt(V):
   """Perform Gram-Schmidt procedure on rows of V, also return the
   stretch/growth factors of each orthogonal vector.  V must be a
   square matrix.
   """
   # Determine the dimension of the vectors
   n = N.size(V[0])
   
   # Initialize a new vector list for storing the orthonormalized vectors, along
   # with a vector to store the norms
   E = N.zeros((n,n))
   d = N.zeros(n)
   
   # Store the first stretch/growth factor, along with the first normalized
   # vector
   d[0] = Norm(V[0])
   E[0] = V[0]/d[0]
   
   # Loop over remaining n-1 vectors, storing the stretch/growth factors as well
   # as the orthonormalized vectors
   for j in range(1,n):
      orthogonal_v = V[j] - N.array([ InnerProduct(V[j],E[k]) * E[k]
                             for k in range(j) ]).sum(axis=0)
      d[j] = Norm(orthogonal_v)
      E[j] = orthogonal_v / d[j]
   return E,d

def ComputeLCE(Dynamic, Jacobian, x0, param, dt, t0, T_pb, N_trans, N_lce,
               atl=1e-8, rtl=1e-8):
   """Compute the Lyapunov Characteristic Exponent for continuous time dynamical
   systems.  The system may be autonomus or non-autonomous.
   
   Inputs:
   
      Dynamic  -- func(x, t, param) computes the derivatives of x at time t
      Jacobian -- the Jacobian matrix of Dynamic, with the same input signature
      x0       -- the initial condition from which to computer the LCE
      param    -- the parameters, if any, that need to be passed to Dynamic and
                  Jacobian
      dt       -- maximum integration stepsize
      t0       -- initial time from which 
      T_pb     -- Period of time to run the differential equations between each
                  pullback
      N_trans  -- Number of pullbacks to perform when allowing transients to
                  decay
      N_lce    -- Number of pullbacks to perform when actually computing the
                  LCE's
      atl      -- Integration absolute error tolerance
      rtl      -- Integration relative error tolerance

   
   Outputs:
      LCE      -- n Lyapunov Characteristic exponents, base - e
   """

   # Determine the dimension of the system
   n = N.size(x0)
   
   # Define the RHS of the n(n+1) dimensional system that includes the flow of
   # the Dynamic as well as the flow of the tangent space. In this formulation,
   # the array x has as its first n entries corresponding to the state of the
   # dynamic, and the remaining n^2 entries corresponding to the state of each
   # of the tangent vectors.
   def FlowAndTangentFlow(x,t,param):
    A = Jacobian(x[0:n],t, param)
    return N.r_[ Dynamic(x[0:n], t, param),
            N.array([N.dot( A, x[(i+1)*n:(i+2)*n]) for i in range(n)]).ravel() ]
   
   # Initial basis vectors for Tangent flow:
   E = N.eye(n)
   
   # Build up the composite state vector for the flow and the tangent flow
   x0 = N.r_[x0, E.ravel()].tolist()
   
   # Iterate the transients out
   for k in range(N_trans):
      #Generate the time points
      t = N.linspace(k*T_pb+t0,(k+1)*T_pb+t0,T_pb/dt+1)
      
      #Integrate the n(n+1) dimensional system of the dynamical system along
      #with each of linearized flows associated with each basis vector
      x = SInt.odeint(FlowAndTangentFlow, x0, t, args = (param,),
                      atol=atl, rtol=rtl)
      
      # Perform Gram-Schmidt -- ignore the growth/shrinkage factors for purpose
      # of LCE calculation here
      V = x[-1,n:].reshape((n,n))
      V_new,d = GramSchmidt(V)
      
      #Get the n-d state after one pullback period, set it as the IC for next
      #loop
      x0 = N.r_[x[-1,0:n], V_new.ravel()]   
   
   t_trans = t[-1].copy()
   LCE_sum = N.zeros(n)
   
   # Iterate and compute the LCE's
   for k in range(N_lce):
      #Generate the time points
      t = N.linspace(k*T_pb+t_trans,(k+1)*T_pb+t_trans,T_pb/dt+1)
      
      #Integrate the n(n+1) dimensional system of the dynamical system along
      #with each of linearized flows associated with each basis vector
      x = SInt.odeint(FlowAndTangentFlow, x0, t, args = (param,),
                      atol=atl, rtol=rtl)
      
      # Perform Gram-Schmidt, accumulate the growth/shrinkage of each basis
      # vector, and sum the natural logarithm of the growth/shrinkage factors
      V = x[-1,n:].reshape((n,n))
      V_new,d = GramSchmidt(V)
      LCE_sum += N.log(d)
      
      #Get the n-d state after one pullback period, set it as the IC for next
      #loop
      x0 = N.r_[x[-1,0:n], V_new.ravel()]
   
   LCE = LCE_sum / (t[-1]-t_trans)
   LCE.sort()
   LCE = N.flipud( LCE )
   
   return LCE

def KaplanYorkeDimension(LCE):
    for j in range(N.size(LCE)):
        if LCE[0:j].sum() < 0:
            break
    
    KYD = j + LCE[0:j-1].sum()/abs(LCE[j])
    return KYD