{-# LANGUAGE
    MultiParamTypeClasses,
    FlexibleContexts, FlexibleInstances,
    UndecidableInstances, GADTs
  #-}

module Data.Random.Distribution.Simplex
    ( StdSimplex(..)
    , stdSimplex
    , stdSimplexT
    , fractionalStdSimplex
    ) where

import Control.Monad
import Data.List
import Data.Random.RVar
import Data.Random.Distribution
import Data.Random.Distribution.Uniform

-- |Uniform distribution over a standard simplex.
newtype StdSimplex as =
    -- | @StdSimplex k@ constructs a standard simplex of dimension @k@
    -- (standard /k/-simplex).
    -- An element of the simplex represents a vector variable @as = (a_0,
    -- a_1, ..., a_k)@. The elements of @as@ are more than or equal to @0@
    -- and @sum as@ is always equal to @1@.
    StdSimplex Int
    deriving (StdSimplex as -> StdSimplex as -> Bool
forall as. StdSimplex as -> StdSimplex as -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: StdSimplex as -> StdSimplex as -> Bool
$c/= :: forall as. StdSimplex as -> StdSimplex as -> Bool
== :: StdSimplex as -> StdSimplex as -> Bool
$c== :: forall as. StdSimplex as -> StdSimplex as -> Bool
Eq, Int -> StdSimplex as -> ShowS
forall as. Int -> StdSimplex as -> ShowS
forall as. [StdSimplex as] -> ShowS
forall as. StdSimplex as -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [StdSimplex as] -> ShowS
$cshowList :: forall as. [StdSimplex as] -> ShowS
show :: StdSimplex as -> String
$cshow :: forall as. StdSimplex as -> String
showsPrec :: Int -> StdSimplex as -> ShowS
$cshowsPrec :: forall as. Int -> StdSimplex as -> ShowS
Show)

instance (Ord a, Fractional a, Distribution StdUniform a) => Distribution StdSimplex [a] where
    rvar :: StdSimplex [a] -> RVar [a]
rvar (StdSimplex Int
k) = forall a.
(Ord a, Fractional a, Distribution StdUniform a) =>
Int -> RVar [a]
fractionalStdSimplex Int
k

-- |@stdSimplex k@ returns a random variable being uniformly distributed over
-- a standard simplex of dimension @k@.
stdSimplex :: Distribution StdSimplex [a] => Int -> RVar [a]
stdSimplex :: forall a. Distribution StdSimplex [a] => Int -> RVar [a]
stdSimplex Int
k = forall (d :: * -> *) t. Distribution d t => d t -> RVar t
rvar (forall as. Int -> StdSimplex as
StdSimplex Int
k)

stdSimplexT :: Distribution StdSimplex [a] => Int -> RVarT m [a]
stdSimplexT :: forall a (m :: * -> *).
Distribution StdSimplex [a] =>
Int -> RVarT m [a]
stdSimplexT Int
k = forall (d :: * -> *) t (n :: * -> *).
Distribution d t =>
d t -> RVarT n t
rvarT (forall as. Int -> StdSimplex as
StdSimplex Int
k)

-- |An algorithm proposed by Rubinstein & Melamed (1998).
-- See, /e.g./, S. Onn, I. Weissman.
-- Generating uniform random vectors over a simplex with implications to
-- the volume of a certain polytope and to multivariate extremes.
-- /Ann Oper Res/ (2011) __189__:331-342.
fractionalStdSimplex :: (Ord a, Fractional a, Distribution StdUniform a) => Int -> RVar [a]
fractionalStdSimplex :: forall a.
(Ord a, Fractional a, Distribution StdUniform a) =>
Int -> RVar [a]
fractionalStdSimplex Int
k = do [a]
us <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
k forall a. Distribution StdUniform a => RVar a
stdUniform
                            let us' :: [a]
us' = forall a. Ord a => [a] -> [a]
sort [a]
us forall a. [a] -> [a] -> [a]
++ [a
1]
                            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [a]
us' (a
0 forall a. a -> [a] -> [a]
: [a]
us')