{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}

{- |
Module      : Granite.Stat
Copyright   : (c) 2025
License     : MIT
Maintainer  : mschavinda@gmail.com

Statistical transforms ('Stat'). Each runs on a layer's mapped X / Y
columns and returns a new frame, generally with the same column names
holding the transformed values. 'Granite.Spec.StatBoxplot' additionally writes
@__ymin@ / @__q1@ / @__median@ / @__q3@ / @__ymax@ for the boxplot
geom to read.
-}
module Granite.Stat (
    applyStat,
    binData,
    kdeData,
    smoothLm,
    smoothMovingAvg,
    smoothLoess,
    boxplotSummary,
    BoxStats (..),
    countByCategory,
    summarizeByCategory,
) where

import Data.List qualified as List
import Data.Text (Text)

import Granite.Data.Frame (
    Column (..),
    DataFrame (..),
    columnAsNum,
    columnAsText,
    fromColumns,
    lookupColumn,
 )
import Granite.Spec (
    BinSpec (..),
    ColumnRef (..),
    Mapping (..),
    SmoothMethod (..),
    Stat (..),
    SummaryFun (..),
 )
import Data.Maybe (listToMaybe, fromMaybe)

applyStat :: Stat -> Mapping -> DataFrame -> DataFrame
applyStat :: Stat -> Mapping -> DataFrame -> DataFrame
applyStat Stat
stat Mapping
m DataFrame
df = case Stat
stat of
    Stat
StatIdentity -> DataFrame
df
    StatBin BinSpec
spec -> BinSpec -> Mapping -> DataFrame -> DataFrame
runBin BinSpec
spec Mapping
m DataFrame
df
    Stat
StatDensity -> Mapping -> DataFrame -> DataFrame
runDensity Mapping
m DataFrame
df
    StatSmooth SmoothMethod
method -> SmoothMethod -> Mapping -> DataFrame -> DataFrame
runSmooth SmoothMethod
method Mapping
m DataFrame
df
    Stat
StatBoxplot -> Mapping -> DataFrame -> DataFrame
runBoxplot Mapping
m DataFrame
df
    Stat
StatCount -> Mapping -> DataFrame -> DataFrame
runCount Mapping
m DataFrame
df
    StatSummary SummaryFun
fn -> SummaryFun -> Mapping -> DataFrame -> DataFrame
runSummary SummaryFun
fn Mapping
m DataFrame
df

runBin :: BinSpec -> Mapping -> DataFrame -> DataFrame
runBin :: BinSpec -> Mapping -> DataFrame -> DataFrame
runBin BinSpec
spec Mapping
m DataFrame
df =
    case Maybe ColumnRef -> DataFrame -> Maybe [Double]
resolveNum (Mapping -> Maybe ColumnRef
aesX Mapping
m) DataFrame
df of
        Maybe [Double]
Nothing -> DataFrame
df
        Just [Double]
xs ->
            let ([Double]
mids, [Double]
counts) = BinSpec -> [Double] -> ([Double], [Double])
binData BinSpec
spec [Double]
xs
                xName :: Text
xName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesX Mapping
m) Text
"x"
                yName :: Text
yName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesY Mapping
m) Text
"count"
             in [(Text, Column)] -> DataFrame
fromColumns
                    [ (Text
xName, [Double] -> Column
ColNum [Double]
mids)
                    , (Text
yName, [Double] -> Column
ColNum [Double]
counts)
                    ]

binData :: BinSpec -> [Double] -> ([Double], [Double])
binData :: BinSpec -> [Double] -> ([Double], [Double])
binData BinSpec
_ [] = ([], [])
binData BinSpec
spec [Double]
xs =
    let lo :: Double
lo = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Double]
xs
        hi :: Double
hi = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
xs
        edges :: [Double]
edges = case BinSpec
spec of
            BinByCount Int
n -> Int -> Double -> Double -> [Double]
evenEdges Int
n Double
lo Double
hi
            BinByWidth Double
w
                | Double
w Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 -> Int -> Double -> Double -> [Double]
evenEdges Int
1 Double
lo Double
hi
                | Bool
otherwise -> Double -> Double -> Double -> [Double]
stepEdges Double
w Double
lo Double
hi
            BinByEdges [Double]
es -> [Double] -> [Double]
forall a. Ord a => [a] -> [a]
List.sort [Double]
es
        nBins :: Int
nBins = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
edges Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        mids :: [Double]
mids =
            [ (Double
e0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
e1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
            | (Double
e0, Double
e1) <- [Double] -> [Double] -> [(Double, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
edges (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 [Double]
edges)
            ]
        counts0 :: [Int]
counts0 = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
nBins (Int
0 :: Int)
        counts :: [Int]
counts =
            (Double -> [Int] -> [Int]) -> [Int] -> [Double] -> [Int]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
                ( \Double
v [Int]
acc ->
                    let ix :: Int
ix = [Double] -> Double -> Int
which [Double]
edges Double
v
                     in [Int] -> Int -> [Int]
update [Int]
acc Int
ix
                )
                [Int]
counts0
                [Double]
xs
     in ([Double]
mids, (Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int]
counts)
  where
    update :: [Int] -> Int -> [Int]
    update :: [Int] -> Int -> [Int]
update [Int]
cs Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
cs = [Int]
cs
        | Bool
otherwise = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
i [Int]
cs [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [[Int]
cs [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
cs

    -- Epsilon tolerance so values landing exactly on an internal edge
    -- go to the right-side bin, matching @[lo, hi)@ even after the
    -- rounding error in @lo + step * i@.
    which :: [Double] -> Double -> Int
    which :: [Double] -> Double -> Int
which [Double]
es Double
v
        | [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
es Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = -Int
1
        | Bool -> (Double -> Bool) -> Maybe Double -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps <) ([Double] -> Maybe Double
forall a. [a] -> Maybe a
listToMaybe [Double]
es) = -Int
1
        | Double
v Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> [Double] -> Double
forall a. HasCallStack => [a] -> a
last [Double]
es Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps = -Int
1
        | Bool
otherwise = Int -> Int
go Int
0
      where
        eps :: Double
eps = Double
1e-9
        n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
es Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        go :: Int -> Int
go Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
            | Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< [Double]
es [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) = Int
i
            | 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 Bool -> Bool -> Bool
&& Double
v Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= [Double]
es [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps = Int
i
            | Bool
otherwise = Int -> Int
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

evenEdges :: Int -> Double -> Double -> [Double]
evenEdges :: Int -> Double -> Double -> [Double]
evenEdges Int
n Double
lo Double
hi
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = [Double
lo, Double
hi]
    | Double
lo Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
hi = [Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
0.5, Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
0.5]
    | Bool
otherwise =
        let step :: Double
step = (Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
         in [Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
step Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i | Int
i <- [Int
0 .. Int
n]]

stepEdges :: Double -> Double -> Double -> [Double]
stepEdges :: Double -> Double -> Double -> [Double]
stepEdges Double
w Double
lo Double
hi
    | Double
lo Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
hi = [Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
w Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2, Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
w Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2]
    | Bool
otherwise =
        let n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling ((Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
w)) :: Int
         in [Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i | Int
i <- [Int
0 .. Int
n]]

runDensity :: Mapping -> DataFrame -> DataFrame
runDensity :: Mapping -> DataFrame -> DataFrame
runDensity Mapping
m DataFrame
df =
    case Maybe ColumnRef -> DataFrame -> Maybe [Double]
resolveNum (Mapping -> Maybe ColumnRef
aesX Mapping
m) DataFrame
df of
        Maybe [Double]
Nothing -> DataFrame
df
        Just [Double]
xs ->
            let ([Double]
gx, [Double]
gy) = Int -> [Double] -> ([Double], [Double])
kdeData Int
128 [Double]
xs
                xName :: Text
xName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesX Mapping
m) Text
"x"
                yName :: Text
yName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesY Mapping
m) Text
"density"
             in [(Text, Column)] -> DataFrame
fromColumns [(Text
xName, [Double] -> Column
ColNum [Double]
gx), (Text
yName, [Double] -> Column
ColNum [Double]
gy)]

-- | Gaussian KDE on @n@ grid points; bandwidth is Silverman's rule.
kdeData :: Int -> [Double] -> ([Double], [Double])
kdeData :: Int -> [Double] -> ([Double], [Double])
kdeData Int
_ [] = ([], [])
kdeData Int
n [Double]
xs =
    let nn :: Int
nn = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs
        mu :: Double
mu = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nn
        var :: Double
var = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Double
x <- [Double]
xs] Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int
nn Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
        sigma :: Double
sigma = Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
var Double
1e-12)
        h :: Double
h = Double
1.06 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sigma Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nn Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (-Double
0.2)
        lo :: Double
lo = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Double]
xs Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h
        hi :: Double
hi = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
xs Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
3 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h
        grid :: [Double]
grid = Int -> Double -> Double -> [Double]
evenSpaced Int
n Double
lo Double
hi
        density :: Double -> Double
density Double
g =
            let k :: Double -> Double
k Double
v = Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double
forall a. Num a => a -> a
negate ((Double
g Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
v) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
2) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
h)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Double
h Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi))
             in [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
k [Double]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nn
     in ([Double]
grid, (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
density [Double]
grid)

evenSpaced :: Int -> Double -> Double -> [Double]
evenSpaced :: Int -> Double -> Double -> [Double]
evenSpaced Int
n Double
lo Double
hi
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = [Double
lo]
    | Bool
otherwise =
        let step :: Double
step = (Double
hi Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lo) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
         in [Double
lo Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
step Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

runSmooth :: SmoothMethod -> Mapping -> DataFrame -> DataFrame
runSmooth :: SmoothMethod -> Mapping -> DataFrame -> DataFrame
runSmooth SmoothMethod
method Mapping
m DataFrame
df =
    case ( Maybe ColumnRef -> DataFrame -> Maybe [Double]
resolveNum (Mapping -> Maybe ColumnRef
aesX Mapping
m) DataFrame
df
         , Maybe ColumnRef -> DataFrame -> Maybe [Double]
resolveNum (Mapping -> Maybe ColumnRef
aesY Mapping
m) DataFrame
df
         ) of
        (Just [Double]
xs, Just [Double]
ys) ->
            let pairs :: [(Double, Double)]
pairs = ((Double, Double) -> Double)
-> [(Double, Double)] -> [(Double, Double)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
List.sortOn (Double, Double) -> Double
forall a b. (a, b) -> a
fst ([Double] -> [Double] -> [(Double, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
xs [Double]
ys)
                ([Double]
sx, [Double]
sy) = [(Double, Double)] -> ([Double], [Double])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Double, Double)]
pairs
                fitted :: [Double]
fitted = case SmoothMethod
method of
                    SmoothMethod
SmoothLm -> [Double] -> [Double] -> [Double]
smoothLm [Double]
sx [Double]
sy
                    SmoothLoess Double
span_ -> Double -> [Double] -> [Double] -> [Double]
smoothLoess Double
span_ [Double]
sx [Double]
sy
                    SmoothMovingAvg Int
w -> Int -> [Double] -> [Double]
smoothMovingAvg Int
w [Double]
sy
                xName :: Text
xName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesX Mapping
m) Text
"x"
                yName :: Text
yName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesY Mapping
m) Text
"y"
             in [(Text, Column)] -> DataFrame
fromColumns [(Text
xName, [Double] -> Column
ColNum [Double]
sx), (Text
yName, [Double] -> Column
ColNum [Double]
fitted)]
        (Maybe [Double], Maybe [Double])
_ -> DataFrame
df

-- | OLS regression, evaluated at the input X positions.
smoothLm :: [Double] -> [Double] -> [Double]
smoothLm :: [Double] -> [Double] -> [Double]
smoothLm [Double]
xs [Double]
ys
    | [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = [Double]
ys
    | Bool
otherwise =
        let n :: Double
n = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs) :: Double
            mx :: Double
mx = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
xs Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n
            my :: Double
my = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
ys Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n
            num :: Double
num = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
x Double
y -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mx) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
my)) [Double]
xs [Double]
ys)
            den :: Double
den = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mx) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
2 | Double
x <- [Double]
xs]
            slope :: Double
slope = if Double
den Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
num Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
den
            intercept :: Double
intercept = Double
my Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
slope Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
mx
         in [Double
intercept Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
slope Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x | Double
x <- [Double]
xs]

{- | Trailing moving average; leading edge is the partial-window mean
so the output length matches the input.
-}
smoothMovingAvg :: Int -> [Double] -> [Double]
smoothMovingAvg :: Int -> [Double] -> [Double]
smoothMovingAvg Int
w [Double]
ys
    | Int
w Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = [Double]
ys
    | Bool
otherwise =
        [ let lo :: Int
lo = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              window :: [Double]
window = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
lo [Double]
ys)
              k :: Double
k = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
window) :: Double
           in if Double
k Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
window Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
k
        | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        ]

{- | LOESS-style local linear regression with tricube weights;
@span_@ is the neighbourhood fraction in @(0, 1]@.
-}
smoothLoess :: Double -> [Double] -> [Double] -> [Double]
smoothLoess :: Double -> [Double] -> [Double] -> [Double]
smoothLoess Double
span_ [Double]
xs [Double]
ys
    | [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = [Double]
ys
    | Bool
otherwise =
        let n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs
            k :: Int
k = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
2 (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double
span_ Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n))
         in [[Double] -> [Double] -> Int -> Double -> Double
loessAt [Double]
xs [Double]
ys Int
k Double
x | Double
x <- [Double]
xs]

loessAt :: [Double] -> [Double] -> Int -> Double -> Double
loessAt :: [Double] -> [Double] -> Int -> Double -> Double
loessAt [Double]
xs [Double]
ys Int
k Double
x0 =
    let dists :: [(Double, Double, Double)]
dists = [(Double -> Double
forall a. Num a => a -> a
abs (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
x0), Double
x, Double
y) | (Double
x, Double
y) <- [Double] -> [Double] -> [(Double, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
xs [Double]
ys]
        nearest :: [(Double, Double, Double)]
nearest = Int -> [(Double, Double, Double)] -> [(Double, Double, Double)]
forall a. Int -> [a] -> [a]
take Int
k (((Double, Double, Double) -> Double)
-> [(Double, Double, Double)] -> [(Double, Double, Double)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
List.sortOn (\(Double
d, Double
_, Double
_) -> Double
d) [(Double, Double, Double)]
dists)
        maxD :: Double
maxD = case [(Double, Double, Double)]
nearest of
            [] -> Double
1
            [(Double, Double, Double)]
ds -> [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double
d | (Double
d, Double
_, Double
_) <- [(Double, Double, Double)]
ds] Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1e-12
        wts :: [Double]
wts = [Double -> Double
tricube (Double
d Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
maxD) | (Double
d, Double
_, Double
_) <- [(Double, Double, Double)]
nearest]
        xv :: [Double]
xv = [Double
x | (Double
_, Double
x, Double
_) <- [(Double, Double, Double)]
nearest]
        yv :: [Double]
yv = [Double
y | (Double
_, Double
_, Double
y) <- [(Double, Double, Double)]
nearest]
        sw :: Double
sw = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
wts
        swx :: Double
swx = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
wts [Double]
xv)
        swy :: Double
swy = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
wts [Double]
yv)
        swxx :: Double
swxx = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
w Double
x -> Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x) [Double]
wts [Double]
xv)
        swxy :: Double
swxy = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Double -> Double -> Double -> Double)
-> [Double] -> [Double] -> [Double] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Double
w Double
x Double
y -> Double
w Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y) [Double]
wts [Double]
xv [Double]
yv)
        denom :: Double
denom = Double
sw Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
swxx Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
swx Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
swx
        slope :: Double
slope = if Double
denom Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else (Double
sw Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
swxy Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
swx Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
swy) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
denom
        intercept :: Double
intercept = if Double
sw Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else (Double
swy Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
slope Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
swx) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sw
     in Double
intercept Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
slope Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x0

tricube :: Double -> Double
tricube :: Double -> Double
tricube Double
u
    | Double -> Double
forall a. Num a => a -> a
abs Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
1 = Double
0
    | Bool
otherwise = (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Num a => a -> a
abs Double
u Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
3) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
3

runBoxplot :: Mapping -> DataFrame -> DataFrame
runBoxplot :: Mapping -> DataFrame -> DataFrame
runBoxplot Mapping
m DataFrame
df =
    case Maybe ColumnRef -> DataFrame -> Maybe [Double]
resolveNum (Mapping -> Maybe ColumnRef
aesY Mapping
m) DataFrame
df of
        Maybe [Double]
Nothing -> DataFrame
df
        Just [Double]
ys ->
            let groupNames :: [Text]
groupNames = case Mapping -> Maybe ColumnRef
aesX Mapping
m of
                    Just (ColumnRef Text
n) -> case Text -> DataFrame -> Maybe Column
lookupColumn Text
n DataFrame
df of
                        Just Column
c -> Column -> [Text]
columnAsText Column
c
                        Maybe Column
Nothing -> Int -> Text -> [Text]
forall a. Int -> a -> [a]
replicate ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ys) Text
"all"
                    Maybe ColumnRef
Nothing -> Int -> Text -> [Text]
forall a. Int -> a -> [a]
replicate ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ys) Text
"all"
                grouped :: [(Text, [(Text, Double)])]
grouped = ((Text, Double) -> Text)
-> [(Text, Double)] -> [(Text, [(Text, Double)])]
forall k a. Eq k => (a -> k) -> [a] -> [(k, [a])]
groupBy (Text, Double) -> Text
forall a b. (a, b) -> a
fst ([Text] -> [Double] -> [(Text, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
groupNames [Double]
ys)
                summaries :: [(Text, BoxStats)]
summaries =
                    [ (Text
name, [Double] -> BoxStats
boxplotSummary [Double
v | (Text
_, Double
v) <- [(Text, Double)]
xs])
                    | (Text
name, [(Text, Double)]
xs) <- [(Text, [(Text, Double)])]
grouped
                    ]
                xName :: Text
xName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesX Mapping
m) Text
"group"
             in [(Text, Column)] -> DataFrame
fromColumns
                    [ (Text
xName, [Text] -> Column
ColCat [Text
n | (Text
n, BoxStats
_) <- [(Text, BoxStats)]
summaries])
                    , (Text
"__ymin", [Double] -> Column
ColNum [BoxStats -> Double
yMin BoxStats
s | (Text
_, BoxStats
s) <- [(Text, BoxStats)]
summaries])
                    , (Text
"__q1", [Double] -> Column
ColNum [BoxStats -> Double
q1 BoxStats
s | (Text
_, BoxStats
s) <- [(Text, BoxStats)]
summaries])
                    , (Text
"__median", [Double] -> Column
ColNum [BoxStats -> Double
med BoxStats
s | (Text
_, BoxStats
s) <- [(Text, BoxStats)]
summaries])
                    , (Text
"__q3", [Double] -> Column
ColNum [BoxStats -> Double
q3 BoxStats
s | (Text
_, BoxStats
s) <- [(Text, BoxStats)]
summaries])
                    , (Text
"__ymax", [Double] -> Column
ColNum [BoxStats -> Double
yMax BoxStats
s | (Text
_, BoxStats
s) <- [(Text, BoxStats)]
summaries])
                    ]

data BoxStats = BoxStats
    { BoxStats -> Double
yMin :: !Double
    , BoxStats -> Double
q1 :: !Double
    , BoxStats -> Double
med :: !Double
    , BoxStats -> Double
q3 :: !Double
    , BoxStats -> Double
yMax :: !Double
    }
    deriving (BoxStats -> BoxStats -> Bool
(BoxStats -> BoxStats -> Bool)
-> (BoxStats -> BoxStats -> Bool) -> Eq BoxStats
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BoxStats -> BoxStats -> Bool
== :: BoxStats -> BoxStats -> Bool
$c/= :: BoxStats -> BoxStats -> Bool
/= :: BoxStats -> BoxStats -> Bool
Eq, Int -> BoxStats -> ShowS
[BoxStats] -> ShowS
BoxStats -> String
(Int -> BoxStats -> ShowS)
-> (BoxStats -> String) -> ([BoxStats] -> ShowS) -> Show BoxStats
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BoxStats -> ShowS
showsPrec :: Int -> BoxStats -> ShowS
$cshow :: BoxStats -> String
show :: BoxStats -> String
$cshowList :: [BoxStats] -> ShowS
showList :: [BoxStats] -> ShowS
Show)

{- | Five-number summary using Tukey hinges: when @n@ is odd the
median is included in both halves, so @boxplotSummary [1..9]@
gives @(1, 3, 5, 7, 9)@.
-}
boxplotSummary :: [Double] -> BoxStats
boxplotSummary :: [Double] -> BoxStats
boxplotSummary [] = Double -> Double -> Double -> Double -> Double -> BoxStats
BoxStats Double
0 Double
0 Double
0 Double
0 Double
0
boxplotSummary [Double]
xs =
    let sorted :: [Double]
sorted = [Double] -> [Double]
forall a. Ord a => [a] -> [a]
List.sort [Double]
xs
        n :: Int
n = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
sorted
        m :: Double
m =
            if Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                then [Double]
sorted [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
                else ([Double]
sorted [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
sorted [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
        half :: Int
half =
            if Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                then (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
                else Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
        lower :: [Double]
lower = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
half [Double]
sorted
        upper :: [Double]
upper = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
half) [Double]
sorted
        qq :: [Double] -> Double
qq [Double]
xs' =
            if [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
xs'
                then Double
m
                else
                    let k :: Int
k = [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs'
                     in if Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                            then [Double]
xs' [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
                            else ([Double]
xs' [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ [Double]
xs' [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2
     in BoxStats
            { yMin :: Double
yMin = Double -> Maybe Double -> Double
forall a. a -> Maybe a -> a
fromMaybe Double
0 ([Double] -> Maybe Double
forall a. [a] -> Maybe a
listToMaybe [Double]
sorted)
            , q1 :: Double
q1 = [Double] -> Double
qq [Double]
lower
            , med :: Double
med = Double
m
            , q3 :: Double
q3 = [Double] -> Double
qq [Double]
upper
            , yMax :: Double
yMax = Double -> Maybe Double -> Double
forall a. a -> Maybe a -> a
fromMaybe Double
0 ([Double] -> Maybe Double
forall a. [a] -> Maybe a
listToMaybe ([Double] -> [Double]
forall a. [a] -> [a]
reverse [Double]
sorted))
            }

runCount :: Mapping -> DataFrame -> DataFrame
runCount :: Mapping -> DataFrame -> DataFrame
runCount Mapping
m DataFrame
df =
    let xs :: [Text]
xs = case Mapping -> Maybe ColumnRef
aesX Mapping
m of
            Just (ColumnRef Text
n) -> [Text] -> (Column -> [Text]) -> Maybe Column -> [Text]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] Column -> [Text]
columnAsText (Text -> DataFrame -> Maybe Column
lookupColumn Text
n DataFrame
df)
            Maybe ColumnRef
Nothing -> []
        counts :: [(Text, Int)]
counts = [Text] -> [(Text, Int)]
countByCategory [Text]
xs
        xName :: Text
xName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesX Mapping
m) Text
"x"
        yName :: Text
yName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesY Mapping
m) Text
"count"
     in [(Text, Column)] -> DataFrame
fromColumns
            [ (Text
xName, [Text] -> Column
ColCat [Text
k | (Text
k, Int
_) <- [(Text, Int)]
counts])
            , (Text
yName, [Double] -> Column
ColNum [Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
v | (Text
_, Int
v) <- [(Text, Int)]
counts])
            ]

countByCategory :: [Text] -> [(Text, Int)]
countByCategory :: [Text] -> [(Text, Int)]
countByCategory [Text]
xs =
    let uniq :: [Text]
uniq = [Text] -> [Text]
forall a. Eq a => [a] -> [a]
List.nub [Text]
xs
     in [(Text
u, [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text
x | Text
x <- [Text]
xs, Text
x Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
u]) | Text
u <- [Text]
uniq]

runSummary :: SummaryFun -> Mapping -> DataFrame -> DataFrame
runSummary :: SummaryFun -> Mapping -> DataFrame -> DataFrame
runSummary SummaryFun
fn Mapping
m DataFrame
df =
    case Maybe ColumnRef -> DataFrame -> Maybe [Double]
resolveNum (Mapping -> Maybe ColumnRef
aesY Mapping
m) DataFrame
df of
        Maybe [Double]
Nothing -> DataFrame
df
        Just [Double]
ys ->
            let groupKeys :: [Text]
groupKeys = case Mapping -> Maybe ColumnRef
aesX Mapping
m of
                    Just (ColumnRef Text
n) -> case Text -> DataFrame -> Maybe Column
lookupColumn Text
n DataFrame
df of
                        Just Column
c -> Column -> [Text]
columnAsText Column
c
                        Maybe Column
Nothing -> Int -> Text -> [Text]
forall a. Int -> a -> [a]
replicate ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ys) Text
"all"
                    Maybe ColumnRef
Nothing -> Int -> Text -> [Text]
forall a. Int -> a -> [a]
replicate ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
ys) Text
"all"
                summaries :: [(Text, Double)]
summaries = SummaryFun -> [(Text, Double)] -> [(Text, Double)]
summarizeByCategory SummaryFun
fn ([Text] -> [Double] -> [(Text, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
groupKeys [Double]
ys)
                xName :: Text
xName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesX Mapping
m) Text
"x"
                yName :: Text
yName = Maybe ColumnRef -> Text -> Text
refName (Mapping -> Maybe ColumnRef
aesY Mapping
m) Text
"y"
             in [(Text, Column)] -> DataFrame
fromColumns
                    [ (Text
xName, [Text] -> Column
ColCat [Text
k | (Text
k, Double
_) <- [(Text, Double)]
summaries])
                    , (Text
yName, [Double] -> Column
ColNum [Double
v | (Text
_, Double
v) <- [(Text, Double)]
summaries])
                    ]

summarizeByCategory :: SummaryFun -> [(Text, Double)] -> [(Text, Double)]
summarizeByCategory :: SummaryFun -> [(Text, Double)] -> [(Text, Double)]
summarizeByCategory SummaryFun
fn [(Text, Double)]
pairs =
    let grouped :: [(Text, [(Text, Double)])]
grouped = ((Text, Double) -> Text)
-> [(Text, Double)] -> [(Text, [(Text, Double)])]
forall k a. Eq k => (a -> k) -> [a] -> [(k, [a])]
groupBy (Text, Double) -> Text
forall a b. (a, b) -> a
fst [(Text, Double)]
pairs
     in [(Text
name, SummaryFun -> [Double] -> Double
forall {a}. (Fractional a, Ord a) => SummaryFun -> [a] -> a
applyFn SummaryFun
fn (((Text, Double) -> Double) -> [(Text, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Double) -> Double
forall a b. (a, b) -> b
snd [(Text, Double)]
vs)) | (Text
name, [(Text, Double)]
vs) <- [(Text, [(Text, Double)])]
grouped]
  where
    applyFn :: SummaryFun -> [a] -> a
applyFn SummaryFun
SumMean [a]
xs = if [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs then a
0 else [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
xs a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)
    applyFn SummaryFun
SumSum [a]
xs = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
xs
    applyFn SummaryFun
SumMedian [a]
xs =
        let s :: [a]
s = [a] -> [a]
forall a. Ord a => [a] -> [a]
List.sort [a]
xs
            n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s
         in if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                then a
0
                else
                    if Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                        then [a]
s [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)
                        else ([a]
s [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a -> a -> a
forall a. Num a => a -> a -> a
+ [a]
s [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)) a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
2
    applyFn (SumQuantile Double
p) [a]
xs =
        let s :: [a]
s = [a] -> [a]
forall a. Ord a => [a] -> [a]
List.sort [a]
xs
            n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s
            ix :: Int
ix = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double
p Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) :: Int))
         in if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then a
0 else [a]
s [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
ix

resolveNum :: Maybe ColumnRef -> DataFrame -> Maybe [Double]
resolveNum :: Maybe ColumnRef -> DataFrame -> Maybe [Double]
resolveNum Maybe ColumnRef
Nothing DataFrame
_ = Maybe [Double]
forall a. Maybe a
Nothing
resolveNum (Just (ColumnRef Text
n)) DataFrame
df = Text -> DataFrame -> Maybe Column
lookupColumn Text
n DataFrame
df Maybe Column -> (Column -> Maybe [Double]) -> Maybe [Double]
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Column -> Maybe [Double]
columnAsNum

refName :: Maybe ColumnRef -> Text -> Text
refName :: Maybe ColumnRef -> Text -> Text
refName (Just (ColumnRef Text
n)) Text
_ = Text
n
refName Maybe ColumnRef
Nothing Text
fallback = Text
fallback

groupBy :: (Eq k) => (a -> k) -> [a] -> [(k, [a])]
groupBy :: forall k a. Eq k => (a -> k) -> [a] -> [(k, [a])]
groupBy a -> k
f [a]
xs =
    let keys :: [k]
keys = [k] -> [k]
forall a. Eq a => [a] -> [a]
List.nub ((a -> k) -> [a] -> [k]
forall a b. (a -> b) -> [a] -> [b]
map a -> k
f [a]
xs)
     in [(k
k, [a
x | a
x <- [a]
xs, a -> k
f a
x k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k]) | k
k <- [k]
keys]