{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module Data.Array.Accelerate.KullbackLiebler ( kullbackLiebler
, entropy
, dropZeroes
, scale
, hellinger
, fDivergence
, alphaDivergence
) where
import qualified Data.Array.Accelerate as A
-- | \( D_f(p \| q) = \displaystyle\int p(x) f\left(\frac{p(x)}{q(x)}\right) dx \)
--
-- @since 0.1.2.0
fDivergence :: (A.Floating e) => (A.Exp e -> A.Exp e) -- ^ \(f\)
-> A.Acc (A.Vector e)
-> A.Acc (A.Vector e)
-> A.Acc (A.Scalar e)
fDivergence f ps qs = A.sum (A.zipWith (\p q -> p * f (p / q)) ps qs)
-- | \( D^{(\alpha)}(p\| q) = \frac{4}{1 - \alpha^2}\left(1 - \displaystyle\int p(x)^{\frac{1-\alpha}{2}} q(x)^{\frac{1+\alpha}{2}} dx\right)\)
-- for \( \alpha \neq \pm 1\)
--
-- @since 0.1.2.0
alphaDivergence :: A.Floating e => A.Exp e -> A.Acc (A.Vector e) -> A.Acc (A.Vector e) -> A.Acc (A.Scalar e)
alphaDivergence α ps qs = A.map (\x -> (4 / (1 - α ** 2)) * (1 - x)) integrand
where integrand = A.sum (A.zipWith (\p q -> p ** ((1 - α)/2) * q ** ((1 + α)/2)) ps qs)
-- | Hellinger distance
--
-- @since 0.1.2.0
hellinger :: (A.Floating e) => A.Acc (A.Vector e) -> A.Acc (A.Vector e) -> A.Acc (A.Scalar e)
hellinger ps qs = A.map (A.sqrt . (2*)) $ A.sum (A.zipWith (\p q -> (A.sqrt p - A.sqrt q) ** 2) ps qs)
-- | Assumes input is nonzero
kullbackLiebler :: (A.Floating e) => A.Acc (A.Vector e) -> A.Acc (A.Vector e) -> A.Acc (A.Scalar e)
kullbackLiebler ps qs = A.sum (A.zipWith (\p q -> p * log (p / q)) ps qs)
-- | Assumes input is nonzero
entropy :: (A.Floating e) => A.Acc (A.Vector e) -> A.Acc (A.Scalar e)
entropy = A.sum . A.map (\p -> p * log p)
dropZeroes :: (A.Eq e, Num (A.Exp e)) => A.Acc (A.Vector e) -> A.Acc (A.Vector e)
dropZeroes = A.afst . A.filter (A./= 0)
-- | Doesn't check for negative values
--
-- @since 0.1.1.0
scale :: A.Floating e => A.Acc (A.Vector e) -> A.Acc (A.Vector e)
scale xs =
let tot = A.the $ A.sum xs
in A.map (/tot) xs