{-# LANGUAGE PatternGuards #-}

module Numeric.Histogram ( Range
                         , binBounds
                         , histValues
                         , histWeightedValues
                         , histWithBins
                         ) where

import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Control.Monad.ST

type Range a = (a,a)

-- | 'binBounds a b n' generates bounds for 'n' bins spaced linearly between
-- 'a' and 'b'
--
-- Examples:
--
-- >>> binBounds 0 3 4
-- [(0.0,0.75),(0.75,1.5),(1.5,2.25),(2.25,3.0)]
binBounds :: RealFrac a => a -> a -> Int -> [Range a]
binBounds :: a -> a -> Int -> [Range a]
binBounds a
a a
b Int
n = (Int -> Range a) -> [Int] -> [Range a]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i->(Int -> a
forall a. Real a => a -> a
lbound Int
i, Int -> a
forall a. Real a => a -> a
lbound (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))) [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
        where lbound :: a -> a
lbound a
i = a
a a -> a -> a
forall a. Num a => a -> a -> a
+ (a
ba -> a -> a
forall a. Num a => a -> a -> a
-a
a) a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
i a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Real a, Fractional b) => a -> b
realToFrac Int
n

-- | 'histValues a b n vs' returns the bins for the histogram of
-- 'vs' on the range from 'a' to 'b' with 'n' bins
histValues :: RealFrac a => a -> a -> Int -> [a] -> V.Vector (Range a, Int)
histValues :: a -> a -> Int -> [a] -> Vector (Range a, Int)
histValues a
a a
b Int
n = Vector (Range a) -> [(Int, a)] -> Vector (Range a, Int)
forall w a.
(Num w, RealFrac a) =>
Vector (Range a) -> [(w, a)] -> Vector (Range a, w)
histWithBins ([Range a] -> Vector (Range a)
forall a. [a] -> Vector a
V.fromList ([Range a] -> Vector (Range a)) -> [Range a] -> Vector (Range a)
forall a b. (a -> b) -> a -> b
$ a -> a -> Int -> [Range a]
forall a. RealFrac a => a -> a -> Int -> [Range a]
binBounds a
a a
b Int
n) ([(Int, a)] -> Vector (Range a, Int))
-> ([a] -> [(Int, a)]) -> [a] -> Vector (Range a, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int -> [Int]
forall a. a -> [a]
repeat Int
1)

-- | 'histValues a b n vs' returns the bins for the weighted histogram of
-- 'vs' on the range from 'a' to 'b' with 'n' bins
histWeightedValues :: RealFrac a => a -> a -> Int -> [(Double,a)] -> V.Vector (Range a, Double)
histWeightedValues :: a -> a -> Int -> [(Double, a)] -> Vector (Range a, Double)
histWeightedValues a
a a
b Int
n = Vector (Range a) -> [(Double, a)] -> Vector (Range a, Double)
forall w a.
(Num w, RealFrac a) =>
Vector (Range a) -> [(w, a)] -> Vector (Range a, w)
histWithBins ([Range a] -> Vector (Range a)
forall a. [a] -> Vector a
V.fromList ([Range a] -> Vector (Range a)) -> [Range a] -> Vector (Range a)
forall a b. (a -> b) -> a -> b
$ a -> a -> Int -> [Range a]
forall a. RealFrac a => a -> a -> Int -> [Range a]
binBounds a
a a
b Int
n)

-- | 'histWithBins bins xs' is the histogram of weighted values 'xs' with 'bins'
--
-- Examples:
--
-- >>> :{
-- histWithBins
--     (V.fromList [(0.0, 0.75), (0.75, 1.5), (1.5, 2.25), (2.25, 3.0)])
--     [(1, 0), (1, 0), (1, 1), (1, 2), (1, 2), (1, 2), (1, 3)]
-- :}
-- [((0.0,0.75),2),((0.75,1.5),1),((1.5,2.25),3),((2.25,3.0),1)]
histWithBins :: (Num w, RealFrac a) => V.Vector (Range a) -> [(w, a)] -> V.Vector (Range a, w)
histWithBins :: Vector (Range a) -> [(w, a)] -> Vector (Range a, w)
histWithBins Vector (Range a)
bins [(w, a)]
xs =
    let n :: Int
n = Vector (Range a) -> Int
forall a. Vector a -> Int
V.length Vector (Range a)
bins
        testBin :: RealFrac a => a -> (Int, Range a) -> Bool
        testBin :: a -> (Int, Range a) -> Bool
testBin a
x (Int
i, (a
a,a
b)) =
            if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                then a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
a Bool -> Bool -> Bool
&& a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
b
                else a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
a Bool -> Bool -> Bool
&& a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
b

        f :: (RealFrac a, Num w)
          => V.Vector (Range a) -> MV.STVector s w -> (w, a)
          -> ST s ()
        f :: Vector (Range a) -> STVector s w -> (w, a) -> ST s ()
f Vector (Range a)
bins1 STVector s w
bs (w
w,a
x) =
            case ((Int, Range a) -> Bool)
-> Vector (Int, Range a) -> Vector (Int, Range a)
forall a. (a -> Bool) -> Vector a -> Vector a
V.dropWhile (Bool -> Bool
not (Bool -> Bool)
-> ((Int, Range a) -> Bool) -> (Int, Range a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> (Int, Range a) -> Bool
forall a. RealFrac a => a -> (Int, Range a) -> Bool
testBin a
x) (Vector (Int, Range a) -> Vector (Int, Range a))
-> Vector (Int, Range a) -> Vector (Int, Range a)
forall a b. (a -> b) -> a -> b
$ Vector (Range a) -> Vector (Int, Range a)
forall a. Vector a -> Vector (Int, a)
V.indexed Vector (Range a)
bins1 of
                Vector (Int, Range a)
v | Vector (Int, Range a) -> Bool
forall a. Vector a -> Bool
V.null Vector (Int, Range a)
v  -> () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                Vector (Int, Range a)
v | (Int
idx,Range a
_) <- Vector (Int, Range a) -> (Int, Range a)
forall a. Vector a -> a
V.head Vector (Int, Range a)
v  -> do
                    w
m <- MVector (PrimState (ST s)) w -> Int -> ST s w
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
MV.read STVector s w
MVector (PrimState (ST s)) w
bs Int
idx
                    MVector (PrimState (ST s)) w -> Int -> w -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
MV.write STVector s w
MVector (PrimState (ST s)) w
bs Int
idx (w -> ST s ()) -> w -> ST s ()
forall a b. (a -> b) -> a -> b
$! w
mw -> w -> w
forall a. Num a => a -> a -> a
+w
w

        counts :: Vector w
counts = (forall s. ST s (Vector w)) -> Vector w
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector w)) -> Vector w)
-> (forall s. ST s (Vector w)) -> Vector w
forall a b. (a -> b) -> a -> b
$ do STVector s w
b <- Int -> w -> ST s (MVector (PrimState (ST s)) w)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MVector (PrimState m) a)
MV.replicate Int
n w
0
                            ((w, a) -> ST s ()) -> [(w, a)] -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Vector (Range a) -> STVector s w -> (w, a) -> ST s ()
forall a w s.
(RealFrac a, Num w) =>
Vector (Range a) -> STVector s w -> (w, a) -> ST s ()
f Vector (Range a)
bins STVector s w
b) [(w, a)]
xs
                            MVector (PrimState (ST s)) w -> ST s (Vector w)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.freeze STVector s w
MVector (PrimState (ST s)) w
b
    in Vector (Range a) -> Vector w -> Vector (Range a, w)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip Vector (Range a)
bins Vector w
counts