{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | Parallel scatter-accumulate aggregation kernel.
module DataFrame.Internal.AggKernelPar (
    scatterReducePar,
    momentScatterPar,
) where

import Control.Concurrent (forkIO, getNumCapabilities)
import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Control.Exception (SomeException, throwIO, try)
import Control.Monad (when)
import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import System.IO.Unsafe (unsafePerformIO)
import Type.Reflection (typeRep)

import DataFrame.Internal.AggKernel (
    Reduction (..),
    scatterColumnToDouble,
    scatterReduce,
 )
import DataFrame.Internal.AggPlan (Moments (..), momentScatter)
import DataFrame.Internal.Column (
    Column (..),
    Columnable,
    fromUnboxedVector,
    materializePacked,
 )

parThreshold :: Int
parThreshold :: Int
parThreshold = Int
200000

capabilities :: Int
capabilities :: Int
capabilities = IO Int -> Int
forall a. IO a -> a
unsafePerformIO IO Int
getNumCapabilities
{-# NOINLINE capabilities #-}

-- | Whether to take the parallel path at this row count.
shouldPar :: Int -> Bool
shouldPar :: Int -> Bool
shouldPar Int
n = Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
parThreshold Bool -> Bool -> Bool
&& Int
capabilities Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1

groupRangeBounds :: VU.Vector Int -> Int -> Int -> VU.Vector Int
groupRangeBounds :: Vector Int -> Int -> Int -> Vector Int
groupRangeBounds Vector Int
offs Int
nGroups Int
caps = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
    MVector s Int
b <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new (Int
caps Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    let !nRows :: Int
nRows = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
nGroups
        !per :: Int
per = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 ((Int
nRows Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
caps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
caps)
        adv :: Int -> Int -> Int
adv !Int
target !Int
gg
            | Int
gg Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = Int
nGroups
            | Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
gg Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
target = Int
gg
            | Bool
otherwise = Int -> Int -> Int
adv Int
target (Int
gg Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        go :: Int -> Int -> ST s ()
go !Int
w !Int
prev
            | Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
caps = MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
b Int
caps Int
nGroups
            | Bool
otherwise = do
                let !target :: Int
target = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
nRows (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
per)
                    !g :: Int
g = Int -> Int -> Int
adv Int
target Int
prev
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
b Int
w Int
g
                Int -> Int -> ST s ()
go (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
g
    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
b Int
0 Int
0
    Int -> Int -> ST s ()
go Int
1 Int
0
    MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
b

forEachRange :: VU.Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange :: Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange Vector Int
bounds Int
caps Int -> Int -> IO ()
act
    | Int
caps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = Int -> Int -> IO ()
act (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
bounds Int
0) (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
bounds Int
caps)
    | Bool
otherwise = do
        [MVar (Either SomeException ())]
vars <- (Int -> IO (MVar (Either SomeException ())))
-> [Int] -> IO [MVar (Either SomeException ())]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Int -> IO (MVar (Either SomeException ()))
spawn [Int
0 .. Int
caps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        [Either SomeException ()]
results <- (MVar (Either SomeException ()) -> IO (Either SomeException ()))
-> [MVar (Either SomeException ())] -> IO [Either SomeException ()]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM MVar (Either SomeException ()) -> IO (Either SomeException ())
forall a. MVar a -> IO a
takeMVar [MVar (Either SomeException ())]
vars
        (Either SomeException () -> IO ())
-> [Either SomeException ()] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((SomeException -> IO ())
-> (() -> IO ()) -> Either SomeException () -> IO ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (SomeException -> IO ()
forall e a. Exception e => e -> IO a
throwIO :: SomeException -> IO ()) () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) [Either SomeException ()]
results
  where
    spawn :: Int -> IO (MVar (Either SomeException ()))
spawn Int
w = do
        MVar (Either SomeException ())
var <- IO (MVar (Either SomeException ()))
forall a. IO (MVar a)
newEmptyMVar
        let !s :: Int
s = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
bounds Int
w
            !e :: Int
e = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
bounds (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
try (Int -> Int -> IO ()
act Int
s Int
e) IO (Either SomeException ())
-> (Either SomeException () -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar (Either SomeException ()) -> Either SomeException () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Either SomeException ())
var)
        MVar (Either SomeException ())
-> IO (MVar (Either SomeException ()))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVar (Either SomeException ())
var

scatterReducePar ::
    Reduction -> VU.Vector Int -> VU.Vector Int -> Int -> Column -> Maybe Column
scatterReducePar :: Reduction
-> Vector Int -> Vector Int -> Int -> Column -> Maybe Column
scatterReducePar Reduction
red Vector Int
vis Vector Int
offs Int
nGroups Column
col
    | Bool -> Bool
not (Int -> Bool
shouldPar (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
vis)) Bool -> Bool -> Bool
|| Int
nGroups Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 =
        Reduction -> Vector Int -> Int -> Column -> Maybe Column
scatterReduce Reduction
red (Vector Int -> Vector Int -> Int -> Vector Int
rtgFromVis Vector Int
vis Vector Int
offs Int
nGroups) Int
nGroups Column
col
    | Bool
otherwise = case Column
col of
        UnboxedColumn Maybe Bitmap
Nothing (Vector a
v :: VU.Vector a) ->
            case TypeRep a -> TypeRep Int -> Maybe (a :~: Int)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Int) of
                Just a :~: Int
Refl -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Reduction
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Idents a
-> Column
forall a.
(Columnable a, Unbox a, Num a, Ord a, Real a) =>
Reduction
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Idents a
-> Column
reduceParTyped Reduction
red Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Idents a
Idents Int
intIdent)
                Maybe (a :~: Int)
Nothing -> case TypeRep a -> TypeRep Double -> Maybe (a :~: Double)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Double) of
                    Just a :~: Double
Refl -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Reduction
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Idents a
-> Column
forall a.
(Columnable a, Unbox a, Num a, Ord a, Real a) =>
Reduction
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Idents a
-> Column
reduceParTyped Reduction
red Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Idents a
Idents Double
dblIdent)
                    Maybe (a :~: Double)
Nothing -> Maybe Column
forall a. Maybe a
Nothing
        p :: Column
p@(PackedText Maybe Bitmap
_ PackedTextData
_) -> Reduction
-> Vector Int -> Vector Int -> Int -> Column -> Maybe Column
scatterReducePar Reduction
red Vector Int
vis Vector Int
offs Int
nGroups (Column -> Column
materializePacked Column
p)
        Column
_ -> Maybe Column
forall a. Maybe a
Nothing
{-# NOINLINE scatterReducePar #-}

rtgFromVis :: VU.Vector Int -> VU.Vector Int -> Int -> VU.Vector Int
rtgFromVis :: Vector Int -> Vector Int -> Int -> Vector Int
rtgFromVis Vector Int
vis Vector Int
offs Int
nGroups = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
    let n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
vis
    MVector s Int
rtg <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
n)
    let go :: Int -> ST s ()
go !Int
g
            | Int
g Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                let !e :: Int
e = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    inner :: Int -> ST s ()
inner !Int
pos
                        | Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
e = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                        | Bool
otherwise = do
                            MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector s Int
MVector (PrimState (ST s)) Int
rtg (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
vis Int
pos) Int
g
                            Int -> ST s ()
inner (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> ST s ()
inner (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
g)
                Int -> ST s ()
go (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> ST s ()
go Int
0
    MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
rtg

data Idents a = Idents {forall a. Idents a -> a
minSeed :: !a, forall a. Idents a -> a
maxSeed :: !a}

intIdent :: Idents Int
intIdent :: Idents Int
intIdent = Int -> Int -> Idents Int
forall a. a -> a -> Idents a
Idents Int
forall a. Bounded a => a
maxBound Int
forall a. Bounded a => a
minBound

dblIdent :: Idents Double
dblIdent :: Idents Double
dblIdent = Double -> Double -> Idents Double
forall a. a -> a -> Idents a
Idents (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0) (Double -> Double
forall a. Num a => a -> a
negate (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0))

reduceParTyped ::
    forall a.
    (Columnable a, VU.Unbox a, Num a, Ord a, Real a) =>
    Reduction ->
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    VU.Vector a ->
    Idents a ->
    Column
reduceParTyped :: forall a.
(Columnable a, Unbox a, Num a, Ord a, Real a) =>
Reduction
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Idents a
-> Column
reduceParTyped Reduction
red Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Idents a
idents =
    let !caps :: Int
caps = Int
capabilities
        !bounds :: Vector Int
bounds = Vector Int -> Int -> Int -> Vector Int
groupRangeBounds Vector Int
offs Int
nGroups Int
caps
     in case Reduction
red of
            Reduction
RCount -> Vector Int -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (IO (Vector Int) -> Vector Int
forall a. IO a -> a
unsafePerformIO (Vector Int
-> Vector Int -> Int -> Int -> Vector Int -> IO (Vector Int)
countPar Vector Int
vis Vector Int
offs Int
nGroups Int
caps Vector Int
bounds))
            Reduction
RSum -> Vector a -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (IO (Vector a) -> Vector a
forall a. IO a -> a
unsafePerformIO (Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector a)
forall a.
(Unbox a, Num a) =>
Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector a)
sumPar Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds))
            Reduction
RMin ->
                Vector a -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector
                    (IO (Vector a) -> Vector a
forall a. IO a -> a
unsafePerformIO ((a -> a -> a)
-> a
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector a)
forall a.
Unbox a =>
(a -> a -> a)
-> a
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector a)
extremaPar a -> a -> a
forall a. Ord a => a -> a -> a
min (Idents a -> a
forall a. Idents a -> a
minSeed Idents a
idents) Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds))
            Reduction
RMax ->
                Vector a -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector
                    (IO (Vector a) -> Vector a
forall a. IO a -> a
unsafePerformIO ((a -> a -> a)
-> a
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector a)
forall a.
Unbox a =>
(a -> a -> a)
-> a
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector a)
extremaPar a -> a -> a
forall a. Ord a => a -> a -> a
max (Idents a -> a
forall a. Idents a -> a
maxSeed Idents a
idents) Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds))
            Reduction
RMean -> Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (IO (Vector Double) -> Vector Double
forall a. IO a -> a
unsafePerformIO (Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
forall a.
(Unbox a, Real a) =>
Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
meanPar Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds))
            Reduction
RVar ->
                Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector
                    (IO (Vector Double) -> Vector Double
forall a. IO a -> a
unsafePerformIO (Bool
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
forall a.
(Unbox a, Real a) =>
Bool
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
varPar Bool
False Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds))
            Reduction
RStd ->
                Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (IO (Vector Double) -> Vector Double
forall a. IO a -> a
unsafePerformIO (Bool
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
forall a.
(Unbox a, Real a) =>
Bool
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
varPar Bool
True Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds))
            Reduction
RTop2Sum -> Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (IO (Vector Double) -> Vector Double
forall a. IO a -> a
unsafePerformIO (Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
forall a.
(Unbox a, Real a) =>
Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
top2Par Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds))
{-# INLINE reduceParTyped #-}

-- | Iterate the rows of groups @[gs, ge)@ in @valueIndices@/group order.
overGroups ::
    VU.Vector Int -> VU.Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
overGroups :: Vector Int
-> Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
overGroups Vector Int
vis Vector Int
offs Int
gs Int
ge Int -> Int -> IO ()
step = Int -> IO ()
grp Int
gs
  where
    grp :: Int -> IO ()
grp !Int
g
        | Int
g Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
ge = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            let !e :: Int
e = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                inner :: Int -> IO ()
inner !Int
pos
                    | Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
e = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                    | Bool
otherwise = Int -> Int -> IO ()
step Int
g (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
vis Int
pos) IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> IO ()
inner (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Int -> IO ()
inner (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
g)
            Int -> IO ()
grp (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE overGroups #-}

countPar ::
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    Int ->
    VU.Vector Int ->
    IO (VU.Vector Int)
countPar :: Vector Int
-> Vector Int -> Int -> Int -> Vector Int -> IO (Vector Int)
countPar Vector Int
_vis Vector Int
offs Int
nGroups Int
caps Vector Int
bounds = do
    MVector RealWorld Int
out <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Int
0 :: Int)
    Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange Vector Int
bounds Int
caps ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
gs Int
ge ->
        let grp :: Int -> IO ()
grp !Int
g
                | Int
g Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
ge = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                | Bool
otherwise = do
                    let !c :: Int
c = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
offs Int
g
                    MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Int
MVector (PrimState IO) Int
out Int
g Int
c
                    Int -> IO ()
grp (Int
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
         in Int -> IO ()
grp Int
gs
    MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Int
MVector (PrimState IO) Int
out

sumPar ::
    (VU.Unbox a, Num a) =>
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    VU.Vector a ->
    Int ->
    VU.Vector Int ->
    IO (VU.Vector a)
sumPar :: forall a.
(Unbox a, Num a) =>
Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector a)
sumPar Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds = do
    MVector RealWorld a
out <- Int -> a -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups a
0
    Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange Vector Int
bounds Int
caps ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
gs Int
ge ->
        Vector Int
-> Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
overGroups Vector Int
vis Vector Int
offs Int
gs Int
ge ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
g Int
row -> do
            a
cur <- MVector (PrimState IO) a -> Int -> IO a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld a
MVector (PrimState IO) a
out Int
g
            MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld a
MVector (PrimState IO) a
out Int
g (a
cur a -> a -> a
forall a. Num a => a -> a -> a
+ Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
row)
    MVector (PrimState IO) a -> IO (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld a
MVector (PrimState IO) a
out
{-# INLINE sumPar #-}

extremaPar ::
    (VU.Unbox a) =>
    (a -> a -> a) ->
    a ->
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    VU.Vector a ->
    Int ->
    VU.Vector Int ->
    IO (VU.Vector a)
extremaPar :: forall a.
Unbox a =>
(a -> a -> a)
-> a
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector a)
extremaPar a -> a -> a
combine a
seed Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds = do
    MVector RealWorld a
out <- Int -> a -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups a
seed
    Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange Vector Int
bounds Int
caps ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
gs Int
ge ->
        Vector Int
-> Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
overGroups Vector Int
vis Vector Int
offs Int
gs Int
ge ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
g Int
row -> do
            a
cur <- MVector (PrimState IO) a -> Int -> IO a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld a
MVector (PrimState IO) a
out Int
g
            MVector (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld a
MVector (PrimState IO) a
out Int
g (a -> a -> a
combine a
cur (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
row))
    MVector (PrimState IO) a -> IO (Vector a)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld a
MVector (PrimState IO) a
out
{-# INLINE extremaPar #-}

meanPar ::
    (VU.Unbox a, Real a) =>
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    VU.Vector a ->
    Int ->
    VU.Vector Int ->
    IO (VU.Vector Double)
meanPar :: forall a.
(Unbox a, Real a) =>
Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
meanPar Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds = do
    MVector RealWorld Double
s <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Double
0 :: Double)
    MVector RealWorld Int
cnt <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Int
0 :: Int)
    Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange Vector Int
bounds Int
caps ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
gs Int
ge ->
        Vector Int
-> Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
overGroups Vector Int
vis Vector Int
offs Int
gs Int
ge ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
g Int
row -> do
            let !x :: Double
x = a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
row)
            Double
cs <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
s Int
g
            MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
s Int
g (Double
cs Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x)
            Int
cc <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Int
MVector (PrimState IO) Int
cnt Int
g
            MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Int
MVector (PrimState IO) Int
cnt Int
g (Int
cc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    MVector RealWorld Double
out <- Int -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
nGroups
    let fin :: Int -> IO ()
fin !Int
k
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                Double
sv <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
s Int
k
                Int
c <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Int
MVector (PrimState IO) Int
cnt Int
k
                MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
out Int
k (if Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Double
sv Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c)
                Int -> IO ()
fin (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> IO ()
fin Int
0
    MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
out
{-# INLINE meanPar #-}

varPar ::
    (VU.Unbox a, Real a) =>
    Bool ->
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    VU.Vector a ->
    Int ->
    VU.Vector Int ->
    IO (VU.Vector Double)
varPar :: forall a.
(Unbox a, Real a) =>
Bool
-> Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
varPar Bool
takeSqrt Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds = do
    MVector RealWorld Int
cnt <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Int
0 :: Int)
    MVector RealWorld Double
meanV <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Double
0 :: Double)
    MVector RealWorld Double
m2 <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Double
0 :: Double)
    Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange Vector Int
bounds Int
caps ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
gs Int
ge ->
        Vector Int
-> Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
overGroups Vector Int
vis Vector Int
offs Int
gs Int
ge ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
g Int
row -> do
            let !x :: Double
x = a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
row)
            Int
c <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Int
MVector (PrimState IO) Int
cnt Int
g
            Double
mu <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
meanV Int
g
            Double
mm <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
m2 Int
g
            let !c' :: Int
c' = Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                !delta :: Double
delta = Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu
                !mu' :: Double
mu' = Double
mu Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
delta Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c'
                !mm' :: Double
mm' = Double
mm Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
delta Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu')
            MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Int
MVector (PrimState IO) Int
cnt Int
g Int
c'
            MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
meanV Int
g Double
mu'
            MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
m2 Int
g Double
mm'
    MVector RealWorld Double
out <- Int -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
nGroups
    let fin :: Int -> IO ()
fin !Int
k
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                Int
c <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Int
MVector (PrimState IO) Int
cnt Int
k
                Double
mm <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
m2 Int
k
                let var :: Double
var = if Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 then Double
0 else Double
mm Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
out Int
k (if Bool
takeSqrt then Double -> Double
forall a. Floating a => a -> a
sqrt Double
var else Double
var)
                Int -> IO ()
fin (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> IO ()
fin Int
0
    MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
out
{-# INLINE varPar #-}

top2Par ::
    (VU.Unbox a, Real a) =>
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    VU.Vector a ->
    Int ->
    VU.Vector Int ->
    IO (VU.Vector Double)
top2Par :: forall a.
(Unbox a, Real a) =>
Vector Int
-> Vector Int
-> Int
-> Vector a
-> Int
-> Vector Int
-> IO (Vector Double)
top2Par Vector Int
vis Vector Int
offs Int
nGroups Vector a
v Int
caps Vector Int
bounds = do
    let ninf :: Double
ninf = Double -> Double
forall a. Num a => a -> a
negate (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0) :: Double
    MVector RealWorld Double
m1 <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups Double
ninf
    MVector RealWorld Double
m2 <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups Double
ninf
    Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange Vector Int
bounds Int
caps ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
gs Int
ge ->
        Vector Int
-> Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
overGroups Vector Int
vis Vector Int
offs Int
gs Int
ge ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
g Int
row -> do
            let !x :: Double
x = a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Vector a -> Int -> a
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector a
v Int
row)
            Double
a1 <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
m1 Int
g
            if Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
a1
                then do
                    MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
m1 Int
g Double
x
                    MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
m2 Int
g Double
a1
                else do
                    Double
a2 <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
m2 Int
g
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
a2) (MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
m2 Int
g Double
x)
    MVector RealWorld Double
out <- Int -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
nGroups
    let fin :: Int -> IO ()
fin !Int
k
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                Double
a1 <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
m1 Int
k
                Double
a2 <- MVector (PrimState IO) Double -> Int -> IO Double
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Double
MVector (PrimState IO) Double
m2 Int
k
                let sm :: Double
sm = (if Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
a1 then Double
0 else Double
a1) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (if Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
a2 then Double
0 else Double
a2)
                MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Double
MVector (PrimState IO) Double
out Int
k Double
sm
                Int -> IO ()
fin (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> IO ()
fin Int
0
    MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
out
{-# INLINE top2Par #-}

-------------------------------------------------------------------------------
-- Parallel fused two-column moments (Q9)
-------------------------------------------------------------------------------

{- | Parallel counterpart of 'momentScatter': one fused pass over both columns,
each group's six sums accumulated within one worker's range. Byte-identical to
'momentScatter'. 'Nothing' unless both columns are non-null unboxed Int/Double.
-}
momentScatterPar ::
    VU.Vector Int -> VU.Vector Int -> Int -> Column -> Column -> Maybe Moments
momentScatterPar :: Vector Int
-> Vector Int -> Int -> Column -> Column -> Maybe Moments
momentScatterPar Vector Int
vis Vector Int
offs Int
nGroups Column
colX Column
colY
    | Bool -> Bool
not (Int -> Bool
shouldPar (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
vis)) Bool -> Bool -> Bool
|| Int
nGroups Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 =
        Vector Int -> Int -> Column -> Column -> Maybe Moments
momentScatter (Vector Int -> Vector Int -> Int -> Vector Int
rtgFromVis Vector Int
vis Vector Int
offs Int
nGroups) Int
nGroups Column
colX Column
colY
    | Bool
otherwise = do
        Vector Double
xs <- Column -> Maybe (Vector Double)
scatterColumnToDouble Column
colX
        Vector Double
ys <- Column -> Maybe (Vector Double)
scatterColumnToDouble Column
colY
        let !caps :: Int
caps = Int
capabilities
            !bounds :: Vector Int
bounds = Vector Int -> Int -> Int -> Vector Int
groupRangeBounds Vector Int
offs Int
nGroups Int
caps
        Moments -> Maybe Moments
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO Moments -> Moments
forall a. IO a -> a
unsafePerformIO (Vector Int
-> Vector Int
-> Int
-> Vector Double
-> Vector Double
-> Int
-> Vector Int
-> IO Moments
momentPar Vector Int
vis Vector Int
offs Int
nGroups Vector Double
xs Vector Double
ys Int
caps Vector Int
bounds))
{-# NOINLINE momentScatterPar #-}

momentPar ::
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    VU.Vector Double ->
    VU.Vector Double ->
    Int ->
    VU.Vector Int ->
    IO Moments
momentPar :: Vector Int
-> Vector Int
-> Int
-> Vector Double
-> Vector Double
-> Int
-> Vector Int
-> IO Moments
momentPar Vector Int
vis Vector Int
offs Int
nGroups Vector Double
xs Vector Double
ys Int
caps Vector Int
bounds = do
    MVector RealWorld Int
cnt <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Int
0 :: Int)
    MVector RealWorld Double
sx <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Double
0 :: Double)
    MVector RealWorld Double
sy <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Double
0 :: Double)
    MVector RealWorld Double
sxx <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Double
0 :: Double)
    MVector RealWorld Double
syy <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Double
0 :: Double)
    MVector RealWorld Double
sxy <- Int -> Double -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Double
0 :: Double)
    let bump :: MVector (PrimState m) a -> Int -> a -> m ()
bump MVector (PrimState m) a
arr Int
g a
d = MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector (PrimState m) a
arr Int
g m a -> (a -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
c -> MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector (PrimState m) a
arr Int
g (a
c a -> a -> a
forall a. Num a => a -> a -> a
+ a
d)
    Vector Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forEachRange Vector Int
bounds Int
caps ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
gs Int
ge ->
        Vector Int
-> Vector Int -> Int -> Int -> (Int -> Int -> IO ()) -> IO ()
overGroups Vector Int
vis Vector Int
offs Int
gs Int
ge ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
g Int
row -> do
            let !x :: Double
x = Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Double
xs Int
row
                !y :: Double
y = Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Double
ys Int
row
            MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead MVector RealWorld Int
MVector (PrimState IO) Int
cnt Int
g IO Int -> (Int -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Int
c -> MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite MVector RealWorld Int
MVector (PrimState IO) Int
cnt Int
g (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall {m :: * -> *} {a}.
(PrimMonad m, Unbox a, Num a) =>
MVector (PrimState m) a -> Int -> a -> m ()
bump MVector RealWorld Double
MVector (PrimState IO) Double
sx Int
g Double
x
            MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall {m :: * -> *} {a}.
(PrimMonad m, Unbox a, Num a) =>
MVector (PrimState m) a -> Int -> a -> m ()
bump MVector RealWorld Double
MVector (PrimState IO) Double
sy Int
g Double
y
            MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall {m :: * -> *} {a}.
(PrimMonad m, Unbox a, Num a) =>
MVector (PrimState m) a -> Int -> a -> m ()
bump MVector RealWorld Double
MVector (PrimState IO) Double
sxx Int
g (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x)
            MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall {m :: * -> *} {a}.
(PrimMonad m, Unbox a, Num a) =>
MVector (PrimState m) a -> Int -> a -> m ()
bump MVector RealWorld Double
MVector (PrimState IO) Double
syy Int
g (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y)
            MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall {m :: * -> *} {a}.
(PrimMonad m, Unbox a, Num a) =>
MVector (PrimState m) a -> Int -> a -> m ()
bump MVector RealWorld Double
MVector (PrimState IO) Double
sxy Int
g (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y)
    Column -> Column -> Column -> Column -> Column -> Column -> Moments
Moments (Column
 -> Column -> Column -> Column -> Column -> Column -> Moments)
-> (Vector Int -> Column)
-> Vector Int
-> Column
-> Column
-> Column
-> Column
-> Column
-> Moments
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Int -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector
        (Vector Int
 -> Column -> Column -> Column -> Column -> Column -> Moments)
-> IO (Vector Int)
-> IO (Column -> Column -> Column -> Column -> Column -> Moments)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Int
MVector (PrimState IO) Int
cnt
        IO (Column -> Column -> Column -> Column -> Column -> Moments)
-> IO Column
-> IO (Column -> Column -> Column -> Column -> Moments)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Double -> Column) -> IO (Vector Double) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
sx)
        IO (Column -> Column -> Column -> Column -> Moments)
-> IO Column -> IO (Column -> Column -> Column -> Moments)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Double -> Column) -> IO (Vector Double) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
sy)
        IO (Column -> Column -> Column -> Moments)
-> IO Column -> IO (Column -> Column -> Moments)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Double -> Column) -> IO (Vector Double) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
sxx)
        IO (Column -> Column -> Moments)
-> IO Column -> IO (Column -> Moments)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Double -> Column) -> IO (Vector Double) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
syy)
        IO (Column -> Moments) -> IO Column -> IO Moments
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Double -> Column) -> IO (Vector Double) -> IO Column
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
sxy)