Haskell Lyapunov exponent calculations using method of Rangarajan et al. (1998)

Lorenz.hs

module Lorenz where

import Data.Default

data LzParam t = LzParam { sigma :: t, r :: t, b :: t } deriving (Eq, Show)

instance Floating t => Default (LzParam t) where
  def = LzParam 10 28 (8/3)

lorenz :: Floating t => t -> [t] -> [t]
lorenz = lorenzWith def

lorenzWith :: Floating t => LzParam t -> t -> [t] -> [t]
lorenzWith (LzParam sigma r b) _ [x, y, z] =
  [sigma * (y - x)
  , r * x - y - x * z
  , x * y - b * z]

lorenztest.hs

import Numeric.LyapunovExponents
import Numeric.LinearAlgebra
import Lorenz

ic :: State
ic = initialState [1,1,1]

ts :: Vector Double
ts = linspace 10001 (0,1000)

soln :: Matrix Double
soln = snd $ lyapunov lorenz ic ts

main = saveMatrix "lorenz.dat" "%f" $ (fromLists . tail . toLists) soln

LyapunovExponents.hs

{-# LANGUAGE Rank2Types #-}

module Numeric.LyapunovExponents
  ( lyapunov
  , odeSolveRhs    
  , State
  , stateFromList
  , stateToList
  , initialState ) where

import Numeric.LinearAlgebra hiding (dot)
import Numeric.GSL.ODE (odeSolve)
import Numeric.AD (jacobian)
import Numeric.AD.Types (AD, Mode)
import Numeric.AD.Internal.Classes (lift)
import qualified Data.List as L
import Data.Default
import Data.MemoTrie

------------------------------------------------------------------------------
--
--  SYMBOLIC REPRESENTATION OF ROTATION MATRICES AND THEIR DERIVATIVES
--
------------------------------------------------------------------------------

-- Entries in the matrices we need to consider are made up of terms
-- composed of factors involving sines, cosines and derivatives of the
-- n(n-1)/2 rotation angles needed to specify a full rotation in
-- n-dimensional space.

-- Throughout, we need to refer to elements of the lower-triangular
-- part of matrices (all the matrices we deal with are either
-- symmetric or skew-symmetric, so we deal with the diagonal and the
-- lower triangular elements separately).  For an n x n matrix, there
-- are n(n-1)/2 elements in the lower-triangular part, which we label
-- with integers 1..n(n-1/2) in the order given here.
--
idxs :: Int -> [((Int,Int),Int)]
idxs n = zip [(j, i) | i <- [1..n], j <- [2..n], i < j] [1..]


-- Individual Factors

-- Representation for individual factors within matrix expressions.
-- Angle theta_{ij} represents a two-dimensional rotation by the given
-- angle in the plane spanned by axes x_i and x_j.  (For simplicity,
-- we actually label the angles using a single index that ranges from
-- 1 to n(n-1)/2 using the scheme defined in idxs above.)
--
-- Factors: C (i,j) = cos theta_{ij}
--          S (i,j) = sin theta_{ij}
--          D (i,j) = d/dt theta_{ij}
--
data Factor = Null | C Int | S Int | D Int deriving (Eq, Ord)


-- Terms
  
-- Terms are products of factors, with an integer coefficient.
--
data Term = Term { coef :: Integer, facs :: [Factor] } deriving (Eq)

-- Defined for convenience when viewing expressions.
--
instance Ord Term where
  compare (Term _ fs1) (Term _ fs2) | l1 /= l2 = compare l1 l2
                                    | otherwise = compare fs1 fs2
    where l1 = length fs1
          l2 = length fs2

-- Multiplication of terms.
--
fmult :: Term -> Term -> Term
fmult (Term c1 fs1) (Term c2 fs2) = Term (c1 * c2) (L.sort $ fs1 ++ fs2)


-- Expressions

-- Expressions are sums of terms.
--
newtype Expr = Expr [Term] deriving (Eq, Ord)

-- Simple Num instance for expressions.
--
instance Num Expr where
  fromInteger n = Expr [Term n []]
  negate (Expr ts) = Expr $ map (\(Term c fs) -> Term (-c) fs) ts
  (Expr ts1) + (Expr ts2) = simplify $ Expr $ L.sort (ts1 ++ ts2)
  (Expr ts1) * (Expr ts2) = 
    simplify $ Expr $ L.sort [fmult t1 t2 | t1 <- ts1, t2 <- ts2]
  signum _ = error "signum not appropriate for Expr"
  abs _ = error "abs not appropriate for Expr"

-- Simplification of expressions.  The simplifications performed are
-- restricted to removal of zero terms, merging of terms involving
-- identical factors, and (the most complex) simplification of
-- trignometric sums using the Pythagoras's theorem.  Simplification
-- is repeated, using these three steps, until no further changes can
-- be made.
--
simplify :: Expr -> Expr
simplify (Expr ts) = Expr $ firstFix (iterate step ts)
  where firstFix (x1:x2:xs) | x1 == x2 = x1
                            | otherwise = firstFix (x2:xs)
        step = dropZeros . mt . trig                                          
        dropZeros = filter ((/=0) . coef)
        mt [] = []              -- Term merging.
        mt [t] = [t]
        mt (t1:(t2:ts)) 
          | facs t1 == facs t2 = mt $ (Term (coef t1 + coef t2) (facs t1)):ts
          | otherwise = t1 : (mt $ t2:ts)

-- Trigonometric simplification (we only do one step here, since
-- simplify repeats simplification steps until we reach a fixed
-- point).
--
trig ts = go ts
  where go [] = []
        go (t:ts) | null reps = t : go ts
                  | otherwise = case match reps t ts of
                    Nothing -> t : go ts
                    Just (f, mt, ts') -> combine f t mt ++ ts'
                    where reps = repeats t
        
        -- Find repeated factors (i.e. cos^2 or sin^2) in a term.
        repeats (Term _ fs) = map head . filter ((>1) . length) $ L.group fs

        -- Match repeated factors in a term against repeated factors
        -- in following terms.
        match [] m ts = Nothing
        match (f:fs) m ts = case matchOne f m ts of
          Nothing -> match fs m ts
          Just (f, mt) -> Just (f, mt, L.delete mt ts)
        
        -- Match a single repeated factor against following terms.
        matchOne f m [] = Nothing
        matchOne f m (t:ts) | matchFac f `elem` repeats t && 
                              filter (/=f) (facs m) == 
                              filter (/=matchFac f) (facs t) &&
                              signum (coef m) == signum (coef t) = Just (f, t)
                            | otherwise = matchOne f m ts
        
        -- cos^2 can be paired off with sin^2, and vice versa.
        matchFac (C x) = S x
        matchFac (S x) = C x
        
        -- Combine terms with compatible repeated elements.
        combine f (Term c1 fs1) (Term c2 fs2) = case compare c1 c2 of
          EQ -> [Term c1 fs1']
          LT -> [Term c1 fs1', Term (c2-c1) fs2]
          GT -> [Term c2 fs2', Term (c1-c2) fs1]
          where fs1' = L.delete f $ L.delete f fs1
                fs2' = L.delete mf $ L.delete mf fs2
                mf = matchFac f

-- Helper functions for creating elementary expressions.
--
c, s, d :: Int -> Expr
c i = Expr [Term 1 [C i]]
s i = Expr [Term 1 [S i]]
d i = Expr [Term 1 [D i]]

-- Symbolic differentiation of expressions using the product rule and
-- Leibnitz's rule.
--
diff :: Expr -> Expr
diff (Expr ts) = L.foldl1' (+) $ map diffTerm ts
  where diffTerm (Term c fs) = L.foldl1' (+) $ map single $
                                   zipWith fmult des leaveOneOuts 
          where des = map diffFac fs
                leaveOneOuts = map (Term c) 
                               [take n fs ++ drop (n + 1) fs | 
                                n <- [0..length fs - 1]]
                single t = Expr [t]
        
        diffFac (D _) = error "No second derivatives!"
        diffFac (S pos) = Term 1 [C pos, D pos]
        diffFac (C pos) = Term (-1) [S pos, D pos]


-- Numeric evaluation of expressions.
--
eval :: (Vector Double, Vector Double) -> Expr -> Double
eval (ctheta, stheta) (Expr ts) = sum $ map evalTerm ts
  where evalTerm (Term c fs) = fromInteger c * (product $ map evalFactor fs)
        evalFactor (C i) = ctheta @> (i-1)
        evalFactor (S i) = stheta @> (i-1)
        evalFactor _ = error "Can't evaluate this factor!"


-- Simple Matrix Operations

-- We call this type (which we use only for our symbolic manipulations
-- to build the T matrix) "Mat" to avoid conflicts with hmatrix types.
--
newtype Mat a = Mat [[a]] deriving (Eq)

instance Show a => Show (Mat a) where
  show (Mat xs) = L.intercalate "\n" $ map show xs

instance Functor Mat where
  fmap f (Mat xs)  = Mat $ map (map f) xs

transpose :: Mat a -> Mat a
transpose (Mat xs) = Mat (L.transpose xs)

dot :: Num a => [a] -> [a] -> a
dot v1 v2 = sum $ zipWith (*) v1 v2

(%*%) :: Num a => Mat a -> Mat a -> Mat a
(Mat m1) %*% (Mat m2) = Mat [[r `dot` c | c <- L.transpose m2] | r <- m1]


-- Rotation Matrices

-- Elementary rotation matrices.
--
elemRot :: Int -> ((Int,Int),Int) -> Mat Expr
elemRot n ((i,j),idx) = Mat [[entry (k,l) | l <- [1..n]] | k <- [1..n]]
  where entry (k,l) | k == l && k /= i && k /= j   = 1
                    | k == l && (k == i || k == j) = c idx
                    | k == i && l == j             = s idx
                    | k == j && l == i             = -s idx
                    | otherwise                    = 0

-- Full rotation matrix.
--
rot :: Int -> Mat Expr
rot = memo rot'
  where rot' n = L.foldl1' (%*%) $ map (elemRot n) (idxs n)

-- Convert from expression form to numeric values by passing in vector
-- of angles.
--
evalRot :: Mat Expr -> Vector Double -> Matrix Double
evalRot q theta = fromLists matrep
  where Mat matrep = fmap (eval cstheta) q 
        cstheta = (mapVector cos theta, mapVector sin theta)


------------------------------------------------------------------------------
--
--  LYAPUNOV EXPONENT ODE EXPRESSIONS
--
------------------------------------------------------------------------------

-- Calculate Q^T \dot{Q} symbolically for a given system dimension.
--
qtqdot :: Int -> Mat Expr
qtqdot n = (transpose q) %*% (fmap diff q)
  where q = rot n

-- Extract lower triangular entries from Q^T \dot{Q} as expressions:
-- used to build T(\theta) matrix to solve for \dot{\theta}.
--
qtqdotEntries :: Int -> [Expr]
qtqdotEntries n = concat [[qq !! r !! c | c <- [0..r-1]] | r <- [1..n-1]]
  where Mat qq = qtqdot n

-- Calculate T(\theta).
--
tMatrix :: Int -> Mat Expr
tMatrix = memo tMatrix'
  where tMatrix' n = Mat $ map makeRow es
          where es = qtqdotEntries n
                makeRow (Expr ts) = map (makeEntry ts) [1..size]
                makeEntry ts i = Expr . map (deleteD i) . filter (hasD i) $ ts
                size = n * (n-1) `div` 2
                hasD i (Term _ fs) = (D i) `elem` fs
                deleteD i (Term c fs) = Term c $ L.delete (D i) fs

-- The full system state for solving for the Lyapunov exponents
-- involves the original system state (x), the scaled Lyapunov
-- exponents (lams) and the rotation angles (ths).
--
data State = State { x :: Vector Double
                   , lams :: Vector Double
                   , ths :: Vector Double }
             deriving (Eq, Show)
                      
-- Calculate Q^T . \nabla F . Q, the right hand side of the evolution
-- equation for the fundamental solution matrix.
--
rhs :: (forall s. Mode s => AD s Double -> [AD s Double] -> [AD s Double]) -> 
       Double -> State -> Matrix Double
rhs f t (State x lams ths) = (trans q) <> jac <> q
  where q = evalRot (rot n) ths
        jac = fromLists $ jacobian (f $ lift t) (toList x)
        n = dim x

-- Calculate the full right hand side for the evolution of our
-- enhanced system state.
--
fullRhs :: (forall t. Floating t => t -> [t] -> [t]) -> 
           Double -> State -> Vector Double
fullRhs f t s@(State x lams ths) = join [xdot, lamdot, thdot]
  where n = dim x
        -- Time derivative from original system.
        xdot = fromList . f t . toList $ x
        
        -- RHS of evolution equation for fundamental solution matrix.
        r = rhs f t s
        
        -- Time derivatives for scaled LEs.
        lamdot = takeDiag r
        
        -- T(\theta) for current \theta values.
        tm = evalRot (tMatrix n) ths
        
        -- Lower triangular entries from the RHS of the evolution
        -- equation for the fundamental solution matrix, corresponding
        -- to \dot{\theta} entries.
        tmrhs = fromList [r @@> (i-1,j-1) | ((i,j),_) <- idxs n]
        
        -- Solve for \dot{\theta} using T(\theta) and the lower
        -- triangular part of r.
        thdot = tm <\> tmrhs

-- Convert State values to and from lists for use with odeSolve
--
stateFromList :: [Double] -> State
stateFromList lx = State (fromList x) (fromList lam) (fromList th)
  where (x,lamth) = splitAt n lx
        (lam,th) = splitAt n lamth
        n = (isqrt (9 + 8 * length lx) - 3) `div` 2
        isqrt = floor . sqrt . fromIntegral

stateToList :: State -> [Double]
stateToList (State x lam th) = concat [toList x, toList lam, toList th]

-- Set up initial State (including LEs and angles) given an initial
-- condition for the underlying system.
--
initialState :: [Double] -> State
initialState lx = State (fromList lx) (constant 1 nlam) (constant 0 nth)
  where nlam = length lx
        nth = nlam * (nlam - 1) `div` 2

-- Produce a function suitable for use with Numeric.GSL.ODE.odeSolve.
--
odeSolveRhs :: (forall t. Floating t => t -> [t] -> [t]) -> 
               (Double -> [Double] -> [Double])
odeSolveRhs sf t lx = toList $ fullRhs sf t (stateFromList lx)


------------------------------------------------------------------------------
--
--  PUBLIC INTERFACE
--
------------------------------------------------------------------------------

lyapunov :: (forall t. Floating t => t -> [t] -> [t]) -> 
            State -> Vector Double -> (State, Matrix Double)
lyapunov f istate ts = (fstate, fromColumns $ ts : cols)
  where odef = odeSolveRhs f
        soln = odeSolve odef (stateToList istate) ts
        cols = zipWith ($) fs (toColumns soln)
        fs = replicate n id ++ 
             replicate n (`divide` ts) ++ 
             replicate (n * (n - 1) `div` 2) id
        n = dim (x istate)
        fstate = stateFromList . toList . last . toRows $ soln

------------------------------------------------------------------------------
--
--  DEBUGGING CODE
--
------------------------------------------------------------------------------

instance Show Factor where
  show Null = error "null factor!"
  show (C i) = "C" ++ show i
  show (S i) = "S" ++ show i
  show (D i) = "D" ++ show i

instance Show Term where
  show (Term c fs) = coefstr ++ L.intercalate " " (map show fs)
    where coefstr | c == 1 = ""
                  | c == -1 = "-"
                  | otherwise = show c ++ " * "
  
instance Show Expr where
  show (Expr []) = "0"
  show (Expr ts) = concat $ zipWith (++) (signs ts) tss
    where tss = map showf ts
          signs [] = []
          signs (t:ts') = (leadingSign t) : map trailingSign ts'
          leadingSign (Term c fs) | c < 0     = "-"
                                  | otherwise = ""
          trailingSign (Term c fs) | c < 0     = " - "
                                   | otherwise = " + "
          showf (Term c []) = show (abs c)
          showf (Term c fs) = coefstr c ++ L.intercalate " " (map show fs)
          coefstr c | abs c == 1 = ""
                    | otherwise = show (abs c) ++ " "
    

VDP.hs

module VDP where

import Data.Default

data VdpParam t = VdpParam { d :: t, b :: t, omega :: t } deriving (Eq, Show)

instance Floating t => Default (VdpParam t) where
  def = VdpParam (-5) 5 2.466

vdp :: Floating t => t -> [t] -> [t]
vdp = vdpWith def

vdpWith :: Floating t => VdpParam t -> t -> [t] -> [t]
vdpWith (VdpParam d b omega) t [z1,z2] = [dz1,dz2]
  where dz1 = z2
        dz2 = -d * (1 - z1**2) * z2 - z1 + b * cos (omega * t)

vdptest.hs

import Numeric.LyapunovExponents
import Numeric.LinearAlgebra
import VDP

ic :: State
ic = initialState [1,1]

ts :: Vector Double
ts = linspace 1001 (0,1000)

soln :: Matrix Double
soln = snd $ lyapunov vdp ic ts

main = saveMatrix "vdp.dat" "%f" $ (fromLists . tail . toLists) soln