module Control.SpaceProbe.Internal.Probe where
import Control.Applicative (
Alternative,
Applicative,
empty,
pure,
(<*>),
(<|>),
(<$>))
import Data.Number.Erf (InvErf, invnormcdf, normcdf)
import Data.Tree(Forest, Tree(..))
data Probe t = forall s . Probe {
_initial :: s,
_partition :: s -> Forest s,
_draw :: s -> Maybe t
}
newPartition :: (s -> [s]) -> (s -> Forest s)
newPartition f = map (\x -> Node x []) . f
tipConcatMap :: (a -> Forest a) -> Tree a -> Tree a
tipConcatMap f (Node x []) = Node x $ f x
tipConcatMap f (Node x xs) = Node x $
case xs of
[] -> f x
_ -> map (tipConcatMap f) xs
instance Functor Probe where
fmap g (Probe x f d) = Probe x f $ fmap g . d
instance Applicative Probe where
pure x = Probe x (const []) (Just . id)
(Probe x0 f0 d0) <*> (Probe x1 f1 d1) = Probe (x0, x1) f d
where f (s0, s1) = map (tipConcatMap (t1 . fst)) $ t0 s1
where t0 s1' = [do {s0' <- t; return (s0', s1')} | t <- f0 s0]
t1 s0' = [do {s1' <- t; return (s0', s1')} | t <- f1 s1]
d (s0, s1) = d0 s0 <*> d1 s1
instance Alternative Probe where
empty = Probe [] (const []) undefined
(Probe x0 f0 d0) <|> (Probe x1 f1 d1) =
Probe {
_initial = Nothing,
_partition = partition,
_draw = draw
} where partition Nothing = map (flip Node [] . Just) [Left x0, Right x1]
partition (Just (Left x)) = map (fmap $ Just . Left) $ f0 x
partition (Just (Right x)) = map (fmap $ Just . Right) $ f1 x
draw m = m >>= either d0 d1
ave :: Num a => (a -> a -> a) -> a -> a -> a
ave divide a b = a + (b a) `divide` 2
floatAve :: Floating a => a -> a -> a
floatAve = ave (/)
intAve :: Integral a => a -> a -> a
intAve = ave quot
distribution :: Floating b => (b -> a) -> Probe a
distribution invcdf = invcdf <$> uniform 0 1
exponential :: Floating a => a -> Probe a
exponential mu = distribution $ \x -> mu * log(1 / (1 x))
normal :: (Eq a, InvErf a) => a -> a -> Probe a
normal mu sigma = (\x -> x * sigma + mu) <$> distribution invnormcdf
uniform :: Floating a => a -> a -> Probe a
uniform a b = Probe {
_initial = (a, floatAve a b, b),
_partition = newPartition $
\(a', x, b') -> [(a', floatAve a' x, x),
(x, floatAve x b', b')],
_draw = \(_, x, _) -> Just x
}
bisect :: (Integral a, Num b, Ord b) => (a -> b) -> b -> a -> a -> a
bisect f y a b = go (intAve a b) a b
where go u a' b'
| b' <= a' = u
| otherwise = let v = intAve a' b'
z = f v
in case compare z y of
LT -> go v (v + 1) b'
EQ -> v
GT -> go v a' v
intDistribution :: (Integral a, Floating b, Ord b) =>
(a -> b) ->
a ->
a ->
Probe a
intDistribution cdf a b = Probe {
_initial = (a, mid a b, b),
_partition = newPartition partition,
_draw = \(_, x, _) -> Just x
} where mid a' b' = bisect cdf (floatAve (cdf a') (cdf b')) a' b'
partition (a', x, b') =
filter (\(u, _, v) -> v > u)
[(a', mid a' x, x),
(x + 1, mid (x + 1) b', b')]
exponentialInt :: (Bounded a, Integral a) => Float -> Probe a
exponentialInt mu =
intDistribution (\x -> 1 exp (fromIntegral x / mu)) 0 maxBound
normalInt :: (Bounded a, Integral a) => Float -> Float -> Probe a
normalInt mu sigma =
intDistribution (\x -> normcdf $ (fromIntegral x mu) / sigma)
(round (mu bound) + 1)
(round (mu + bound) 1)
where bound = 6 * sigma
uniformInt :: (Eq a, Integral a) => a -> a -> Probe a
uniformInt a b = Probe {
_initial = (a, intAve a b, b),
_partition = newPartition partition',
_draw = \(_, x, _) -> Just x
} where partition' (a', x, b') =
filter (\(u, _, v) -> v > u) [(a', intAve a' x, x),
(x + 1, intAve (x + 1) b', b')]
constants :: [a] -> Probe a
constants xs = Probe Nothing partition id
where partition Nothing = map (flip Node [] . Just) xs
partition _ = []
permutation :: [Integer] -> [a] -> Integer -> [a]
permutation _ [] _ = []
permutation [] _ _ = error "permutation: factorial list too short"
permutation (u:facs) xs n = y : permutation facs ys r
where (q, r) = quotRem n u
(y:ys) = let (us, v:vs) = splitAt (fromIntegral q) xs
in v:us ++ vs
permute :: [a] -> Probe [a]
permute xs = permutation facs xs <$> uniformInt 0 u
where (u:facs) =
reverse $ scanl (*) 1 [1..fromIntegral $ length xs] :: [Integer]
extractElem :: [a] -> [(a, [a])]
extractElem [] = []
extractElem (x:xs) = (x, xs) : map (\(y, ys) -> (y, x:ys)) (extractElem xs)
sizedSublist :: Int -> [a] -> Probe [a]
sizedSublist k xs = Probe {
_initial = (0, xs, []),
_partition = newPartition $
\(n, xs', ys) ->
if n == k
then []
else [(n + 1, zs, x:ys) | (x, zs) <- extractElem xs'],
_draw = \(m, xs', ys) -> Just $ ys ++ take (k m) xs'
}
sizedWithReplacement :: Int -> [a] -> Probe [a]
sizedWithReplacement k xs = Probe {
_initial = (0, []),
_partition = newPartition $ \(n, ys) -> if n == k
then []
else [(n + 1, x:ys) | x <- xs],
_draw = Just . snd
}
sublist :: [a] -> Probe [a]
sublist xs = Probe {
_initial = (xs, []),
_partition = newPartition $ \(xs', ys) ->
case xs' of
[] -> []
(x:xs'') -> [(xs'', x:ys), (xs'', ys)],
_draw = Just . reverse . snd
}
withReplacement :: [a] -> Probe [a]
withReplacement xs = Probe {
_initial = [],
_partition = newPartition $ \ys -> map (:ys) xs,
_draw = Just . id
}