text:
{-# language TypeSynonymInstances #-}
{-# language LiberalTypeSynonyms #-}

module FrequencyResponse
  (Analyzer (..),
  AnalyzerAcyclic (..),
  Spectrum (..),
  SpecLin (..),
  SpecWide ,
  getResponse,
  getResponseLin,
  getResponseAcyclic)

where

import CLaSH.Prelude (Default(..))
import Prelude
import CLaSH.Signal.MultiSignal
import Data.Complex as C
import Data.Maybe
import Control.Applicative
import Data.Map.Strict as Map

-- | Type representing input and output signal in circuitry
-- p of kind (* -> *) represents underlaying Signal of type Prependable (ZipList for example)
-- s is type of Spectrum
data Analyzer p s = Analyzer {unAnalyzer :: ([Maybe (s -> s)] , p s)}

instance (Functor p, Prependable p) => Prependable (Analyzer p) where
     prepend a da = Analyzer (Nothing : f , (prepend a $ fmap fx d))
      where
        (f,d) = unAnalyzer da
        fx = fromMaybe id $ foldr1 (<|>) f

instance (Applicative p,Num n) => Num (Analyzer p n) where
    (+) = liftA2FA (+)
    (-) = liftA2FA (-)
    (*) = liftA2FA (*)
    fromInteger = makeAnalyzer . fromInteger
    negate = liftA1FA negate
    abs = liftA1FA abs
    signum = liftA1FA signum

instance (Applicative p,Fractional n) => Fractional (Analyzer p n) where
    (/) = liftA2FA (/)
    fromRational = makeAnalyzer . fromRational

liftA2FA op da db = Analyzer (zipWith (<|>) fa fb, liftA2 op a b) where
       (fa,a) = unAnalyzer da
       (fb,b) = unAnalyzer db

liftA1FA op f = Analyzer (repeat Nothing , fmap op s)   where
   (_,s) = unAnalyzer f

makeAnalyzer s = Analyzer ( repeat Nothing , pure s )
{- inline liftA2FA, liftA1FA-}

-- | spectrum is result of frequency mixing two spectrums
-- n represents type represented as harmonics (Int or better, Rational)
-- and a is type of complex number (frequently Double)
data Spectrum n a = Spectrum {getSpectrum :: Map.Map n (Complex a)}

instance (Ord n,Num n,Default a,RealFloat a) => Default (Spectrum n a) where
    def = dcSpectrum def

instance (Ord n,Num n,RealFloat a) => Num (Spectrum n a) where
    (+) = dotMergeSpec (+)
    (-) = dotMergeSpec (-)
    (*) = mixSpectrum
    negate = fmapSpec negate
    abs = error "no abs function for spectrum available"
    signum = error "no sig num function for spectrum available"
    fromInteger = dcSpectrum . fromInteger

instance (Show n, Show a, Ord n,Num n,Fractional a,RealFloat a) => Fractional (Spectrum n a) where
    (/) = spectrumDiv
    fromRational = dcSpectrum . fromRational

dotMergeSpec op (Spectrum sa) (Spectrum sb) = Spectrum $ (Map.unionWith op) sa sb
fmapSpec f = Spectrum . fmap f . getSpectrum
constSpectrum freq a = Spectrum $ Map.singleton freq (a :+ 0 )
dcSpectrum a = constSpectrum 0 a

mixSpectrum (Spectrum sa) (Spectrum sb) = Spectrum
        $ Map.fromListWith (+)
        $ concat [mixFeq a b | a <- Map.toList sa , b <- Map.toList sb]

mixFeq (fa,pa) (fb,pb)
   | fa > fb = mixFeqSorted (fb,pb) (fa,pa)
   | True =  mixFeqSorted (fa,pa) (fb,pb)

mixFeqSorted (0,pa) (fb,pb) = [(fb,pa*pb)]
mixFeqSorted (fa,pa) (fb,pb) = [(fb - fa,r1),(fb+fa,r2)] where
   (ma,pha) = polar pa
   (mb,phb) = polar pb
   k = ma * mb / 2
   r1
      | fa == fb  = mkPolar (k * sin (phb - pha + pi/2)) 0
      | otherwise = mkPolar k (phb - pha + pi/2)
   r2 = mkPolar ((-1)*k) (phb + pha + pi/2)

delaySpec ph (Spectrum s) = Spectrum $ Map.mapWithKey f s where
    f 0 a = a
    f k a = a * mkPolar 1 (fromRational k * ph)

spectrumDiv (Spectrum a) sb = Spectrum (fmap (/ dc) a) where
    dc = getDcOrError sb

getDcOrError (Spectrum a) =
     case toList a of
            ((0,dc):[]) -> dc
            s  -> error ("spectrum is not dc, but" ++ show s )

makeDcSig :: (Applicative p, Num n, Num a) => a -> Analyzer p (Spectrum n a)
makeDcSig = makeAnalyzer . dcSpectrum

-- | spectrum for linear system, faster version of Spectrum
-- underlaying type is tuple insted of Map
newtype SpecLin = SpecLin {getSpecLin :: (Double,Complex Double)}

instance  Default SpecLin where
    def = SpecLin (def,def :+ def)

instance  Num SpecLin where
    (+) = dotMergeSpecLin (+)
    (-) = dotMergeSpecLin (-)
    (*) = mixSpecLin
    negate = fmapSpecLin negate
    abs = error "no abs function for spectrum available"
    signum = error "no signum function for spectrum available"
    fromInteger = dcSpecLin . fromInteger

instance Fractional SpecLin where
    (/) = undefined -- dcSpecLin
    fromRational = dcSpecLin . fromRational

dotMergeSpecLin op (SpecLin (dcA,acA)) (SpecLin (dcB,acB)) = SpecLin (op dcA dcB , acr) where
    acr = op (realPart acA) (realPart acB) :+ op (imagPart acA) (imagPart acB)

mixSpecLin (SpecLin (dcA,acA)) (SpecLin (dcB,acB))
   | (acA /= 0 ) && (acB /= 0) = error "can not mix two linear spectrum"
   | otherwise = SpecLin (dcr , acr) where
       dcr = dcA * dcB
       acr = fmap (* dcA) acB + fmap (* dcB) acA

dcSpecLin dc = SpecLin (dc,0)

fmapSpecLin f (SpecLin (dc,ac)) =  SpecLin (f dc,fmap f ac)

delayLinSpec ph (SpecLin (dc,ac)) = SpecLin (dc,ac * mkPolar 1 ph)
-- end of SpecLin

-- for acyclic circuits we can implement faster version of Prependable
-- It has nice property that first value of output is correct.
data AnalyzerAcyclic s = AnalyzerAcyclic {unAnalyzerAcyclic :: (Maybe (s -> s), s)}

-- Prependable instance has funny laws (no laws?)
-- it does not follow list isomorphism, because it is single value
instance Prependable AnalyzerAcyclic where
     prepend _ da = AnalyzerAcyclic (f , fx d)
      where
        (f,d) = unAnalyzerAcyclic da
        fx = fromMaybe id f

instance (Num n) => Num (AnalyzerAcyclic n) where
    (+) = opAcyclic (+)
    (-) = opAcyclic (-)
    (*) = opAcyclic (*)
    fromInteger = pureAcyclic . fromInteger
    negate = liftAcyclic negate
    abs = liftAcyclic abs
    signum = liftAcyclic signum

instance (Fractional n) => Fractional (AnalyzerAcyclic n) where
    (/) = opAcyclic (/)
    fromRational = pureAcyclic . fromRational

pureAcyclic n = AnalyzerAcyclic (Nothing,n)
opAcyclic op (AnalyzerAcyclic (fa,a)) (AnalyzerAcyclic (fb,b)) = (AnalyzerAcyclic (fa <|> fb ,op a b))
liftAcyclic f (AnalyzerAcyclic (fa,a)) = (AnalyzerAcyclic (fa,f a))

-- end of AnalyzerAcyclic

type SpecWide = Spectrum Rational Double

-- | run analysis for both acyclic\/cyclic and both linear\/nonlinear circuits
-- It is required to skip first N resulting values
-- where N represent delay of circuitry (approx total number of registers)
getResponse :: (Analyzer ZipList SpecWide -> Analyzer ZipList SpecWide) -- ^ circuitry
            -> Double  -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
            -> [Spectrum Rational Double] -- ^ response as stream
getResponse f ph = getZipList $ snd $ unAnalyzer $ f $ Analyzer ([ Just (delaySpec ph)] , pure (constSpectrum 1 1))

-- | run analysis for both acyclic\/cyclic and only linear circuits
-- faster version of getResponse for acyclic linear ciruits
getResponseLin :: (Analyzer ZipList SpecLin -> Analyzer ZipList SpecLin) -- ^ circuitry
          -> Double -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
          -> [SpecLin] -- ^ response as stream
getResponseLin f ph = getZipList $ snd $ unAnalyzer $ f $ Analyzer ([Just (delayLinSpec ph)] , pure (SpecLin (0,1 :+ 0)))

-- | run analysis both linear\/nonlinear and only acyclic circuit
-- calculate response of acyclic circuit
getResponseAcyclic :: (AnalyzerAcyclic SpecWide -> AnalyzerAcyclic SpecWide) -- ^ circuitry
                      -> Double -- ^ phase in radians that represents Nyquist frequency. (phase pi ≡ 1.0 of sampling frequency)
                      -> Spectrum Rational Double -- ^ response as stream
getResponseAcyclic f ph =  snd $ unAnalyzerAcyclic $ f $ AnalyzerAcyclic (Just (delaySpec ph) ,constSpectrum 1 1)


text:
{-# OPTIONS -cpp #-}
module Main where

-- import MathObj.FFT.Radix4Par as R4 (fft,fftDo2,fftDo3)
import MathObj.FFT.Radix4Seq as R4 (fft)
import MathObj.FFT.Reference as R2 (fft, ifft)
import MathObj.FFT.MultiPar (twoDimFft, twoDimFftG, twoDimFftRD)
import Prelude
-- import System.CPUTime
import Data.Complex
import MathObj.FFT.Complex
import Control.Concurrent
#ifdef TEST
import qualified MathObj.FFT.Reference as Ref
#endif

import System.Environment (getArgs)
import System.IO
import Control.Parallel.Eden (rnf, Trans)
import Control.Parallel.Eden.Map (farmS)

instance (RealFloat a, Trans a) => Trans (Complex a) -- ???

main = do
  hSetBuffering stdout NoBuffering
  putStrLn "usage: testfftTest <size or base for power of 4(1024) of one side> <version 0,1,2,3,4(3)> <chunksize(512), no effect?> <prefetch(20)>"
  args <- getArgs
  let f = mkLists n
--      n = if null args then 512 else read (head args)
      n = if b > 15 then b else 4^b
      b = if null args then 1024 else read (head args)
      v = if length args < 2 then 3 else read (args!!1)
      s = if length args < 3 then 512 else read (args!!2)
      p = if length args < 4 then 20 else read (args!!3)
--  let r = fftDo3 s (n, f)
--      r2 = fftDo2 s (n,f)
      r4 = twoDimFft n n f
      r5 = twoDimFftG map R4.fft n n f
      r6 = twoDimFftG (farmS p) R2.fft n n f
      r7 = twoDimFftG (farmS p) R4.fft n n f
      r8 = twoDimFftRD (farmS p) R4.fft n n f
#ifdef TEST
  let (len,resultlist) = r
  putStrLn ("Length given as " ++ show len ++
	    ", in reality " ++ show (length resultlist))
  let ref = Ref.fft (n,f)
  if (r /= ref) then putStrLn "WRONG RESULT" else putStrLn "CORRECT"
  rnf r2 `seq`
--  (length (snd r)) `seq`
--  (fst r ) `seq`
    putStrLn "Done version fftDo2"
  threadDelay 100000
  rnf r `seq`
--  (length (snd r)) `seq`
--  (fst r ) `seq`
    putStrLn "Done version fft"
  if r == r2
     then putStrLn "same result"
     else outputDiff 5 r r2
#else
  case v of
--     0 -> rnf r `seq` putStrLn ("Done, version fft")
--     1 -> rnf r2 `seq` putStrLn ("Done, version fftDo2")
     2 -> putStrLn ("not defined in this version, look in code/small/")
     3 -> rnf r4 `seq` putStrLn ("Done, version twoDimFft with seqmap R2 (sequential ground truth for R2)")
     4 -> rnf r5 `seq` putStrLn ("Done, version twoDimFft with seqmap R4 (sequential ground truth for R4)")
     5 -> rnf r6 `seq` putStrLn ("Done, version twoDimFft with farm R2")
     6 -> rnf r7 `seq` putStrLn ("Done, version twoDimFft with farm R4")
     7 ->rnf r7 `seq` putStrLn ("Done, version twoDimFftRD(!) with farm R4")
     _ -> putStrLn ("no version " ++ show v ++ " defined")
#endif

mkLists :: Int -> [[Complex Double]]
mkLists n = map (\k -> lift $ [k..k+n]) [1..n]

outputDiff :: (Show a,Eq a,Num a) => Int -> (Int,[a]) -> (Int,[a]) -> IO ()
outputDiff n (n1,l1) (n2,l2)
    | n1 /= n2  = putStrLn "different length"
    | otherwise = do putStrLn (show (length diffs) ++ " different elements")
		     putStrLn (unlines (map show (take n diffs)))
    where diffs = [ (k,x,y,x-y) | (k,x,y) <- zip3 [0..] l1 l2, x/=y ]

ext: hs

max_stars_repo_name: Parrows/Parrows

max_issues_repo_name: Parrows/Parrows

max_forks_repo_name: Parrows/Parrows


text:
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
{-# LANGUAGE StandaloneKindSignatures #-}
module Inference.Conjugate where

import           Control.Monad.Primitive        ( PrimMonad
                                                , PrimState
                                                )
import           Control.Monad.Reader           ( ReaderT
                                                , runReaderT
                                                )
import           Control.Monad.Reader.Class     ( MonadReader(ask) )
import           Control.Monad.State            ( State
                                                , StateT
                                                , evalState
                                                , evalStateT
                                                , execStateT
                                                , runStateT
                                                )
import           Control.Monad.State.Class      ( get
                                                , modify
                                                , put
                                                )
import           Control.Monad.Trans.Maybe      ( MaybeT(MaybeT)
                                                , runMaybeT
                                                )
import           Control.Monad.Writer
import           Data.Dynamic                   ( Dynamic
                                                , Typeable
                                                , fromDynamic
                                                , toDyn
                                                )
import           Data.Kind
import           Data.Maybe                     ( fromMaybe )
import           Data.MultiSet                 as MS
import qualified Data.Sequence                 as S
import           Data.Typeable                  ( Proxy(Proxy)
                                                , typeRep
                                                )
import qualified Data.Vector                   as V
import qualified Debug.Trace                   as DT
import           GHC.Float                      ( int2Double )
import           GHC.Generics
import           GHC.TypeNats
import           Lens.Micro
import           Lens.Micro.Extras
import           Lens.Micro.TH                  ( makeLenses )
import           Numeric.SpecFunctions          ( logBeta
                                                , logChoose
                                                , logFactorial
                                                , logGamma
                                                )
import           System.Random.MWC.Probability
                                         hiding ( Uniform )

----------------------------------------------------------------
-- type classes and families for distributions and conjugates --
----------------------------------------------------------------

-- | Describes a family of distributions with a fixed form.
-- For example, a 'Bernoulli' distribution is parameterized by a probability @p@
-- and produces binary samples
-- (@True@ with probability @p@, @False@ with probability @1-p@).
--
-- Its 'Distribution' instance is:
-- > instance Distribution Bernoulli where
-- >   type Params Bernoulli = Double
-- >   type Support Bernoulli = Bool
-- >   distSample _ = uncurry bernoulli
-- >   distLogP _ p True = log p
-- >   distLogP _ p False = log (1 - p)
class Distribution a where
  type Params a :: Type
  type Support a :: Type
  distSample :: (PrimMonad m) => a -> Params a -> Prob m (Support a)
  distLogP :: a -> Params a -> Support a -> Double

-- -- | Used as a type-lifted kind for conjugate distribution pairs.
-- data Conj a b = Conj a bb

-- | A type-level marker for treating a distribution as a prior.
newtype AsPrior p = AsPrior p

-- -- | A type-level marker for treating a distribution as a likelihood.
-- newtype AsLk l = AsLk l

-- | Marks two distributions as a conjugate pair of prior and likelihood.
-- The property of such a pair is that the posterior has the same form as the prior
-- (including the same 'Params' and 'Support'),
-- and that its parameters can be obtained analytically from the parameters of the prior
-- and a set of observations.
--
-- The class method 'updatePrior' returns the parameters of the posterior
-- given the prior parameters after a single observation.
class (Distribution p, Distribution l, Support p ~ Params l) => Conjugate p l where
  priorSingleton :: p -- ^ provides a singleton instance of the prior distribution
                      -- in order to make sampling from priors easier.
  updatePrior :: l -> Params p -> Support l -> Params p
  predLogP :: l -> Params p -> Support l -> Double

type family Hyper (a :: k) :: Type
type instance Hyper (AsPrior p) = Params p

type family Probs (a :: k) :: Type
type instance Probs (AsPrior p) = Support p

-- type family Value (a :: k) :: Type
-- type instance Value (Conj p l) = Support l

-- helper types for instantiating hyperparameters, parameters, and values
-- ----------------------------------------------------------------------

newtype HyperRep p = HyperRep { runHyper :: Hyper (AsPrior p) }
deriving instance Show (Hyper (AsPrior p)) => Show (HyperRep p)

type instance Hyper (a :: (Type -> Type) -> Type) = a HyperRep

newtype ProbsRep p = ProbsRep { runProbs :: Probs (AsPrior p)}
deriving instance Show (Probs (AsPrior p)) => Show (ProbsRep p)

type instance Probs (a :: (Type -> Type) -> Type) = a ProbsRep

-- newtype ValueRep p l = ValueRep { runValue :: Value (Conj p l) }
-- deriving instance Show (Value (Conj p l)) => Show (ValueRep p l)

-- type instance Value (a :: (Type -> Type -> Type) -> Type) = a ValueRep

-----------------------------------------------------
-- generic magic for conjugates and parameter sets --
-----------------------------------------------------

-- Jeffreys prior
-- ---------------

class Jeffreys a where
  jeffreysPrior :: Hyper a

class GJeffreys t where
  gjeffreysPrior :: forall p. t p

instance GJeffreys V1 where
  gjeffreysPrior = undefined -- ok

instance GJeffreys U1 where
  gjeffreysPrior = U1

-- base case: k is a conjugate distribution
instance (Jeffreys (AsPrior p)) => GJeffreys (K1 i (HyperRep p)) where
  gjeffreysPrior = K1 $ HyperRep $ jeffreysPrior @(AsPrior p)

-- recursive case: k is another record
instance (Jeffreys k, k HyperRep ~ Hyper k) => GJeffreys (K1 i (k HyperRep)) where
  gjeffreysPrior = K1 $ jeffreysPrior @k

instance (GJeffreys t) => GJeffreys (M1 i c (t :: Type -> Type)) where
  gjeffreysPrior = M1 gjeffreysPrior

instance (GJeffreys ta, GJeffreys tb) => GJeffreys (ta :*: tb) where
  gjeffreysPrior = gjeffreysPrior @ta :*: gjeffreysPrior @tb

instance (Generic (t HyperRep), GJeffreys (Rep (t HyperRep))) => Jeffreys (t :: (Type -> Type) -> Type) where
  jeffreysPrior = GHC.Generics.to (gjeffreysPrior @(Rep (t HyperRep)))

-- uniform prior
-- -------------

class Uniform a where
  uniformPrior :: Hyper a

class GUniform t where
  guniformPrior :: forall p. t p

instance GUniform V1 where
  guniformPrior = undefined -- ok

instance GUniform U1 where
  guniformPrior = U1

-- base case: k is a conjugate distribution
instance (Uniform (AsPrior p)) => GUniform (K1 i (HyperRep p)) where
  guniformPrior = K1 $ HyperRep $ uniformPrior @(AsPrior p)

-- recursive case: k is another record
instance (Uniform k, k HyperRep ~ Hyper k) => GUniform (K1 i (k HyperRep)) where
  guniformPrior = K1 $ uniformPrior @k

instance (GUniform t) => GUniform (M1 i c (t :: Type -> Type)) where
  guniformPrior = M1 guniformPrior

instance (GUniform ta, GUniform tb) => GUniform (ta :*: tb) where
  guniformPrior = guniformPrior @ta :*: guniformPrior @tb

instance (Generic (t HyperRep), GUniform (Rep (t HyperRep))) => Uniform (t :: (Type -> Type) -> Type) where
  uniformPrior = GHC.Generics.to (guniformPrior @(Rep (t HyperRep)))

-- sampling from prior
-- ------------------

class Prior a where
  sampleProbs :: (PrimMonad m) => Hyper a -> Prob m (Probs a)

class GPrior i o where
  gsampleProbs :: forall m p. PrimMonad m => i p -> Prob m (o p)

instance GPrior V1 V1 where
  gsampleProbs = undefined -- ok

instance GPrior U1 U1 where
  gsampleProbs _ = pure U1

-- base case: k is a conjugate distribution
instance (Prior (AsPrior p)) => GPrior (K1 i (HyperRep p)) (K1 i (ProbsRep p)) where
  gsampleProbs (K1 (HyperRep hyper)) =
    K1 . ProbsRep <$> sampleProbs @(AsPrior p) hyper

-- recursive case: k is another record
instance (Prior k, k HyperRep ~ Hyper k, k ProbsRep ~ Probs k) =>
         GPrior (K1 i (k HyperRep)) (K1 i (k ProbsRep)) where
  gsampleProbs (K1 hyper) = K1 <$> sampleProbs @k hyper

instance (GPrior ti to) => GPrior (M1 i c ti) (M1 i' c' to) where
  gsampleProbs (M1 x) = M1 <$> gsampleProbs x

instance (GPrior ia oa, GPrior ib ob) => GPrior (ia :*: ib) (oa :*: ob) where
  gsampleProbs (a :*: b) = (:*:) <$> gsampleProbs a <*> gsampleProbs b

instance (GPrior ia oa, GPrior ib ob) => GPrior (ia :+: ib) (oa :+: ob) where
  gsampleProbs (L1 a) = L1 <$> gsampleProbs a
  gsampleProbs (R1 b) = R1 <$> gsampleProbs b

instance ( Generic (a HyperRep)
         , Generic (a ProbsRep)
         , GPrior (Rep (a HyperRep)) (Rep (a ProbsRep))
         ) => Prior (a :: (Type -> Type) -> Type) where
  sampleProbs hyper = GHC.Generics.to <$> gsampleProbs (from hyper)

-----------------------
-- likelihood monads --
-----------------------

type Accessor r p = forall f . Lens' (r f) (f p)

class Monad m => RandomInterpreter m r | m -> r where
  type SampleCtx m a :: Constraint
  sampleValue :: (Conjugate p l, SampleCtx m l) => String -> l -> Accessor r p -> m (Support l)
  sampleConst :: (Distribution d, SampleCtx m d) => String -> d -> Params d -> m (Support d)
  permutationPlate :: (Ord a) => Int -> m a -> m [a]

newtype Trace (r :: (Type -> Type) -> Type)  = Trace { runTrace :: S.Seq Dynamic }
  deriving (Show)

observeValue
  :: (Conjugate p l, Typeable (Support l), Monad m)
  => String
  -> l
  -> Accessor r p
  -> Support l
  -> StateT (Trace r) m ()
observeValue _ _ _ val = modify $ \(Trace st) -> Trace $ st S.|> toDyn val

observeConst
  :: (Distribution d, Typeable (Support d), Monad m)
  => String
  -> d
  -> Params d
  -> Support d
  -> StateT (Trace r) m ()
observeConst _ _ _ val = modify $ \(Trace st) -> Trace $ st S.|> toDyn val

takeTrace :: Typeable a => Trace r -> Maybe (a, Trace r)
takeTrace (Trace t) = do
  (valDyn, rest) <- case S.viewl t of
    S.EmptyL         -> Nothing
    valDyn S.:< rest -> Just (valDyn, rest)
  val <- fromDynamic valDyn
  pure (val, Trace rest)

-- just sample
-- -----------

newtype SampleI m r a = SampleI (ReaderT (r ProbsRep) (Prob m) a)
  deriving (Functor, Applicative, Monad)

instance (PrimMonad m) => RandomInterpreter (SampleI m r) r where
  type SampleCtx (SampleI m r) a = ()
  sampleValue
    :: forall p l
     . (Conjugate p l)
    => String
    -> l
    -> Accessor r p
    -> SampleI m r (Support l)
  sampleValue _ lk getProbs = SampleI $ do
    probs <- ask
    lift $ distSample lk $ runProbs $ view getProbs probs
  sampleConst
    :: forall d
     . (Distribution d)
    => String
    -> d
    -> Params d
    -> SampleI m r (Support d)
  sampleConst _ dist params = SampleI $ lift $ distSample dist params
  permutationPlate = replicateM

sampleResult :: p ProbsRep -> SampleI m p a -> Gen (PrimState m) -> m a
sampleResult probs (SampleI a) = sample (runReaderT a probs)

-- sample and trace the execution process
-- --------------------------------------

newtype TraceI m r a = TraceI (ReaderT (r ProbsRep) (StateT (Trace r) (Prob m)) a)
  deriving (Functor, Applicative, Monad)

instance (PrimMonad m) => RandomInterpreter (TraceI m r) r where
  type SampleCtx (TraceI m r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> TraceI m r (Support l)
  sampleValue _ lk getProbs = TraceI $ do
    probs <- ask
    val   <- lift $ lift $ distSample lk $ runProbs $ view getProbs probs
    modify $ \(Trace obs) -> Trace $ obs S.|> toDyn val
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> TraceI m r (Support d)
  sampleConst _ dist params = TraceI $ do
    val <- lift $ lift $ distSample dist params
    modify $ \(Trace obs) -> Trace $ obs S.|> toDyn val
    pure val
  permutationPlate = replicateM

sampleTrace
  :: r ProbsRep -> TraceI m r a -> Gen (PrimState m) -> m (a, Trace r)
sampleTrace probs (TraceI a) = do
  let st = runReaderT a probs
      pr = runStateT st (Trace mempty)
  sample pr

-- evaluate the probability of a trace
-- -----------------------------------

newtype EvalTraceI r a = EvalTraceI (ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (EvalTraceI r) r where
  type SampleCtx (EvalTraceI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> EvalTraceI r (Support l)
  sampleValue _ lk getProbs = EvalTraceI $ do
    probs              <- ask
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP lk (runProbs $ view getProbs probs) val
    put (trace', totalLogP + logP)
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> EvalTraceI r (Support d)
  sampleConst _ dist params = EvalTraceI $ do
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP dist params val
    put (trace', totalLogP + logP)
    pure val
  permutationPlate n submodel = EvalTraceI $ do
    probs                     <- ask
    (trace  , totalLogP     ) <- get
    (results, (trace', logP)) <- lift $ lift $ runTraceLogP
      probs
      trace
      (replicateM n submodel)
    let unique = MS.fromList results
        permutations =
          logFactorial n - sum (logFactorial . snd <$> MS.toOccurList unique)
    put (trace', totalLogP + logP + permutations)
    pure results

runTraceLogP
  :: r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, (Trace r, Double))
runTraceLogP probs trace (EvalTraceI model) = do
  runStateT (runReaderT model probs) (trace, 0)

evalTraceLogP :: r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, Double)
evalTraceLogP probs trace model = do
  (val, (_trace, logp)) <- runTraceLogP probs trace model
  pure (val, logp)

-- evaluate the predictive probability of a trace
-- -----------------------------------

newtype EvalPredTraceI r a = EvalPredTraceI (ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (EvalPredTraceI r) r where
  type SampleCtx (EvalPredTraceI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> EvalPredTraceI r (Support l)
  sampleValue _ lk getHyper = EvalPredTraceI $ do
    hyper              <- ask
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = predLogP @p lk (runHyper $ view getHyper hyper) val
    put (trace', totalLogP + logP)
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> EvalPredTraceI r (Support d)
  sampleConst _ dist params = EvalPredTraceI $ do
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP dist params val
    put (trace', totalLogP + logP)
    pure val
  permutationPlate n submodel = EvalPredTraceI $ do
    probs                     <- ask
    (trace  , totalLogP     ) <- get
    (results, (trace', logP)) <- lift $ lift $ runTracePredLogP
      probs
      trace
      (replicateM n submodel)
    let unique = MS.fromList results
        permutations =
          logFactorial n - sum (logFactorial . snd <$> MS.toOccurList unique)
    put (trace', totalLogP + logP + permutations)
    pure results

runTracePredLogP
  :: r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, (Trace r, Double))
runTracePredLogP hyper trace (EvalPredTraceI model) = do
  runStateT (runReaderT model hyper) (trace, 0)

evalTracePredLogP
  :: r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, Double)
evalTracePredLogP hyper trace model = do
  (val, (_trace, logp)) <- runTracePredLogP hyper trace model
  pure (val, logp)

-- update priors
-- -------------

newtype UpdatePriorsI r a = UpdatePriorsI (StateT (Trace r, r HyperRep) Maybe a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (UpdatePriorsI r) r where
  type SampleCtx (UpdatePriorsI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> UpdatePriorsI r (Support l)
  sampleValue _ lk accessor = UpdatePriorsI $ do
    (trace, priors) <- get
    (val  , trace') <- lift $ takeTrace trace
    let priors' :: r HyperRep
        priors' = over
          accessor
          (\(HyperRep pr) -> HyperRep $ updatePrior @p @l lk pr val)
          priors
    put (trace', priors')
    pure val
  sampleConst _ _ _ = UpdatePriorsI $ do
    (trace, priors) <- get
    (val  , trace') <- lift $ takeTrace trace
    put (trace', priors)
    pure val
  permutationPlate = replicateM

getPosterior
  :: r HyperRep -> Trace r -> UpdatePriorsI r a -> Maybe (r HyperRep)
getPosterior priors trace (UpdatePriorsI model) = do
  (_trace, posteriors) <- execStateT model (trace, priors)
  pure posteriors

-- show how a trace is generated by a model
-- ----------------------------------------

newtype ShowTraceI r a = ShowTraceI (MaybeT (WriterT String (State (Trace r))) a)
  deriving (Functor, Applicative, Monad)

showTraceItem
  :: forall l r
   . (Show (Support l), Typeable l, Typeable (Support l))
  => String
  -> ShowTraceI r (Support l)
showTraceItem name = ShowTraceI $ do
  trace         <- get
  (val, trace') <- MaybeT $ pure $ takeTrace trace
  put trace'
  let distName = show (typeRep (Proxy :: Proxy l))
  tell
    $  "Sampled value "
    <> show val
    <> " from a "
    <> distName
    <> " at "
    <> name
    <> ".\n"
  pure val

instance RandomInterpreter (ShowTraceI r) r where
  type SampleCtx (ShowTraceI r) l
    = (Typeable (Support l), Typeable l, Show (Support l))
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l), Typeable l, Show (Support l))
    => String
    -> l
    -> Accessor r p
    -> ShowTraceI r (Support l)
  sampleValue name _ acc = showTraceItem @l name
  sampleConst
    :: forall d
     . (Distribution d, SampleCtx (ShowTraceI r) d)
    => String
    -> d
    -> Params d
    -> ShowTraceI r (Support d)
  sampleConst name _ _ = showTraceItem @d name
  permutationPlate = replicateM

showTrace :: Trace r -> ShowTraceI r a -> (Maybe a, String)
showTrace trace (ShowTraceI model) =
  evalState (runWriterT (runMaybeT model)) trace

printTrace :: Trace r -> ShowTraceI r a -> IO ()
printTrace trace model = do
  let (res, txt) = showTrace trace model
  putStrLn txt
  case res of
    Nothing -> do
      putStrLn "Trace does not match the model (stops here)"
    Just _ -> pure ()

-- tracing interpreter (for debugging)
-- -----------------------------------

newtype TraceTraceI r a = TraceTraceI (State (Trace r) a)
  deriving (Functor, Applicative, Monad)

traceTraceItem
  :: forall l r
   . (Show (Support l), Typeable l, Typeable (Support l))
  => String
  -> TraceTraceI r (Support l)
traceTraceItem name = TraceTraceI $ do
  trace <- get
  let loc = show (typeRep (Proxy :: Proxy l)) <> " at " <> name
  case takeTrace trace of
    Just (val, trace') -> do
      put trace'
      DT.traceM $ "Sampled value " <> show val <> " from a " <> loc <> "."
      pure val
    Nothing -> error $ "Incompatible trace at " <> loc <> "."

instance RandomInterpreter (TraceTraceI r) r where
  type SampleCtx (TraceTraceI r) l
    = (Typeable (Support l), Typeable l, Show (Support l))
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l), Typeable l, Show (Support l))
    => String
    -> l
    -> Accessor r p
    -> TraceTraceI r (Support l)
  sampleValue name _ acc = traceTraceItem @l name
  sampleConst
    :: forall d
     . (Distribution d, SampleCtx (TraceTraceI r) d)
    => String
    -> d
    -> Params d
    -> TraceTraceI r (Support d)
  sampleConst name _ _ = traceTraceItem @d name
  permutationPlate = replicateM

traceTrace :: Trace r -> TraceTraceI r a -> a
traceTrace trace (TraceTraceI model) = evalState model trace

-------------------
-- distributions --
-------------------

-- Beta
-- ----

data Beta = Beta
  deriving (Eq, Ord, Show, Generic)

instance Distribution Beta where
  type Params Beta = (Double, Double)
  type Support Beta = Double
  distSample _ = uncurry beta
  distLogP _ (a, b) p =
    log (p ** (a - 1)) + log ((1 - p) ** (b - 1)) - logBeta a b

instance Jeffreys (AsPrior Beta) where
  jeffreysPrior = (0.5, 0.5)

instance Uniform (AsPrior Beta) where
  uniformPrior = (1, 1)

instance Prior (AsPrior Beta) where
  sampleProbs = distSample Beta

-- Bernoulli
-- ---------

data Bernoulli = Bernoulli
  deriving (Eq, Ord, Show, Generic)

instance Distribution Bernoulli where
  type Params Bernoulli = Double
  type Support Bernoulli = Bool
  distSample _ = bernoulli
  distLogP _ p True  = log p
  distLogP _ p False = log (1 - p)

-- Binomial
-- --------

newtype Binomial = Binomial Int
  deriving (Eq, Ord, Show, Generic)

instance Distribution Binomial where
  type Params Binomial = Double
  type Support Binomial = Int
  distSample (Binomial n) = binomial n
  distLogP (Binomial n) p k =
    logChoose n k + (log p * k') + (log (1 - p) * (n' - k'))
   where
    k' = int2Double k
    n' = int2Double n

-- Categorical
-- -----------

type Categorical :: Nat -> Type
data Categorical n = Categorical
  deriving (Eq, Ord, Show, Generic)

instance Distribution (Categorical n) where
  type Params (Categorical n) = V.Vector Double
  type Support (Categorical n) = Int
  distSample _ = categorical
  distLogP _ ps cat = log $ fromMaybe 0 $ ps V.!? cat

-- Dirichlet
-- ---------

type Dirichlet :: Nat -> Type
data Dirichlet n = Dirichlet
  deriving (Eq, Ord, Show, Generic)

instance Distribution (Dirichlet n) where
  type Params (Dirichlet n) = V.Vector Double
  type Support (Dirichlet n) = V.Vector Double
  distSample _ = dirichlet
  distLogP _ counts probs = logp + logz
   where
    logp = sum (V.zipWith (\a x -> log x * (a - 1)) counts probs)
    logz = logGamma (sum counts) - sum (logGamma <$> counts)

instance KnownNat n => Jeffreys (AsPrior (Dirichlet n)) where
  jeffreysPrior = V.replicate (fromIntegral $ natVal (Proxy :: Proxy n)) 0.5

instance KnownNat n => Uniform (AsPrior (Dirichlet n)) where
  uniformPrior = V.replicate (fromIntegral $ natVal (Proxy :: Proxy n)) 1

instance Prior (AsPrior (Dirichlet n)) where
  sampleProbs = distSample Dirichlet

-- Geometric (from 0)
-- ------------------

data Geometric0 = Geometric0
  deriving (Eq, Ord, Show, Generic)

instance Distribution Geometric0 where
  type Params Geometric0 = Double
  type Support Geometric0 = Int
  distSample _ = geometric0
   where
    geometric0 p = do
      coin <- bernoulli p
      if coin then pure 0 else (1 +) <$> geometric0 p
  distLogP _ p val | val >= 0  = (log (1 - p) * int2Double val) + log p
                   | otherwise = log 0

-- Geometric (from 1)
-- ------------------

data Geometric1 = Geometric1
  deriving (Eq, Ord, Show, Generic)

instance Distribution Geometric1 where
  type Params Geometric1 = Double
  type Support Geometric1 = Int
  distSample _ = geometric1
   where
    geometric1 p = do
      coin <- bernoulli p
      if coin then pure 1 else (1 +) <$> geometric1 p
  distLogP _ p val | val >= 1  = (log (1 - p) * int2Double (val - 1)) + log p
                   | otherwise = log 0

-----------------------------
-- conjugate distributions --
-----------------------------

-- beta bernoulli
-- --------------

instance Conjugate Beta Bernoulli where
  priorSingleton = Beta
  updatePrior _ (a, b) False = (a, b + 1)
  updatePrior _ (a, b) True  = (a + 1, b)
  predLogP _ (a, b) False = log a - log (a + b)
  predLogP _ (a, b) True  = log b - log (a + b)

-- B(1+a, b)   G(1+a) G(b) G(a+b)   G(1+a) G(a+b)   a G(a) G(a+b)        a
-- --------- = ------------------ = ------------- = ----------------- = ---
-- B(a, b)     G(a) G(b) G(a+b+1)   G(a) G(a+b+1)   G(a) (a+b) G(a+b)   a+b

-- beta binomial
-- -------------

instance Conjugate Beta Binomial where
  priorSingleton = Beta
  updatePrior (Binomial n) (a, b) x = (a + x', b + (n' - x'))
   where
    x' = int2Double x
    n' = int2Double n
  predLogP (Binomial n) (a, b) k =
    logChoose n k + logBeta (k' + a) (n' - k' + a) - logBeta a b
   where
    n' = int2Double n
    k' = int2Double k

-- beta geometric0
-- ---------------

instance Conjugate Beta Geometric0 where
  priorSingleton = Beta
  updatePrior _ (a, b) k = (a + 1, b + int2Double k)
  predLogP _ (a, b) k =
    (log a + logGamma (a + b) + logGamma (k' + b))
      - (logGamma b + logGamma (a + b + k' + 1))
    where k' = int2Double k

-- beta geometric1
-- ---------------

instance Conjugate Beta Geometric1 where
  priorSingleton = Beta
  updatePrior _ (a, b) k = (a + 1, b + int2Double (k - 1))
  predLogP _ (a, b) k =
    (log a + logGamma (a + b) + logGamma (k' + b))
      - (logGamma b + logGamma (a + b + k' + 1))
    where k' = int2Double k - 1

-- dirichlet categorical
-- ---------------------

instance Conjugate (Dirichlet n) (Categorical n) where
  priorSingleton = Dirichlet
  updatePrior _ counts obs
    | obs >= 0, obs < V.length counts
    = counts V.// [(obs, (counts V.! obs) + 1)]
    | otherwise
    = counts
  predLogP _ counts obs =
    log (fromMaybe 0 $ counts V.!? obs) - log (V.sum counts)

max_stars_repo_name: DCMLab/conjugate-haskell-programs

max_issues_repo_name: DCMLab/conjugate-haskell-programs

max_forks_repo_name: DCMLab/conjugate-haskell-programs


text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_HADDOCK show-extensions #-}

-- |
-- Module      : Swarm
-- Description : Particle Swarm Optimisation
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental
module PSO.Swarm
  ( -- * Idea
    --
    -- | This module implements different variants of Particle Swarm
    -- Optimisation (PSO), so if you're completely unfamiliar with PSO, go read
    -- https://en.wikipedia.org/wiki/Particle_swarm_optimization first, it's a
    -- nice introduction.
    --
    -- Let there be \(M\in\mathbb{N}\) 'Bee's in the swarm. Bees live in a phase
    -- space (which is a vector space). What this means is that each bee is
    -- described by a phase \(\psi\). The two most common choices of the phase
    -- are:
    --
    -- * \((q, p)\in\mathcal{H}^2\) with \(q\) being the canonical coordinate
    -- and \(p\) being the canonical momentum.
    -- * \(|\psi\rangle\in\mathcal{H}^2\), i.e. a vector in Hilbert space.
    --
    -- While flying around, bees may pick up some information about the
    -- surrounding environment. We call the part of it that bees remember a
    -- local (or private) /guide/. It's a guide because it helps bees decide
    -- where to move next. The simplest and most common guide is 'BeeGuide'.
    --
    -- The whole point of having a swarm of bees rather than one bee is
    -- collective behavior. In case of PSO this is achieved by having a /global
    -- guide/. Each bee then uses this global guide, its own local guide, and
    -- its current phase to compute the next phase. This is represented by
    -- 'PhaseUpdater'.
    --
    -- The whole optimisation process then looks more or less like this:
    --
    --    1. __Initialisation__ (see 'mkSwarm' function).
    --
    --        * Generate initial phases \((\psi^{(0)}_i)_{i\in\{1\dots M\}}\).
    --        * Evaluate fitness in these phases, i.e. compute
    --          \((f(\psi^{(0)}_i))_{i\in\{1\dots M\}}\) (see 'mkBee' function).
    --        * Initialise local guides \((g^{(0)}_{l,i})_{i\in\{1\dots M\}}\)
    --          (see 'mkLocalG' function).
    --        * Construct bees using phases, fitness values, and guides (see
    --          'mkBee' function).
    --        * Initialise global guide \(g^{(0)}_g\) (see 'mkGlobalG'
    --          function).
    --        * Construct swarm using bees and global guide.
    --
    --    2. __Main Loop__. While the termination condition is not met, do the
    --    following (see 'updateSwarm' function):
    --
    --        * Update each bee (see 'updateBee' function), i.e.
    --
    --            * Using global guide \(g^{(n)}_g\), local guide
    --              \(g^{(n)}_{l,i}\), current phase \(\psi^{(n)}_i\), and
    --              \(f(\psi^{(n)}_i)\), come up with a new phase
    --              \(\psi^{(n+1)}_i\) (see 'PhaseUpdater').
    --            * Evaluate the fitness in this new phase, i.e. compute
    --              \(f(\psi^{(n+1)}_i)\).
    --            * Using \(\psi^{(n+1)}_i\) and \(f(\psi^{(n+1)}_i)\), compute
    --              the local guide in the next iteration \(g^{(n+1)}_i\) (see
    --              'updateLocalG' function).
    --
    --        * Update the global guide (see 'updateGlobalG' function).
    --

    -- * Basic types
    Bee(..)
  , CMState(..)
  , QMState(..)
  , LocalGuide(..)
  , BeeGuide(..)
  , GlobalGuide(..)
  , SwarmGuide(..)
  , Swarm(..)
  , PhaseUpdater(..)

    -- * Important functions
  , wpgUpdater
  -- , kickUpdater
  , deltaUpdater
  , absorbUpdater
  , mkCMState
  , mkQMState
  , optimiseND
  , optimiseNDFromList

    -- * Details
  , VelocityUpdater(..)
  , upScale
  , upLocal
  , upGlobal
  , upStep
  , updateBee
  , updateSwarm
  , mkBee
  , mkSwarm
  , fromPure
  , iterateNM
  , iterateWhileM
  , Absorbable(..)
  , DeltaWell(..)
  , Scalable(..)

    -- * Lenses
  , HasPos(..)
  , HasVal(..)
  , HasVar(..)
  , HasGuide(..)
  , HasState(..)
  , HasRun(..)
  , HasBees(..)
  , HasUpdater(..)
  , HasIteration(..)
  ) where

import Data.Complex
import Data.Ord(comparing)
import qualified Data.List as List
-- import Data.Monoid
import Data.Semigroup
import qualified Data.Vector.Generic as GV
import qualified Data.Vector.Storable as V

import Control.Arrow
-- import Control.Lens
import Lens.Micro
import Lens.Micro.TH
import Lens.Micro.Extras
import Control.DeepSeq
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.Reader

import qualified Foreign
import Foreign.Storable.Tuple

import GHC.Float (float2Double)
import GHC.Generics(Generic)

-- import qualified Numeric.LinearAlgebra as LA
-- import Numeric.LinearAlgebra.Devel(matrixFromVector, orderOf, MatrixOrder(..))
-- import Numeric.LinearAlgebra.Data(flatten, tr)
-- import qualified Numeric.LinearAlgebra.Devel as LADevel

import PSO.VectorSpace
import PSO.Random

-- | A bee in our swarm.
--
-- * @α@ is the type of a point in our phase space (i.e. position-velocity pair
-- in the classical case and wave function in the quantum mechanical)
-- * @β@ is the type of the local guide, i.e. \"memory\" of the bee.
-- * @r@ return type of the function we're minimising.
--
-- /Note:/ use 'state', 'guide', and 'val' lenses to access the fields.
data Bee α β r = Bee
  { _beeState :: !α
  , _beeGuide :: !β
  , _beeVal   :: !r
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''Bee

-- | Point in the classical phase space.
--
-- @χ@ -- vector in coordinate space. We always use time steps of 1, so physical
-- dimension of velocity is equal to the one of coordinate. The type of velocity
-- is thus also @χ@.
--
-- /Note:/ use 'pos' and 'vel' lenses to access the fields.
data CMState χ = CMState
  { _cmstatePos :: !χ
  , _cmstateVel :: !χ
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''CMState

-- | Point in our quantum mechanical phase space (i.e. Hilbert space).
--
-- /Note:/ use 'pos' lens to access the field.
data QMState χ = QMState
  { _qmstatePos :: !χ
  } deriving (Generic, Show, NFData)

makeLensesWith abbreviatedFields ''QMState

-- | Rather than limiting ourselves to a single type of local guide (because who
-- knows which info we might want to keep track of), we define a type class.
-- @LocalGuide β α r@ means that @β@ is a local guide for a bee with phase of
-- type @α@ and fitness function returning @r@.
class LocalGuide β α r where
  -- | Initialisation of a guide. Given a phase \(\psi^{(0)}\) and its fitness
  -- \(f(\psi^{(0)})\), @mkLocalG@ creates local guide \(g^{(0)}_l\).
  mkLocalG     :: α -> r -> β
  -- | Update step. After computing the phase \(\psi^{(n+1)}\) (of type @α@) and
  -- its fitness \(f(\psi^{(n+1)})\) (of type @r@) in the @(n+1)@'st iteration,
  -- we use @updateLocalG@ to compute \(g^{(n+1)}_l\) given \(g^{(n)}_l\).
  updateLocalG :: α -> r -> β -> β

-- | Similar to 'LocalGuide' we define a class of global guides rather than a
-- single one. Here, @GlobalGuide γ β@ means that @γ@ can serve as a global
-- guide for local guides of type @β@.
class (LocalGuide β α r) => GlobalGuide γ β α r where
  -- | Initialisation of a guide. Given a collection of bees, @mkLocalG@
  -- initialises global guide \(g^{(0)}_g\).
  mkGlobalG     :: (Foldable t, Functor t) => t (Bee α β r) -> γ
  -- | Update step. After computing phases, values, and local guides of bees in
  -- the @(n+1)@'st iteration, we use @updateLocalG@ to compute \(g^{(n+1)}_g\)
  -- given \(g^{(n)}_g\).
  updateGlobalG :: (Foldable t, Functor t) => t (Bee α β r) -> γ -> γ

-- | A guide that remembers position and its corresponding fitness.
--
-- /Note:/ use 'pos' and 'val' lenses to access the fields.
data BeeGuide χ r = BeeGuide
  { _beeguidePos :: !χ -- ^ Position.
  , _beeguideVal :: !r -- ^ Fitness value.
  } deriving (Generic, Show)

makeLensesWith abbreviatedFields ''BeeGuide

instance (NFData χ, NFData r) => NFData (BeeGuide χ r)

-- | We compare guides by their fitness values.
instance Eq r => Eq (BeeGuide χ r) where
  x == y = (x ^. val) == (y ^. val)

-- | We compare guides by their fitness values.
instance Ord r => Ord (BeeGuide χ r) where
  compare = comparing (view val)

-- | 'BeeGuide' is the simplest and, probably, the most useful example of a
-- 'LocalGuide'. It keeps track of the best position and its fitness. That is,
-- if we've run @n@ iterations (calls to 'updateLocalG'), then the best position
-- is given by \(\operatorname{argmin}_{k\in\{0\dots n\}}f(\psi^{(k)}_i)\), and
-- its fitness is just \(\operatorname{min}_{k\in\{0\dots n\}}f(\psi^{(k)}_i)\).
instance (HasPos α χ, Ord r) => LocalGuide (BeeGuide χ r) α r where
  mkLocalG ψ fψ = BeeGuide (ψ ^. pos) fψ
  updateLocalG ψ fψ g
    | fψ < g ^. val = g & (pos .~ (ψ ^. pos))
                        & (val .~ fψ)
    | otherwise = g

-- | 'BeeGuide' can also serve as a 'GlobalGuide'. It keeps track of the overall
-- best position among all bees and iterations and its fitness. So after @n@
-- iterations, the best position is given by
-- \(\operatorname{argmin}_{k\in\{0\dots n\},i\in\{1\dots M\}}f(\psi^{(k)}_i)\),
-- and its fitness is
-- \(\operatorname{min}_{k\in\{0\dots n\},i\in\{1\dots M\}}f(\psi^{(k)}_i)\).
instance (LocalGuide (BeeGuide χ r) α r, Ord (BeeGuide χ r))
  => GlobalGuide (BeeGuide χ r) (BeeGuide χ r) α r where
    mkGlobalG = minimum . fmap (view guide)
    updateGlobalG xs _ = minimum . fmap (view guide) $ xs

-- | A guide that remembers
--
--    * Best position throughout the whole swarm.
--    * Fitness value at this position.
--    * Current iteration.
--    * Variance in the fitness values of the local guides.
--
-- /Note:/ use 'pos', 'val', 'iteration', and 'var' lenses to access the fields.
data SwarmGuide χ r = SwarmGuide
  { _swarmguidePos       :: !χ
  , _swarmguideVal       :: !r
  , _swarmguideIteration :: !Int
  , _swarmguideVar       :: () -- !r
  } deriving (Generic, Show)

makeLensesWith abbreviatedFields ''SwarmGuide

instance (NFData χ, NFData r) => NFData (SwarmGuide χ r)

-- | Calculates the mean value of a sequence.
mean :: (Fractional a, Foldable v) => v a -> a
mean xs = sum xs / fromIntegral (length xs)

-- | Calculates the variance in a sequence.
variance :: (Fractional a, Foldable v, Functor v) => v a -> a
variance xs = let m = mean xs in mean . fmap (\x -> (x - m)^^2) $ xs

-- | 'SwarmGuide' is another example of a 'GlobalGuide'. On top of keeping track
-- of the best position and its fitness value (what 'BeeGuide' does), we
-- remember the current iteration ('mkGlobalG' sets it to @0@, and each call to
-- 'updateGlobalG' increases the iteration by @1@). And finally, each iteration
-- we compute the variance \(\operatorname{Var}\{f(\psi^{(n)}_i)\}_i\) in the
-- fitness values of bees.
instance (LocalGuide (BeeGuide χ r) α r, Ord (BeeGuide χ r))
  => GlobalGuide (SwarmGuide χ r) (BeeGuide χ r) α r where
    mkGlobalG xs = SwarmGuide (x ^. pos) (x ^. val) 0 var
      where x   = minimum . fmap (view guide) $ xs
            var = () -- variance . fmap (view val) $ xs
    updateGlobalG xs g = mkGlobalG xs & iteration .~ (g ^. iteration + 1)

-- | Given global guide and bee in the @n@'th iteration, calculates the phase of
-- the bee in the @(n+1)@'st iteration. This calculation is done in some 'Monad'
-- @m@ which allows us to, for example, use random numbers when updating the
-- phase.
--
-- /Note:/ use 'run' lens to unpack the function.
newtype PhaseUpdater m γ β α r = PhaseUpdater
  { _phaseupdaterRun :: γ -> Bee α β r -> m α
  }

makeLensesWith abbreviatedFields ''PhaseUpdater

-- | 'PhaseUpdater's can be combined to build more difficult ones.
instance (Monad m) => Semigroup (PhaseUpdater m γ β α r) where
  (PhaseUpdater a) <> (PhaseUpdater b) = PhaseUpdater update
    where update guide bee = a guide bee >>= \ψ ->
            b guide (bee & state .~ ψ)

-- | Our swarm.
data Swarm m γ β α r = Swarm
  { _swarmBees      :: ![Bee α β r] -- ^ Collection of bees.
  , _swarmGuide     :: !γ -- ^ Global guide.
  , _swarmFunc      :: !(α -> m r) -- ^ Fitness function.
  , _swarmUpdater   :: !(PhaseUpdater m γ β α r) -- ^ Updater.
  }

makeLensesWith abbreviatedFields ''Swarm

-- | Updates the canonical momentum in the classical phase. Given a global guide
-- and a bee described by classical mechanics, returns the momentum of the bee
-- in the next iteration. 'VelocityUpdater's can be combined using '<>' operator
-- and afterwards converted to 'PhaseUpdater's using 'upStep' function.
newtype VelocityUpdater m γ χ r = VelocityUpdater
  { _velocityupdaterRun :: γ -> Bee (CMState χ) (BeeGuide χ r) r -> m χ
  }

makeLensesWith abbreviatedFields ''VelocityUpdater

-- | 'VelocityUpdater's can be combined to form more complex ones.
instance (Monad m) => Semigroup (VelocityUpdater m γ χ r) where
  (VelocityUpdater a) <> (VelocityUpdater b) = VelocityUpdater $ update
    where update guide bee =
            a guide bee >>= \v -> b guide (bee & (state . vel) .~ v)

-- | Given a constant \(\varphi_l\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = p^{(n)} + \varphi_l U(0, 1) \circ (g^{(n)}_l - q^{(n)}) \;,
-- \]
-- where \(\circ\) denotes the Hadamard product, \(U(0, 1)\) is a vector
-- whose elements are uniformly distributed in the \([0, 1)\) interval, and
-- \(g^{(n)}_l\) is local best position in the @n@'th iteration.
upLocal :: (RandomScalable m χ, VectorSpace χ λ) => λ -> VelocityUpdater m γ χ r
upLocal φl = VelocityUpdater $ upLocalImpl φl

upLocalImpl ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos α χ, HasVel α χ, HasPos β χ)
  => λ -> γ -> Bee α β r -> m χ
upLocalImpl φ _ bee =
  let δq = bee ^. guide . pos - bee ^. state . pos
      p  = bee ^. state . vel
   in liftM (scale φ) (randScale δq) >>= \δp -> return (p + δp)

-- | Given a constant \(\varphi_g\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = p^{(n)} + \varphi_g U(0, 1) \circ (g^{(n)}_g - q^{(n)}) \;,
-- \]
-- where \(\circ\) denotes the Hadamard product, \(U(0, 1)\) is a vector
-- whose elements are uniformly distributed in the \([0, 1)\) interval, and
-- \(g^{(n)}_g\) is global best position in the @n@'th iteration.
upGlobal ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => λ -> VelocityUpdater m γ χ r
upGlobal φg = VelocityUpdater $ upGlobalImpl φg

upGlobalImpl ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos α χ, HasVel α χ, HasPos γ χ)
  => λ -> γ -> Bee α β r -> m χ
upGlobalImpl φ g bee = do
  let δq = g ^. pos - bee ^. state . pos
      p  = bee ^. state . vel
   in liftM (scale φ) (randScale δq) >>= \δp -> return (p + δp)

-- | Given a constant \(\omega\), creates a 'VelocityUpdater' that updates
-- momentum according to
-- \[
--    p^{(n+1)} = \omega p^{(n)} \;.
-- \]
upScale :: (Monad m, Scalable λ χ) => λ -> VelocityUpdater m γ χ r
upScale ω = VelocityUpdater $ upScaleImpl ω

upScaleImpl :: (Monad m, Scalable λ χ, HasVel α χ) => λ -> γ -> Bee α β r -> m χ
upScaleImpl ω _ = return . scale ω . view (state . vel)

-- | Converts a 'VelocityUpdater' into a 'PhaseUpdater'. This is done by first
-- running the velocity updater to compute \(p^{(n+1)}\). \(q^{(n+1)}\) is then
-- calculated as \(q^{(n)} + p^{(n+1)}\).
upStep ::
     (Monad m, Num χ)
  => VelocityUpdater m γ χ r -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
upStep (VelocityUpdater upVel) = PhaseUpdater update
  where update guide bee = do
          δq <- upVel guide bee
          let x = bee ^. state
              q = x ^. pos
          return $ x & (vel .~ δq)
                     & (pos .~ q + δq)

-- | Creates the standard \"WPG\" updater. Given constants
-- \(\omega,\varphi_l,\varphi_g\), creates a 'PhaseUpdater' that updates the
-- phase according to
-- \[
--  \left\{
--  \begin{aligned}
--    p^{(n+1)} &= \omega p^{(n)}
--              + \varphi_l U_l(0, 1) \circ (g^{(n)}_l - q^{(n)})
--              + \varphi_g U_g(0, 1) \circ (g^{(n)}_g - q^{(n)}) \;, \\
--    q^{(n+1)} &= q^{(n)} + p^{(n+1)} \;,
--  \end{aligned}
--  \right.
-- \]
-- where \(\circ\) denotes the Hadamard product and \(U_l(0, 1)\) and
-- \(U_g(0,1)\) are vectors whose elements are uniformly distributed in the
-- \([0, 1)\) interval.
wpgUpdater ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => (λ, λ, λ) -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
wpgUpdater (ω, φl, φg) =
  upStep $ (upScale ω) <> (upLocal φl) <> (upGlobal φg)

-- | Creates the standard updater which uses constriction rather than inertia
-- parameter. Given constants
-- \(\kappa,\varphi_l,\varphi_g\), creates a 'PhaseUpdater' that updates the
-- phase according to
-- \[
--  \left\{
--  \begin{aligned}
--    p^{(n+1)} &= \kappa \left( p^{(n)}
--              + \varphi_l U_l(0, 1) \circ (g^{(n)}_l - q^{(n)})
--              + \varphi_g U_g(0, 1) \circ (g^{(n)}_g - q^{(n)}) \right) \;, \\
--    q^{(n+1)} &= q^{(n)} + p^{(n+1)} \;,
--  \end{aligned}
--  \right.
-- \]
-- where \(\circ\) denotes the Hadamard product and \(U_l(0, 1)\) and
-- \(U_g(0,1)\) are vectors whose elements are uniformly distributed in the
-- \([0, 1)\) interval.
constrictedUpdater ::
     (RandomScalable m χ, VectorSpace χ λ, HasPos γ χ)
  => (λ, λ, λ) -> PhaseUpdater m γ (BeeGuide χ r) (CMState χ) r
constrictedUpdater (κ, φl, φg) =
  upStep $ (upLocal φl) <> (upGlobal φg) <> (upScale κ)

-- randomWalkUpdater :: (UniformDist p gen m)
--                   => (p, p) -> Updater m gen s g p r
-- randomWalkUpdater bounds = Updater $ f
--   where f _ _ = do
--           randGen <- ask
--           v' <- lift $ uniform bounds randGen
--           return v'
--
-- kickUpdater ::
--      ( RandomScalable p gen m
--      , UniformDist p gen m
--      , VectorSpace p a
--      , Scalable r p
--      , HasVar s r
--      , HasPos s p
--      , HasPos g p
--      , Floating r
--      , Ord r
--      , s ~ SwarmGuide p r
--      )
--   => (r, r) -> (p, p) -> (a, a, a) -> Updater m gen s g p r
-- kickUpdater cutoff bounds wpg = Updater $ kickUpdaterImpl cutoff bounds wpg
--
-- kickUpdaterImpl ::
--      ( RandomScalable p gen m
--      , UniformDist p gen m
--      , VectorSpace p a
--      , Scalable r p
--      , HasVar s r
--      , HasPos s p
--      , HasPos g p
--      , Floating r
--      , Ord r
--      )
--   => (r, r) -> (p, p) -> (a, a, a) -> s -> Bee p r g -> ReaderT gen m p
-- kickUpdaterImpl (cutoff, c) bounds wpg stats x = runUpdater updater stats x
--   where updater = if (stats ^. var) < cutoff
--                      then randomWalkUpdater bounds
--                           <> upScale (c * sqrt (stats ^.var))
--                      else standardUpdater wpg

-- | Update the position according to QDPSO (Quantum Delta well Particle Swarm
-- Optimisation).
class (Monad m) => DeltaWell m λ χ where
  upDeltaWell :: λ -> χ -> χ -> m χ

-- | 1D QDPSO updater. Given a constant \(\kappa\in\mathbb{R}\), center
-- \(g^{(n)}\in\mathbb{R}\) of the density distribution, and bees current
-- position \(q^{(n)}\in\mathbb{R}\), its position in the next iteration is
-- computed as
-- \[
--  q^{(n+1)} =
--    q^{(n)} + \sgn(2\cdot\operatorname{rand}(0,1) - 1) \kappa
--                \log(\frac{1}{\operatorname{rand}(0, 1)}) \cdot
--                  |q^{(n)} - g^{(n)}| \;.
-- \]
instance (Monad m, RealFloat λ, Ord λ, Randomisable m λ)
  => DeltaWell m λ λ where
    upDeltaWell κ p x = do
      u    <- random
      sign <- (\c -> signum (2 * c - 1)) <$> random
      return $ p + (sign * (log . recip $ u) / κ) * abs (x - p)

-- | Extending QDPSO to \(\mathbb{C}\). 1D real version is simply applied to
-- both real and imaginary parts.
instance (RealFloat λ, DeltaWell m λ λ)
  => DeltaWell m λ (Complex λ) where
    upDeltaWell κ p x =
      let f getter = uncurry (upDeltaWell κ) . over each getter
       in return (:+) `ap` f realPart (p, x)
                      `ap` f imagPart (p, x)

-- | Extending QDPSO to multidimensional case. 1D version is simply applied to
-- each component.
instance (Monad m, Foreign.Storable ξ, DeltaWell m λ ξ)
  => DeltaWell m λ (V.Vector ξ) where
    upDeltaWell κ p x = V.zipWithM (upDeltaWell κ) p x

{-
liftMatrix2M ::
     (Monad m, LA.Element α, LA.Element β, LA.Element γ,
      LA.Container V.Vector α, Num α,
      LA.Container V.Vector β, Num β,
      LA.Transposable (LA.Matrix α) (LA.Matrix α),
      LA.Transposable (LA.Matrix β) (LA.Matrix β))
  => (LA.Vector α -> LA.Vector β -> m (LA.Vector γ))
  -> LA.Matrix α -> LA.Matrix β -> m (LA.Matrix γ)
liftMatrix2M f m1@(LA.size -> (r,c)) m2
  | LA.size m1 /= LA.size m2 = error "nonconformant matrices in liftMatrix2M"
  | orderOf m1 == RowMajor =
      return (matrixFromVector RowMajor r c) `ap` f (flatten m1) (flatten m2)
  | otherwise =
      return (matrixFromVector ColumnMajor r c) `ap`
        f (flatten . tr $ m1) (flatten . tr $ m2)

instance (Monad m, Num ξ, LA.Container V.Vector ξ, LA.Transposable (LA.Matrix ξ)(LA.Matrix ξ), DeltaWell m λ ξ)
  => DeltaWell m λ (LA.Matrix ξ) where
    upDeltaWell κ = liftMatrix2M (V.zipWithM (upDeltaWell κ))
-}

-- | Creates a QDPSO updater.
deltaUpdater :: forall m g γ β χ λ r.
     ( Randomisable m λ
     , RealFloat λ
     , Ord λ
     , VectorSpace χ λ
     , DeltaWell m λ χ
     , HasPos β χ
     , HasPos γ χ
     )
  => λ -> γ -> Bee (QMState χ) β r -> m (QMState χ)
deltaUpdater κ g bee = do
  φ₁ <- random :: m λ
  φ₂ <- random :: m λ
  -- lift $ putStrLn $ "φ₁ = " ++ show φ₁ ++ ", φ₂ = " ++ show φ₂
  let center = (φ₁ / (φ₁ + φ₂)) `scale` (g ^. pos)
             + (φ₂ / (φ₁ + φ₂)) `scale` (bee ^. guide . pos)
  x <- upDeltaWell κ center (bee ^. state . pos)
  -- lift $ putStrLn $ "p - x = " ++ show (bee ^. state . pos - p)
  return $ (bee ^. state) & (pos .~ x)

data RelativePosition = Inside | LeftOf | RightOf

class RelativelyPositioned α β where
  getRelPos :: α -> β -> RelativePosition

instance (Ord λ) => RelativelyPositioned (λ, λ) λ where
  getRelPos (l, h) x
    | x < l     = LeftOf
    | x > h     = RightOf
    | otherwise = Inside

class Absorbable λ α where
  absorb :: (λ, λ) -> α -> α

class Reflectable λ α where
  reflect :: (λ, λ) -> α -> α

instance (Ord χ) => Absorbable χ χ where
  absorb (l, h) x
    | x < l     = l
    | x > h     = h
    | otherwise = x

instance (Ord λ) => Absorbable λ (Complex λ) where
  absorb bs = over each (absorb bs)

instance (Ord λ, Num λ)
  => Absorbable λ (λ, λ) where
    absorb (l, h) (q, p) = case getRelPos (l, h) q of
                             Inside  -> (q, p)
                             LeftOf  -> (l, 0)
                             RightOf -> (h, 0)

zipWithTuple :: (a -> b -> c) -> (a, a) -> (b, b) -> (c, c)
zipWithTuple f (x₁, x₂) (y₁, y₂) = (f x₁ y₁, f x₂ y₂)

instance (Ord λ, Num λ)
  => Absorbable λ (Complex λ, Complex λ) where
    absorb bounds ψ =
      let f getter = absorb bounds . over each getter
       in zipWithTuple (:+) (f realPart ψ) (f imagPart ψ)

instance (Absorbable λ (ξ, ξ), GV.Vector χ ξ, GV.Vector χ (ξ, ξ))
  => Absorbable λ (CMState (χ ξ)) where
    absorb bounds ψ =
      let (q, p) = GV.unzip . GV.map (absorb bounds)
                 $ GV.zip (ψ ^. pos) (ψ ^. vel)
       in ψ & (pos .~ q) & (vel .~ p)

instance (Absorbable λ ξ, GV.Vector χ ξ)
  => Absorbable λ (QMState (χ ξ)) where
    absorb bounds = pos %~ GV.map (absorb bounds)

absorbUpdater :: (Monad m, Absorbable λ α) => (λ, λ) -> PhaseUpdater m γ β α r
absorbUpdater bounds = PhaseUpdater $ upAbsorbImpl bounds

upAbsorbImpl :: (Monad m, Absorbable λ α) => (λ, λ) -> γ -> Bee α β r -> m α
upAbsorbImpl bounds _ = return . absorb bounds . view state

-- | Given a swarm and bee in the @n@'th iteration, returns a computation that
-- returns the bee in the @(n+1)@'st iteration.
updateBee ::
     (Monad m, LocalGuide β α r)
  => Swarm m γ β α r -> Bee α β r -> m (Bee α β r)
updateBee xs x = do
  ψ'  <- (xs ^. updater . run) (xs ^. guide) x
  fψ' <- (xs ^. func) ψ'
  return $ x & (state .~ ψ')
             & (val   .~ fψ')
             & (guide %~ updateLocalG ψ' fψ')

-- | Given a swarm in the @n@'th iteration, returns a computation that returns
-- the swarm in the @(n+1)@'st iteration.
updateSwarm :: (Monad m, GlobalGuide γ β α r)
            => Swarm m γ β α r -> m (Swarm m γ β α r)
updateSwarm swarm = do
  xs <- mapM (updateBee swarm) (swarm ^. bees)
  return $ swarm & (bees .~ xs)
                 & (guide %~ updateGlobalG xs)

-- | Applies \(f\) to \(x\) \(\max(0, n - 1)\) times and accumulates the
-- results. Returned list will have \(n\) elements.
iterateNM ::
     (Monad m)
  => Int -- ^ \(n\), number of times to apply \(f\)
  -> (a -> m a) -- ^ function \(f\) to apply
  -> a -- ^ initial value \(x\)
  -> m [a]
iterateNM n f x
  | n == 0 = return [x]
  | otherwise = do
    y <- f x
    ys <- iterateNM (n - 1) f y
    return $ x : ys

iterateWhileM ::
     (Monad m)
  => (a -> Bool)
  -> (a -> m a) -- ^ function \(f\) to apply
  -> a -- ^ initial value \(x\)
  -> m [a]
iterateWhileM predicate func x
  | predicate x = do
      y  <- func x
      ys <- iterateWhileM predicate func y
      return $ x : ys
  | otherwise   = return [x]

mkCMState :: (UniformDist m χ, Num χ) => (χ, χ) -> m (CMState χ)
mkCMState bounds = liftM (\q -> CMState q 0) $ uniform bounds

mkQMState :: (UniformDist m χ, Num χ) => (χ, χ) -> m (QMState χ)
mkQMState bounds = liftM (\q -> QMState q)   $ uniform bounds

-- | @mkSwarm updater initialiser inside func n@ creates a new 'Swarm' where
--
-- * @updater@ is the update strategy used to obtain velocities of bees in the
-- next iteration.
-- * @initialiser@ creates a new random 'Bee'. It is assumed that position of a
-- bee created using @initialiser@ will satisfy @inside@.
-- * Given a position \(x\in\mathcal{H}\) returns whether \(x\in\mathcal{V}\)
-- holds.
-- * @func@ is the function \(f\) we're trying to minimise.
-- * @n@ is number of bees in the newly created swarm.
mkSwarm ::
     (Monad m, GlobalGuide γ β α r)
  => m α
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> Int
  -> m (Swarm m γ β α r)
mkSwarm newState updater func n = (sequence . replicate n $ newState) >>=
  mapM (mkBee func) >>= \xs ->
    return $ Swarm xs (mkGlobalG xs) func updater

mkSwarmFromList ::
     (Monad m, GlobalGuide γ β α r)
  => [α]
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> m (Swarm m γ β α r)
mkSwarmFromList states updater func =
  mapM (mkBee func) states >>= \xs ->
    return $ Swarm xs (mkGlobalG xs) func updater

fromPure :: (Monad m, HasPos α χ) => (χ -> r) -> α -> m r
fromPure f ψ = return $ f (ψ ^. pos)
--
-- class Between a b where
--   isBetween :: (a, a) -> b -> Bool
--
-- instance (Ord a) => Between a a where
--   isBetween (l, h) x = l <= x && x <= h
--
-- instance {-# OVERLAPS #-} (Between a a) => Between (Complex a) (Complex a) where
--   isBetween (rl :+ il, rh :+ ih) (rx :+ ix) =
--     isBetween (rl, rh) rx && isBetween (il, ih) ix
--
-- instance {-# OVERLAPS #-} (Foreign.Storable a, Between a a)
--   => Between (V.Vector a) (V.Vector a) where
--     isBetween (low, high) x =
--       V.and $ V.zipWith3 (\l h x -> isBetween (l, h) x) low high x
--
-- -- | Creates a new 'Bee' with a random position inside a given 1D interval and
-- -- zero velocity.
-- randomBee1D ::
--      (Num p, UniformDist p gen m)
--   => (p, p) -- ^ @[low, high]@ interval. It is assumed that @low <= high@.
--   -> (p -> ReaderT gen m r) -- ^ Function we're minimising.
--   -> ReaderT gen m (Bee p r (BeeGuide p r))
-- randomBee1D bounds func = do
--   g <- ask
--   x <- lift $ uniform bounds g
--   fx <- func x
--   return $ Bee x 0 fx (BeeGuide x fx)

-- | Given a fitness function \(f\) and phase \(\psi^{(0)}\), creates a new
-- 'Bee'.
mkBee :: (Monad m, LocalGuide β α r)
  => (α -> m r) -> α -> m (Bee α β r)
mkBee f ψ = do
  fψ <- f ψ
  return $ Bee ψ (mkLocalG ψ fψ) fψ

-- -- | Simplified @mkSwarm@ for 1D pure functions.
-- -- simpleInit1D ::
-- --      ( Num a
-- --      , RandomScalable a gen m
-- --      , UniformDist a gen m
-- --      , Between a a
-- --      , Ord r
-- --      , s ~ g
-- --      , g ~ BeeGuide a r
-- --      )
-- --   => (a, a, a) -- ^ WPG parameters \((\omega,\varphi_l,\varphi_g)\).
-- --   -> (a, a) -- ^ Boundaries \([a, b]\).
-- --   -> (a -> r) -- ^ Function \(f\) we're minimising.
-- --   -> Int -- ^ Number of bees.
-- --   -> ReaderT gen m (Swarm m gen s g a r)
-- -- simpleInit1D wpg (low, high) func =
-- --   let updater = standardUpdater wpg
-- --       initialiser = randomBee1D (low, high)
-- --       boundaries = isBetween (low, high)
-- --   in mkSwarm updater initialiser boundaries (fromPure func)
--
--
-- unpackToV :: (Foreign.Storable a) => [(a, a)] -> (V.Vector a, V.Vector a)
-- unpackToV xs = let a = V.fromList . map fst $ xs
--                    b = V.fromList . map snd $ xs
--                 in (a, b)
--
-- -- | Simplified @mkSwarm@ for N-dimensional functions.
-- simpleInitND ::
--      ( Foreign.Storable a
--      , VectorSpace p a
--      , RandomScalable p gen m
--      , UniformDist p gen m
--      , Between p p
--      , Ord r
--      , p ~ V.Vector a
--      , GlobalGuide s g p r
--      )
--   => Updater m gen s g p r -- ^ PWG parameters \((\omega,\varphi_l,\varphi_g)\)
--   -> [(a, a)]
--   -> [(a, a)] -- ^ List of intervals \([a_i,b_i]\). \(\mathcal{V}\) is then
--   -- \(\prod_i [a_i,b_i]\).
--   -> (p -> r) -- ^ Function \(f\) we're minimising.
--   -> Int -- ^ Number of bees.
--   -> ReaderT gen m (Swarm m gen s g p r)
-- simpleInitND updater initBounds bounds func =
--   let initialiser = randomBeeND (unpackToV initBounds)
--       boundaries = isBetween (unpackToV bounds)
--   in mkSwarm updater initialiser boundaries (fromPure func)
--
-- -- optimise1D ::
-- --      ( Num a
-- --      , RandomScalable a gen m
-- --      , UniformDist a gen m
-- --      , Between a a
-- --      , Ord r
-- --      , s ~ g
-- --      , g ~ BeeGuide a r
-- --      )
-- --   => (a, a, a)
-- --   -> (a, a)
-- --   -> (a -> r)
-- --   -> Int
-- --   -> (Swarm m gen s g a r -> Bool)
-- --   -> ReaderT gen m [Swarm m gen s g a r]
-- -- optimise1D wpg bounds func n predicate = do
-- --   swarm <- simpleInit1D wpg bounds func n
-- --   iterateWhileM (not . predicate) updateSwarm swarm

optimiseND ::
     ( Monad m
     , HasPos α χ
     , GlobalGuide γ β α r
     )
  => m α
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> Int
  -> (Swarm m γ β α r -> Bool)
  -> m [Swarm m γ β α r]
optimiseND newState updater func n predicate = do
  swarm <- mkSwarm newState updater func n
  iterateWhileM (not . predicate) updateSwarm swarm

optimiseNDFromList ::
     ( Monad m
     , HasPos α χ
     , GlobalGuide γ β α r
     )
  => [α]
  -> PhaseUpdater m γ β α r
  -> (α -> m r)
  -> (Swarm m γ β α r -> Bool)
  -> (Swarm m γ β α r -> m ())
  -> m (Swarm m γ β α r)
optimiseNDFromList states updater func predicate process =
  do
    swarm <- mkSwarmFromList states updater func
    go swarm
  where
        go x
          | predicate x = return x
          | otherwise   = process x >> updateSwarm x >>= go

max_stars_repo_name: twesterhout/tcm-swarm

max_issues_repo_name: twesterhout/tcm-swarm

max_forks_repo_name: twesterhout/tcm-swarm


text:
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE Strict          #-}

module Optimization.Types where

import Control.Lens
import Numeric.LinearAlgebra            ( Numeric )
import Numeric.LinearAlgebra.Data
import Data.Vector.Storable       as VS ( map
                                        , singleton
                                        , zipWith )

{- |Basic "zippy" Num instance for storable vectors. Note that, while this is defined
    for some types of Vector by hmatrix, it unfortunately does not account for any
    Num (Vector a). -}
instance Numeric a => Num (Vector a) where
  a + b    = VS.zipWith (+)   a b
  a - b    = VS.zipWith (-)   a b
  a * b    = VS.zipWith (*)   a b
  abs a    = VS.map  (abs)    a
  signum a = VS.map  (signum) a
  fromInteger i = VS.singleton (fromInteger i)

{- |Termination criterion for a numerical optimization method.
    We terminate either on maxit iterations or once our guesses are within
    threshold of one another. -}
data TerminationCriterion a = TerminateAt {
  _maxit     :: Maybe Int, -- ^Maximum number of iterations (if any).
  _threshold :: Maybe a    -- ^Threshold used to declare convergence (if any).
  } deriving (Show, Eq)

-- |Step parameters for Nelder-Mead simplex search. a must be Storable and Num.
data NMStepParameters a = NMStepParameters {
  _epsilon :: a, -- ^Offset around initial guess to construct starting simplex.
  _alpha   :: a, -- ^Reflection coefficient for Nelder-Mead algorithm.
  _gamma   :: a, -- ^Expansion coefficient for Nelder-Mead algorithm.
  _rho     :: a, -- ^Contraction coefficient for Nelder-Mead algorithm.
  _sigma   :: a  -- ^Shrink coefficient for Nelder-Mead algorithm.
  } deriving (Show, Eq)

-- |Standard Nelder-Mead parameters.
defaultStep :: NMStepParameters Double
defaultStep = NMStepParameters 1.0 1.0 2.0 0.5 0.5

-- |The Nelder-Mead simplex.
type Simplex a = Matrix a

-- |Create a TerminationCriterion with just a convergence threshold.
terminateAt ::
  a ->                   -- ^Convengence threshold
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAt t = TerminateAt Nothing (Just t)

-- |Create a TerminationCriterion with just a maximum iteration threshold.
terminateAfter ::
  Int               ->   -- ^Maximum number of iterations.
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAfter m = TerminateAt (Just m) Nothing

-- |Create a TerminationCriterion with a convergence threshold and maximum iteration
-- |limit.
terminateAtOrAfter ::
  a      ->              -- ^Convergence threshold
  Int    ->              -- ^Maximum number of iterations
  TerminationCriterion a -- ^A TerminationCriterion for optimization.
terminateAtOrAfter t m = TerminateAt (Just m) (Just t)

makeLenses ''TerminationCriterion
makeLenses ''NMStepParameters

max_stars_repo_name: agbrooks/buzzwords

max_issues_repo_name: agbrooks/buzzwords

max_forks_repo_name: agbrooks/buzzwords


text:
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module      : Graphics.Color.Algebra.Elevator
-- Copyright   : (c) Alexey Kuleshevich 2018-2020
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Graphics.Color.Algebra.Elevator
  ( Elevator(..)
  , module Data.Word
  , clamp01
  ) where

import Data.Complex
import Data.Int
import Data.Typeable
import Data.Vector.Storable (Storable)
import Data.Vector.Unboxed (Unbox)
import Data.Word
import GHC.Float
import Text.Printf

infixl 7 //

defFieldFormat :: FieldFormat
defFieldFormat = FieldFormat Nothing Nothing Nothing Nothing False "" 'v'

-- | A class with a set of functions that allow for changing precision by shrinking and
-- streatching the values.
class (Show e, Eq e, Num e, Typeable e, Unbox e, Storable e) => Elevator e where
  maxValue :: e

  minValue :: e

  fieldFormat :: e -> FieldFormat
  fieldFormat _ = defFieldFormat

  -- | This is a pretty printer for the value.
  toShowS :: e -> ShowS
  default toShowS :: PrintfArg e => e -> ShowS
  toShowS e = formatArg e (fieldFormat e)

  -- | Values are scaled to @[0, 255]@ range.
  toWord8 :: e -> Word8

  -- | Values are scaled to @[0, 65535]@ range.
  toWord16 :: e -> Word16

  -- | Values are scaled to @[0, 4294967295]@ range.
  toWord32 :: e -> Word32

  -- | Values are scaled to @[0, 18446744073709551615]@ range.
  toWord64 :: e -> Word64

  -- | Values are scaled to @[0.0, 1.0]@ range.
  toRealFloat :: (Elevator a, RealFloat a) => e -> a

  -- | Values are scaled from @[0.0, 1.0]@ range.
  fromRealFloat :: (Elevator a, RealFloat a) => a -> e

  -- | Values are scaled to @[0.0, 1.0]@ range.
  toFloat :: e -> Float
  toFloat = toRealFloat

  -- | Values are scaled to @[0.0, 1.0]@ range.
  toDouble :: e -> Double
  toDouble = toRealFloat

  -- | Values are scaled from @[0.0, 1.0]@ range.
  fromDouble :: Double -> e
  fromDouble = fromRealFloat

  -- | Division that works for integral types as well as floating points. May throw an exception.
  (//) :: e -> e -> e

-- | Lower the precision
dropDown :: forall a b. (Integral a, Bounded a, Integral b, Bounded b) => a -> b
dropDown !e = fromIntegral $ fromIntegral e `div` ((maxBound :: a) `div`
                                                   fromIntegral (maxBound :: b))
{-# INLINE dropDown #-}

-- | Increase the precision
raiseUp :: forall a b. (Integral a, Bounded a, Integral b, Bounded b) => a -> b
raiseUp !e = fromIntegral e * ((maxBound :: b) `div` fromIntegral (maxBound :: a))
{-# INLINE raiseUp #-}

-- | Convert to fractional with value less than or equal to 1.
squashTo1 :: forall a b. (Fractional b, Integral a, Bounded a) => a -> b
squashTo1 !e = fromIntegral e / fromIntegral (maxBound :: a)
{-# INLINE squashTo1 #-}

-- | Convert to integral streaching it's value up to a maximum value.
stretch :: forall a b. (RealFloat a, Integral b, Bounded b) => a -> b
stretch !e = round (fromIntegral (maxBound :: b) * clamp01 e)
{-# INLINE stretch #-}

-- | Clamp a value to @[0, 1]@ range.
clamp01 :: RealFloat a => a -> a
clamp01 !x = min (max 0 x) 1
{-# INLINE clamp01 #-}

float2Word32 :: Float -> Word32
float2Word32 d'
  | d' <= 0 = 0
  | d > 4.294967e9 = maxBound
  | otherwise = round d
  where
    d = maxWord32 * d'
{-# INLINE float2Word32 #-}

-- | Same as:
-- >>> fromIntegral (maxBound :: Word32) :: Float
-- 4.2949673e9
--
maxWord32 :: Float
maxWord32 = F# 4.2949673e9#
{-# INLINE maxWord32 #-}

double2Word64 :: Double -> Word64
double2Word64 d'
  | d' <= 0 = 0
  | d > 1.844674407370955e19 = maxBound
  | otherwise = round d
  where
    d = maxWord64 * d'
{-# INLINE double2Word64 #-}

-- | Differs from `fromIntegral` due to: [GHC #17782](https://gitlab.haskell.org/ghc/ghc/issues/17782)
--
-- >>> fromIntegral (maxBound :: Word64) :: Double
-- 1.844674407370955e19
--
maxWord64 :: Double
maxWord64 = D# 1.8446744073709552e19##
{-# INLINE maxWord64 #-}

{-# RULES
"fromRealFloat :: Double -> Word" fromRealFloat = fromDouble :: Double -> Word
"fromRealFloat :: Double -> Word64" fromRealFloat = fromDouble :: Double -> Word64
"fromRealFloat :: Float -> Word32" fromRealFloat = float2Word32
 #-}

-- | Values between @[0, 255]]@
instance Elevator Word8 where
  maxValue = maxBound
  minValue = minBound
  fieldFormat _ = defFieldFormat {fmtWidth = Just 3, fmtChar = 'd'}
  toWord8 = id
  {-# INLINE toWord8 #-}
  toWord16 = raiseUp
  {-# INLINE toWord16 #-}
  toWord32 = raiseUp
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp
  {-# INLINE toWord64 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  fromDouble = toWord8
  {-# INLINE fromDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = toWord8
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 65535]]@
instance Elevator Word16 where
  maxValue = maxBound
  minValue = minBound
  fieldFormat _ = defFieldFormat { fmtWidth = Just 5, fmtChar = 'd'}
  toWord8 = dropDown
  {-# INLINE toWord8 #-}
  toWord16 = id
  {-# INLINE toWord16 #-}
  toWord32 = raiseUp
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp
  {-# INLINE toWord64 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  fromDouble = toWord16
  {-# INLINE fromDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = toWord16
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 4294967295]@
instance Elevator Word32 where
  maxValue = maxBound
  minValue = minBound
  fieldFormat _ = defFieldFormat { fmtWidth = Just 10, fmtChar = 'd'}
  toWord8 = dropDown
  {-# INLINE toWord8 #-}
  toWord16 = dropDown
  {-# INLINE toWord16 #-}
  toWord32 = id
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp
  {-# INLINE toWord64 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  fromDouble = toWord32
  {-# INLINE fromDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = toWord32
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 18446744073709551615]@
instance Elevator Word64 where
  maxValue = maxBound
  minValue = minBound
  fieldFormat _ = defFieldFormat { fmtWidth = Just 20, fmtChar = 'd'}
  toWord8 = dropDown
  {-# INLINE toWord8 #-}
  toWord16 = dropDown
  {-# INLINE toWord16 #-}
  toWord32 = dropDown
  {-# INLINE toWord32 #-}
  toWord64 = id
  {-# INLINE toWord64 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  fromDouble = double2Word64
  {-# INLINE fromDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = toWord64
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 18446744073709551615]@ on 64bit
instance Elevator Word where
  maxValue = maxBound
  minValue = minBound
#if WORD_SIZE_IN_BITS < 64
  fieldFormat _ = defFieldFormat { fmtWidth = Just 10, fmtChar = 'd'}
  toWord64 = dropDown
  {-# INLINE toWord64 #-}
  fromDouble = stretch
  {-# INLINE fromDouble #-}
#else
  fieldFormat _ = defFieldFormat { fmtWidth = Just 20, fmtChar = 'd'}
  toWord64 (W64# w#) = (W# w#)
  {-# INLINE toWord64 #-}
  fromDouble = toWord64 . double2Word64
  {-# INLINE fromDouble #-}
#endif
  toWord8 = dropDown
  {-# INLINE toWord8 #-}
  toWord16 = dropDown
  {-# INLINE toWord16 #-}
  toWord32 = dropDown
  {-# INLINE toWord32 #-}
  toFloat = squashTo1
  {-# INLINE toFloat #-}
  toDouble = squashTo1
  {-# INLINE toDouble #-}
  toRealFloat = squashTo1
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 127]@
instance Elevator Int8 where
  maxValue = maxBound
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 3, fmtChar = 'd'}
  toWord8 = fromIntegral . max 0
  {-# INLINE toWord8 #-}
  toWord16 = raiseUp . max 0
  {-# INLINE toWord16 #-}
  toWord32 = raiseUp . max 0
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp . max 0
  {-# INLINE toWord64 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 32767]@
instance Elevator Int16 where
  maxValue = maxBound
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 5, fmtChar = 'd'}
  toWord8 = dropDown . max 0
  {-# INLINE toWord8 #-}
  toWord16 = fromIntegral . max 0
  {-# INLINE toWord16 #-}
  toWord32 = raiseUp . max 0
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp . max 0
  {-# INLINE toWord64 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 2147483647]@
instance Elevator Int32 where
  maxValue = maxBound
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 10, fmtChar = 'd'}
  toWord8 = dropDown . max 0
  {-# INLINE toWord8 #-}
  toWord16 = dropDown . max 0
  {-# INLINE toWord16 #-}
  toWord32 = fromIntegral . max 0
  {-# INLINE toWord32 #-}
  toWord64 = raiseUp . max 0
  {-# INLINE toWord64 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 9223372036854775807]@
instance Elevator Int64 where
  maxValue = maxBound
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 19, fmtChar = 'd'}
  toWord8 = dropDown . max 0
  {-# INLINE toWord8 #-}
  toWord16 = dropDown . max 0
  {-# INLINE toWord16 #-}
  toWord32 = dropDown . max 0
  {-# INLINE toWord32 #-}
  toWord64 = fromIntegral . max 0
  {-# INLINE toWord64 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0, 9223372036854775807]@ on 64bit
instance Elevator Int where
  maxValue = maxBound
  minValue = 0
#if WORD_SIZE_IN_BITS < 64
  fieldFormat _ = defFieldFormat { fmtWidth = Just 10, fmtChar = 'd'}
  toWord64 = dropDown . max 0
  {-# INLINE toWord64 #-}
#else
  fieldFormat _ = defFieldFormat { fmtWidth = Just 19, fmtChar = 'd'}
  toWord64 = fromIntegral . max 0
  {-# INLINE toWord64 #-}
#endif
  toWord8 = dropDown . max 0
  {-# INLINE toWord8 #-}
  toWord16 = dropDown . max 0
  {-# INLINE toWord16 #-}
  toWord32 = dropDown . max 0
  {-# INLINE toWord32 #-}
  toFloat = squashTo1 . max 0
  {-# INLINE toFloat #-}
  toRealFloat = squashTo1 . max 0
  {-# INLINE toRealFloat #-}
  fromRealFloat = stretch
  {-# INLINE fromRealFloat #-}
  (//) = div
  {-# INLINE (//) #-}

-- | Values between @[0.0, 1.0]@
instance Elevator Float where
  maxValue = 1
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 11, fmtPrecision = Just 8, fmtChar = 'f'}
  toWord8 = stretch
  {-# INLINE toWord8 #-}
  toWord16 = stretch
  {-# INLINE toWord16 #-}
  toWord32 = float2Word32
  {-# INLINE toWord32 #-}
  toWord64 = stretch
  {-# INLINE toWord64 #-}
  toFloat = id
  {-# INLINE toFloat #-}
  toDouble = float2Double
  {-# INLINE toDouble #-}
  fromDouble = toFloat
  {-# INLINE fromDouble #-}
  toRealFloat = uncurry encodeFloat . decodeFloat
  {-# INLINE toRealFloat #-}
  fromRealFloat = uncurry encodeFloat . decodeFloat
  {-# INLINE fromRealFloat #-}
  (//) = (/)
  {-# INLINE (//) #-}

-- | Values between @[0.0, 1.0]@
instance Elevator Double where
  maxValue = 1
  minValue = 0
  fieldFormat _ = defFieldFormat { fmtWidth = Just 19, fmtPrecision = Just 16, fmtChar = 'f' }
  toWord8 = stretch
  {-# INLINE toWord8 #-}
  toWord16 = stretch
  {-# INLINE toWord16 #-}
  toWord32 = stretch
  {-# INLINE toWord32 #-}
  toWord64 = double2Word64
  {-# INLINE toWord64 #-}
  toFloat = double2Float
  {-# INLINE toFloat #-}
  toDouble = id
  {-# INLINE toDouble #-}
  fromDouble = id
  {-# INLINE fromDouble #-}
  toRealFloat = uncurry encodeFloat . decodeFloat
  {-# INLINE toRealFloat #-}
  fromRealFloat = uncurry encodeFloat . decodeFloat
  {-# INLINE fromRealFloat #-}
  (//) = (/)
  {-# INLINE (//) #-}

{-# RULES
"toRealFloat   :: Double -> Double / Float -> Float" toRealFloat = id
"toRealFloat   :: Double -> Float"                   toRealFloat = double2Float
"toRealFloat   :: Float -> Double"                   toRealFloat = float2Double
"fromRealFloat :: Double -> Double / Float -> Float" fromRealFloat = id
"fromRealFloat :: Double -> Float"                   fromRealFloat = double2Float
"fromRealFloat :: Float -> Double"                   fromRealFloat = float2Double
 #-}

-- | Discards imaginary part and changes precision of real part.
instance (PrintfArg e, Elevator e, RealFloat e) => Elevator (Complex e) where
  maxValue = maxValue :+ maxValue
  minValue = minValue :+ minValue
  toShowS (r :+ i) = toShowS r . formatArg i ((fieldFormat i) {fmtSign = Just SignPlus}) . ('i' :)
  toWord8 = toWord8 . realPart
  {-# INLINE toWord8 #-}
  toWord16 = toWord16 . realPart
  {-# INLINE toWord16 #-}
  toWord32 = toWord32 . realPart
  {-# INLINE toWord32 #-}
  toWord64 = toWord64 . realPart
  {-# INLINE toWord64 #-}
  toFloat = toFloat . realPart
  {-# INLINE toFloat #-}
  toDouble = toDouble . realPart
  {-# INLINE toDouble #-}
  fromDouble = (:+ 0) . fromDouble
  {-# INLINE fromDouble #-}
  toRealFloat = toRealFloat . realPart
  {-# INLINE toRealFloat #-}
  fromRealFloat = (:+ 0) . fromRealFloat
  {-# INLINE fromRealFloat #-}
  (//) = (/)
  {-# INLINE (//) #-}

max_stars_repo_name: TravisCardwell/Color

max_issues_repo_name: TravisCardwell/Color

max_forks_repo_name: TravisCardwell/Color


text:
{-# LANGUAGE FlexibleContexts #-}
{-# language BangPatterns        #-}
{-# language LambdaCase          #-}
{-# options_ghc -Wno-unused-imports #-}
module Data.Vector.FFT (
  fft, ifft
  -- * Useful results
  , crossCorrelation
  ) where

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad(..))

import Control.Monad.ST (runST)
import Data.Bits (shiftR,shiftL,(.&.),(.|.))
import Data.Bool (Bool,otherwise)
import Data.Complex (Complex(..),conjugate)
import Data.Foldable (forM_)

import Data.Vector.Unboxed as V (Vector, Unbox, map, zipWith, length, unsafeFreeze, (!))
import qualified Data.Vector.Unboxed.Mutable as VM (MVector, read, write, new, length)
import qualified Data.Vector.Generic as VG (Vector(..), copy)

import Prelude hiding (read)

{-# RULES
"fft/ifft" forall x. fft (ifft x) = x
"ifft/fft" forall x. ifft (fft x) = x
  #-}

-- | (Circular) cross-correlation of two vectors
--
-- Defined via the FFT and IFFT for computational efficiency
--
-- NB the source vectors should have matching length for meaningful results
crossCorrelation :: Vector (Complex Double)
                 -> Vector (Complex Double)
                 -> Vector (Complex Double)
crossCorrelation v1 v2 = ifft $ (cmap conjugate v1hat) `prod` v2hat
  where
    prod = V.zipWith (*)
    v1hat = fft v1
    v2hat = fft v2

-- | Radix-2 decimation-in-time fast Fourier Transform.
--
--   The given array (and therefore the output as well) is zero-padded to the next power of two if necessary.
fft :: Vector (Complex Double) -> Vector (Complex Double)
fft arr = runST $ do
  marr <- copyPadded arr
  mfft marr
  V.unsafeFreeze marr
{-# inlinable [1] fft #-}

-- | Inverse fast Fourier transform.
--
--   The given array (and therefore the output as well) is zero-padded to the next power of two if necessary.
ifft :: Vector (Complex Double) -> Vector (Complex Double)
ifft arr = do
  let lenComplex = intToComplexDouble (V.length arr)
  cmap ((/ lenComplex) . conjugate) . fft . cmap conjugate $ arr
{-# inlinable [1] ifft #-}

-- | Copy the source vector into a zero-padded mutable one
copyPadded :: (PrimMonad m, Num a, Unbox a) =>
              Vector a -> m (VM.MVector (PrimState m) a)
copyPadded arr = do
  let
    len = V.length arr
    l2 = nextPow2 len
  marr <- VM.new l2
  forM_ [0 .. l2 - 1] $ \i -> do
    let x | i < len = arr V.! i
          | otherwise = 0
    VM.write marr i x
  pure marr
{-# inline copyPadded #-}

-- | Radix-2 decimation-in-time fast Fourier Transform.
--   The given array must have a length that is a power of two,
--   though this property is not checked.
mfft :: (PrimMonad m) => VM.MVector (PrimState m) (Complex Double) -> m ()
mfft mut = do {
    let len = VM.length mut
  ; let bitReverse !i !j = do {
          ; if i == len - 1
              then stage 0 1
              else do {
                  when (i < j) $ swap mut i j
                ; let inner k l = if k <= l
                        then inner (k `shiftR` 1) (l - k)
                        else bitReverse (i + 1) (l + k)
                ; inner (len `shiftR` 1) j
              }
        }
        stage l l1 = if l == (log2 len)
          then pure ()
          else do {
              let !l2 = l1 `shiftL` 1
                  !e = (negate twoPi) / (intToDouble l2)
                  flight j !a = if j == l1
                    then stage (l + 1) l2
                    else do {
                        let butterfly i = if i >= len
                              then flight (j + 1) (a + e)
                              else do {
                                  let i1 = i + l1
                                ; xi1 :+ yi1 <- VM.read mut i1
                                ; let !co = cos a
                                      !si = sin a
                                      d = (co * xi1 - si * yi1) :+ (si * xi1 + co * yi1)
                                ; ci <- VM.read mut i
                                ; VM.write mut i1 (ci - d)
                                ; VM.write mut i (ci + d)
                                ; butterfly (i + l2)
                              }
                      ; butterfly j
                    }
            ; flight 0 0
         }
  ; bitReverse 0 0
}

-- wildcard cases should never happen. if they do, really bad things will happen.
b,s :: Int -> Int
b = \case { 0 -> 0x02; 1 -> 0x0c; 2 -> 0xf0; 3 -> 0xff00; 4 -> wordToInt 0xffff0000; 5 -> wordToInt 0xffffffff00000000; _ -> 0; }
s = \case { 0 -> 1; 1 -> 2; 2 -> 4; 3 -> 8; 4 -> 16; 5 -> 32; _ -> 0; }
{-# inline b #-}
{-# inline s #-}

-- | Next power of 2
nextPow2 :: Int -> Int
nextPow2 n
  | mod n 2 == 0 = n
  | otherwise = (2 :: Int) ^ (log2 n + 1)

log2 :: Int -> Int
log2 v0 = if v0 <= 0
  then error $ "Data.Vector.FFT: nonpositive input, got " ++ show v0
  else go 5 0 v0
  where
    go !i !r !v
      | i == -1 = r
      | v .&. b i /= 0 =
          let si = s i
          in go (i - 1) (r .|. si) (v `shiftR` si)
      | otherwise = go (i - 1) r v

{-# inline swap #-}
swap :: (PrimMonad m, Unbox a) =>
        VM.MVector (PrimState m) a -> Int -> Int -> m ()
swap mut i j = do
  atI <- VM.read mut i
  atJ <- VM.read mut j
  VM.write mut i atJ
  VM.write mut j atI

twoPi :: Double
{-# inline twoPi #-}
twoPi = 6.283185307179586

intToDouble :: Int -> Double
{-# inline intToDouble #-}
intToDouble = fromIntegral

wordToInt :: Word -> Int
{-# inline wordToInt #-}
wordToInt = fromIntegral

intToComplexDouble :: Int -> Complex Double
{-# inline intToComplexDouble #-}
intToComplexDouble = fromIntegral

{-# inline cmap #-}
cmap :: (Floating a, Unbox a) => (Complex a -> Complex a) -> V.Vector (Complex a) -> V.Vector (Complex a)
cmap = V.map

--

-- {-# inline copyWhole #-}
-- copyWhole :: (PrimMonad m, VG.Vector Vector a, Unbox a) => V.Vector a -> m (VM.MVector (PrimState m) a)
-- copyWhole arr = do
--   let len = V.length arr
--   marr <- VM.new len
--   VG.copy marr arr
--   pure marr

-- {-# inline arrOK #-}
-- arrOK :: Unbox a => Vector a -> Bool
-- arrOK arr =
--   let n = V.length arr
--   in (1 `shiftL` log2 n) == n

max_stars_repo_name: ocramz/vector-fft

max_issues_repo_name: ocramz/vector-fft

max_forks_repo_name: ocramz/vector-fft


text:
module Statistics.FastBayes where

import Statistics.FastBayes.Linear()

max_stars_repo_name: cscherrer/fastbayes

max_issues_repo_name: cscherrer/fastbayes

max_forks_repo_name: cscherrer/fastbayes


text:
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Packed.Statistics
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Basic multivariate statistics.
--

module Numeric.LinearAlgebra.Packed.Statistics (
    defaultCovUplo,

    -- * Immutable interface
    cov,
    covWithMean,
    weightedCov,
    weightedCovWithMean,

    -- * Mutable interface
    covTo,
    covWithMeanTo,
    weightedCovTo,
    weightedCovWithMeanTo,

    ) where

import Control.Monad( when )
import Control.Monad.ST( ST )
import Data.List( foldl' )
import Text.Printf( printf )

import Numeric.LinearAlgebra.Types
import Numeric.LinearAlgebra.Packed.Base( Packed, STPacked )
import qualified Numeric.LinearAlgebra.Packed.Base as P

import Numeric.LinearAlgebra.Vector( Vector, RVector )
import qualified Numeric.LinearAlgebra.Vector as V

import qualified Numeric.LinearAlgebra.Matrix as M

-- | Returns the default storage scheme for covariance matrices.
defaultCovUplo :: Uplo
defaultCovUplo = Lower

-- | Returns the sample covariance matrix hermitian matrix (in packed form)
-- with storage scheme equal to 'defaultCovUplo'.  The first argument gives
-- the dimension of the vectors.
cov :: (BLAS2 e)
    => Int -> CovMethod -> [Vector e] -> Herm Packed e
cov p t xs = P.hermCreate $ do
    c <- Herm uplo `fmap` P.new_ p
    covTo c t xs
    return c
  where
    uplo = defaultCovUplo

-- | Given the pre-computed mean, returns the sample covariance matrix
-- (in packed form) with storage scheme equal to 'defaultCovUplo'.
covWithMean :: (BLAS2 e)
            => Vector e -> CovMethod -> [Vector e] -> Herm Packed e
covWithMean mu t xs = P.hermCreate $ do
    c <- Herm uplo `fmap` P.new_ p
    covWithMeanTo c mu t xs
    return c
  where
    p = V.dim mu
    uplo = defaultCovUplo

-- | Returns the weighed sample covariance matrix (in packed form) with
-- storage scheme equal to 'defaultCovUplo'. The first argument gives the
-- dimension of the vectors.
weightedCov :: (BLAS2 e)
            => Int -> CovMethod -> [(Double, Vector e)] -> Herm Packed e
weightedCov p t wxs = P.hermCreate $ do
    c <- Herm uplo `fmap` P.new_ p
    weightedCovTo c t wxs
    return c
  where
    uplo = defaultCovUplo

-- | Given the pre-computed mean, returns the weighed sample covariance matrix
-- (in packed form) with storage scheme equal to 'defaultCovUplo'.
weightedCovWithMean :: (BLAS2 e)
                    => Vector e -> CovMethod -> [(Double, Vector e)]
                    -> Herm Packed e
weightedCovWithMean mu t wxs = P.hermCreate $ do
    c <- Herm uplo `fmap` P.new_ p
    weightedCovWithMeanTo c mu t wxs
    return c
  where
    p = V.dim mu
    uplo = defaultCovUplo

-- | Computes and copies the sample covariance matrix (in packed form)
-- to the given destination.
covTo :: (RVector v, BLAS2 e)
      => Herm (STPacked s) e -> CovMethod -> [v e] -> ST s ()
covTo c@(Herm _ a) t xs = do
    p <- P.getDim a
    mu <- V.new p 1
    V.meanTo mu xs
    covWithMeanTo c mu t xs

-- | Given the pre-computed mean, computes and copies the sample covariance
-- matrix (in packed form) to the given destination.
covWithMeanTo :: (RVector v1, RVector v2, BLAS2 e)
              => Herm (STPacked s) e
              -> v1 e -> CovMethod -> [v2 e]
              -> ST s ()
covWithMeanTo c@(Herm _ a) mu t xs = do
    pa <- P.getDim a
    p <- V.getDim mu

    when (pa /= p) $ error $
        printf ("covWithMeanTo"
                ++ " (Herm _ <packed matrix with dim %d>)"
                ++ " <vector with dim %d>"
                ++ " _ _"
                ++ ": dimension mismatch")
               pa p

    xt <- M.new_ (p,n)
    M.withColsM xt $ \xs' ->
        sequence_ [ V.subTo x' mu x
                  | (x,x') <- zip xs xs'
                  ]
    P.withVectorM a V.clear
    M.withColsM xt $ \xs' ->
        sequence_ [ P.hermRank1UpdateM_ scale x' c | x' <- xs' ]
  where
    n = length xs
    df = fromIntegral $ case t of { MLCov -> n ; UnbiasedCov -> n - 1 }
    scale = 1/df

-- | Computes and copies the weighed sample covariance matrix (in packed
-- form) to the given destination.
weightedCovTo :: (RVector v, BLAS2 e)
              => Herm (STPacked s) e
              -> CovMethod -> [(Double, v e)]
              -> ST s ()
weightedCovTo c@(Herm _ a) t wxs = do
    p <- P.getDim a
    mu <- V.new p 1
    V.weightedMeanTo mu wxs
    weightedCovWithMeanTo c mu t wxs

-- | Given the pre-computed mean, computes and copies the weighed sample
-- covariance matrix (in packed form) to the given destination.
weightedCovWithMeanTo :: (RVector v1, RVector v2, BLAS2 e)
                      => Herm (STPacked s) e
                      -> v1 e -> CovMethod -> [(Double, v2 e)]
                      -> ST s ()
weightedCovWithMeanTo c@(Herm _ a) mu t wxs = do
    pa <- P.getDim a
    p <- V.getDim mu

    when (pa /= p) $ error $
        printf ("weightedCovWithMeanTo"
                ++ " (Herm _ <packed matrix with dim %d>)"
                ++ " <vector with dim %d>"
                ++ " _ _"
                ++ ": dimension mismatch")
               pa p

    xt <- M.new_ (p,n)
    M.withColsM xt $ \xs' ->
        sequence_ [  V.subTo x' mu x
                  >> V.scaleM_ (realToFrac $ sqrt (w / invscale)) x'
                  |  (w,x,x') <- zip3 ws xs xs'
                  ]
    P.withVectorM a V.clear
    M.withCols xt $ \xs' ->
        sequence_ [ P.hermRank1UpdateM_ 1 x' c | x' <- xs' ]
  where
    (ws0,xs) = unzip wxs
    w_sum = foldl' (+) 0 ws0
    ws = if w_sum == 0 then ws0 else map (/w_sum) ws0
    w2s_sum = foldl' (+) 0 $ map (^^(2::Int)) ws
    invscale = case t of
                   MLCov -> 1
                   UnbiasedCov -> (1 - w2s_sum)
    n = length ws0

max_stars_repo_name: patperry/hs-linear-algebra

max_issues_repo_name: cartazio/hs-cblas

max_forks_repo_name: cartazio/hs-cblas