W' balance modelling with Haskell

README.md

Haskell W′ modelling

An implementation of Skiba et al. (2012)'s W′ model in Haskell.

test.sh

#!/bin/bash
stack wprime.hs | gnuplot -p -e 'plot "/dev/stdin" using 0:1 with lines'

wprime.hs

#!/usr/bin/env stack
{- stack script --resolver lts-9.13 -}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NamedFieldPuns             #-}
import qualified Data.DList as DL

infixl 0 #
(#) :: a -> (a -> c) -> c
(#) = flip ($)

-- DATA -----------------------------------------------------------------------

-- Aliases for SI units.
-- NOTE: Not deriving from Num and friends because callers
-- should have to explicitly use value constructors, rather
-- than use numeric literals via fromInteger/fromRational!
newtype Power  = Watts   { watts   :: Double  } deriving (Eq, Ord)
newtype Energy = Joules  { joules  :: Double  } deriving (Eq, Ord)
newtype Time   = Seconds { seconds :: Integer } deriving (Enum)

-- Helpful synonym.
type Timestamp = Time
type DiffTime  = Time  -- NOTE: we don't need Data.Time.Clock's precision

--  A "segment of an exercise session".
data Segment = Segment { segStart :: Timestamp
                       , segEnd   :: Timestamp
                       , segPwr   :: Power     -- mean
                       }

-- Exercise intensity domains. We don't actually use moderate here,
-- it's included for completeness.
data Domain
    = Moderate
    | Heavy
    | Severe
    deriving (Eq, Ord)

-- Athlete descriptors.
data Athlete = Athlete { athleteCP  :: Power
                       , athleteLT  :: Power  -- not currently used
                       , athleteW'  :: Energy
                       , athleteTau :: Power -> Double -- :: DCP -> Tau
                       }

-- DATA HELPERS ---------------------------------------------------------------

-- Returns the length of a segment in seconds.
segLen :: Segment -> DiffTime
segLen Segment{ segStart, segEnd } =
    Seconds $ seconds segEnd - seconds segStart

-- Returns the intensity domain of a segment for a given athlete.
segDomain :: Athlete -> Segment -> Domain
segDomain Athlete{ athleteCP, athleteLT } Segment{ segPwr }
    | segPwr > athleteCP = Severe
    | segPwr > athleteLT = Heavy
    | otherwise          = Moderate

-- Skiba et al., 2012
defaultTau :: Power -> Double
defaultTau (Watts dcp) = 546 * exp (negate 0.01 * dcp) + 316

-- CALCULATIONS ---------------------------------------------------------------

-- Calculate W' balance for consecutive segments of an exercise session.
calcBalance :: Athlete -> [Segment] -> [Energy]
calcBalance athlete@Athlete{ athleteW' } = DL.toList . calcFrom athleteW'
    where
        calcFrom :: Energy -> [Segment] -> DL.DList Energy
        calcFrom _ [] = DL.empty
        calcFrom bal segs@(seg1:_) =
            let
                (section, rest) =
                    span ((== domain seg1). domain) segs

                bals =
                    if domain seg1 >= Severe
                        then scanl (expendW' athlete) bal section
                        else map (recoverW' athlete seg1 bal) section
            in
                DL.fromList bals `DL.append` calcFrom (last bals) rest

        domain :: Segment -> Domain
        domain = segDomain athlete   -- shorthand

-- Discharge function.
-- NOTE: can't drop below zero.
expendW' :: Athlete -> Energy -> Segment -> Energy
expendW' a (Joules bal) seg = Joules $ max 0 $ bal - (pwr - cp) * dt
    where pwr = watts . segPwr $ seg
          cp  = watts . athleteCP $ a
          dt  = fromIntegral . seconds . segLen $ seg

-- Recharge function.
recoverW' :: Athlete -> Segment -> Energy -> Segment -> Energy
recoverW' Athlete{ athleteCP, athleteW', athleteTau }
          Segment{ segStart = u }
          (Joules bal)
          Segment{ segEnd = t, segPwr }
    | bal < w'  = Joules $ w' - wexp * exp (negate tu / athleteTau dcp)
    | otherwise = Joules w'
    where
        pwr   = segPwr    # watts
        cp    = athleteCP # watts
        tu    = fromIntegral $ seconds t - seconds u
        dcp   = Watts $ cp - pwr
        w'    = athleteW' # joules
        wexp  = w' - bal  -- W' expended

-- TESTING --------------------------------------------------------------------

exampleAthlete :: Athlete
exampleAthlete =
    Athlete { athleteCP  = 300  # Watts
            , athleteLT  = 200  # Watts
            , athleteW'  = 21e3 # Joules
            , athleteTau = defaultTau
            }

-- Second time addition.
tadd :: Time -> Time -> Time
tadd (Seconds x) (Seconds y) = Seconds (x + y)

exampleSegments :: [Segment]
exampleSegments =
    segs (Watts 350) (Seconds 0) (Seconds 800) (Seconds 1)
        ++ segs (Watts 100) (Seconds 801) (Seconds 1200) (Seconds 1)
    where segs pwr start stop step =
            [ Segment{segStart = i, segEnd = i `tadd` step, segPwr = pwr}
            | i <- [start, start `tadd` step..start `tadd` stop]
            ]

instance Show Energy where
    show (Joules e) = (++" kJ") . show $ e / 1000

main :: IO ()
main = do
    let bals = calcBalance exampleAthlete exampleSegments
    mapM_ print bals