content: {- | Provides learning rate functions - functions that return coefficient which modulates how big are updates of parameters in training.
  -}
  
  
  module Synapse.NN.LearningRates
      ( -- * 'LearningRateFn' type alias and 'LearningRate' newtype
  
        LearningRateFn
      , LearningRate (LearningRate, unLearningRate)
  
        -- * Learning rate decay functions
      
      , exponentialDecay
      , inverseTimeDecay
      , polynomialDecay
      , cosineDecay
      , piecewiseConstantDecay
      ) where
  
  
  -- | 'LearningRateFn' type alias represents functions that return coefficient which modulates how big are updates of parameters in training.
  type LearningRateFn a = Int -> a
  
  
  -- | 'LearningRate' newtype wraps 'LearningRateFn's - functions that modulate how big are updates of parameters in training.
  newtype LearningRate a = LearningRate
      { unLearningRate :: LearningRateFn a  -- ^ Unwraps 'LearningRate' newtype.
      }
  
  
  -- Learning rate decay functions
  
  -- | Takes initial learning rate, decay steps and decay rate and calculates exponential decay learning rate (@initial * decay_rate ^ (step / decay_steps)@).
  exponentialDecay :: Num a => a -> Int -> a -> LearningRateFn a
  exponentialDecay initial steps rate step = initial * rate ^ (step `div` steps)
  
  -- | Takes initial learning rate, decay steps and decay rate and calculates inverse time decay learning rate (@initial / (1 + rate * step / steps))@).
  inverseTimeDecay :: Fractional a => a -> Int -> a -> LearningRateFn a
  inverseTimeDecay initial steps rate step = initial / (1.0 + rate * fromIntegral (step `div` steps))
  
  {- | Takes initial learning rate, decay steps, polynomial power and end decay and calculates polynomial decay learning rate
  (@if step < steps then initial * (1 - step / steps) ** power else end@).
  -}
  polynomialDecay :: Floating a => a -> Int -> a -> a -> LearningRateFn a
  polynomialDecay initial steps power end step
      | step < steps = initial * fromIntegral (1 - min step steps `div` steps) ** power
      | otherwise    = end
  
  {- | Takes initial learning rate, decay steps, alpha coefficient and warmup steps and target (optional) and calculates cosine decay learning rate
  (@(1 - alpha) * (0.5 * (1.0 + cos (pi * step / steps))) + alpha) *
  (if warmup then (if step < warmupSteps then (warmupLR - initial) * step / warmupSteps else warmupLR) else initial)@).
  -}
  cosineDecay :: Floating a => a -> Int -> a -> Maybe (Int, a) -> LearningRateFn a
  cosineDecay initial steps alpha warmup step =
      ((1.0 - alpha) * (0.5 * (1.0 + cos (pi * fromIntegral (step `div` steps)))) + alpha) *
      case warmup of
          Nothing -> initial
          Just (warmupSteps, warmupLR) -> if step < warmupSteps
                                          then (warmupLR - initial) * fromIntegral (step `div` warmupSteps)
                                          else warmupLR
  
  {- | Takes list of boundaries and learning rate values and last rate value for those boundaries and calculates piecewise constant decay learning rate
  (@if step < bound1 then value1 else if step < bound2 then value2 else lastRate@ for @[(bound1, value1), (bound2, value2)]@).
  -}
  piecewiseConstantDecay :: [(Int, a)] -> a -> LearningRateFn a
  piecewiseConstantDecay [] lastRate _ = lastRate
  piecewiseConstantDecay ((stepBound, rateValue):xs) lastRate step
      | step < stepBound = rateValue
      | otherwise        = piecewiseConstantDecay xs lastRate step
  